jprdonnelly/airbyte
1
0
Fork 0
mirror of synced 2026-01-20 12:07:14 -05:00
Code Issues Packages Projects Releases Wiki Activity

Run MyPy on CDK/base-python and fix issues. (#3175)

Browse Source
This commit is contained in:
Davin Chia
2021-05-04 11:02:53 +08:00
committed by GitHub
parent dd45537239
commit 72e7fe35a3
11 changed files with 54 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
airbyte-cdk/python/airbyte_cdk/base_python/cdk/abstract_source.py
View File

@@ -166,9 +166,10 @@ class AbstractSource(Source, ABC):
yield self._checkpoint_state(stream_name, stream_state, connector_state, logger)
def _read_full_refresh(self, stream_instance: Stream, configured_stream: ConfiguredAirbyteStream) -> Iterator[AirbyteMessage]:
args = {"sync_mode": SyncMode.full_refresh, "cursor_field": configured_stream.cursor_field}
for slices in stream_instance.stream_slices(**args):
for record in stream_instance.read_records(stream_slice=slices, **args):
slices = stream_instance.stream_slices(sync_mode=SyncMode.full_refresh, cursor_field=configured_stream.cursor_field)
for slice in slices:
records = stream_instance.read_records(stream_slice=slice, sync_mode=SyncMode.full_refresh, cursor_field=configured_stream.cursor_field)
for record in records:
yield self._as_airbyte_record(configured_stream.stream.name, record)
def _checkpoint_state(self, stream_name, stream_state, connector_state, logger):

8
airbyte-cdk/python/airbyte_cdk/base_python/cdk/streams/auth/oauth.py
View File

@@ -21,7 +21,7 @@
# SOFTWARE.
from typing import Any, Mapping, Tuple
from typing import Any, Mapping, Tuple, List, Union, MutableMapping
import pendulum
import requests
@@ -34,7 +34,7 @@ class Oauth2Authenticator(HttpAuthenticator):
The generated access token is attached to each request via the Authorization header.
"""
def __init__(self, token_refresh_endpoint: str, client_id: str, client_secret: str, refresh_token: str, scopes: [str] = None):
def __init__(self, token_refresh_endpoint: str, client_id: str, client_secret: str, refresh_token: str, scopes: List[str] = None):
self.token_refresh_endpoint = token_refresh_endpoint
self.client_secret = client_secret
self.client_id = client_id
@@ -59,9 +59,9 @@ class Oauth2Authenticator(HttpAuthenticator):
def token_has_expired(self) -> bool:
return pendulum.now() > self._token_expiry_date
def get_refresh_request_body(self) -> Mapping[str, any]:
def get_refresh_request_body(self) -> Mapping[str, Any]:
""" Override to define additional parameters """
payload = {
payload: MutableMapping[str, Any] = {
"grant_type": "refresh_token",
"client_id": self.client_id,
"client_secret": self.client_secret,

8
airbyte-cdk/python/airbyte_cdk/base_python/cdk/streams/core.py
View File

@@ -33,7 +33,7 @@ from airbyte_cdk.base_python.schema_helpers import ResourceSchemaLoader
def package_name_from_class(cls: object) -> str:
"""Find the package name given a class name"""
module = inspect.getmodule(cls)
module: Any = inspect.getmodule(cls)
return module.__name__.split(".")[0]
@@ -57,7 +57,7 @@ class Stream(ABC):
self,
sync_mode: SyncMode,
cursor_field: List[str] = None,
stream_slice: Mapping[str, any] = None,
stream_slice: Mapping[str, Any] = None,
stream_state: Mapping[str, Any] = None,
) -> Iterable[Mapping[str, Any]]:
"""
@@ -79,7 +79,7 @@ class Stream(ABC):
if self.supports_incremental:
stream.source_defined_cursor = self.source_defined_cursor
stream.supported_sync_modes.append(SyncMode.incremental)
stream.supported_sync_modes.append(SyncMode.incremental) # type: ignore
stream.default_cursor_field = self._wrapped_cursor_field()
return stream
@@ -110,7 +110,7 @@ class Stream(ABC):
def stream_slices(
self, sync_mode: SyncMode, cursor_field: List[str] = None, stream_state: Mapping[str, Any] = None
) -> Iterable[Optional[Mapping[str, any]]]:
) -> Iterable[Optional[Mapping[str, Any]]]:
"""
Override to define the slices for this stream. See the stream slicing section of the docs for more information.

21
airbyte-cdk/python/airbyte_cdk/base_python/cdk/streams/http.py
View File

@@ -204,28 +204,27 @@ class HttpStream(Stream, ABC):
def read_records(
self,
sync_mode: SyncMode,
stream_slice: Optional[Mapping[str, Any]] = None,
stream_state: Optional[Mapping[str, Any]] = None,
cursor_field: List[str] = None,
stream_slice: Mapping[str, Any] = None,
stream_state: Mapping[str, Any] = None,
) -> Iterable[Mapping[str, Any]]:
stream_state = stream_state or {}
args = {"stream_state": stream_state, "stream_slice": stream_slice}
pagination_complete = False
while not pagination_complete:
next_page_token = None
request_headers = self.request_headers(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token)
request = self._create_prepared_request(
path=self.path(**args),
headers=dict(self.request_headers(**args), **self.authenticator.get_auth_header()),
params=self.request_params(**args),
json=self.request_body_json(**args),
path=self.path(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token),
headers=dict(request_headers, **self.authenticator.get_auth_header()),
params=self.request_params(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token),
json=self.request_body_json(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token),
)
response = self._send_request(request)
yield from self.parse_response(response, **args)
yield from self.parse_response(response, stream_state=stream_state, stream_slice=stream_slice)
next_page_token = self.next_page_token(response)
if next_page_token:
args["next_page_token"] = next_page_token
else:
if not next_page_token:
pagination_complete = True
# Always return an empty generator just in case no records were ever yielded

4
airbyte-cdk/python/airbyte_cdk/base_python/client.py
View File

@@ -32,7 +32,7 @@ from .schema_helpers import ResourceSchemaLoader
def package_name_from_class(cls: object) -> str:
"""Find the package name given a class name"""
module = inspect.getmodule(cls)
module: Any = inspect.getmodule(cls)
return module.__name__.split(".")[0]
@@ -60,7 +60,7 @@ class BaseClient(StreamStateMixin, ABC):
self._schema_loader = self.schema_loader_class(package_name)
self._stream_methods = self._enumerate_methods()
def _enumerate_methods(self) -> Mapping[str, callable]:
def _enumerate_methods(self) -> Mapping[str, Callable]:
"""Detect available streams and return mapping"""
prefix = "stream__"
mapping = {}

26
airbyte-cdk/python/airbyte_cdk/base_python/integration.py
View File

@@ -25,7 +25,7 @@ import json
import os
import pkgutil
from collections import defaultdict
from typing import Dict, Generator
from typing import Dict, Iterable, Any, Optional, Mapping, MutableMapping
from airbyte_cdk.models import AirbyteCatalog, AirbyteConnectionStatus, AirbyteMessage, ConfiguredAirbyteCatalog, ConnectorSpecification
@@ -34,8 +34,8 @@ from .logger import AirbyteLogger
class AirbyteSpec(object):
@staticmethod
def from_file(file: str):
with open(file) as file:
def from_file(file_name: str):
with open(file_name) as file:
spec_text = file.read()
return AirbyteSpec(spec_text)
@@ -46,7 +46,7 @@ class AirbyteSpec(object):
class Integration(object):
# can be overridden to change an input config
def configure(self, config: json, temp_dir: str) -> json:
def configure(self, config: Mapping[str, Any], temp_dir: str) -> Mapping[str, Any]:
"""
Persist config in temporary directory to run the Source job
"""
@@ -55,13 +55,13 @@ class Integration(object):
return config
@staticmethod
def read_config(config_path: str) -> json:
def read_config(config_path: str) -> Mapping[str, Any]:
with open(config_path, "r") as file:
contents = file.read()
return json.loads(contents)
@staticmethod
def write_config(config: json, config_path: str):
def write_config(config: Mapping[str, Any], config_path: str):
with open(config_path, "w") as fh:
fh.write(json.dumps(config))
@@ -70,7 +70,7 @@ class Integration(object):
return ConfiguredAirbyteCatalog.parse_obj(self.read_config(catalog_path))
# can be overridden to change an input state
def read_state(self, state_path: str) -> Dict[str, any]:
def read_state(self, state_path: str) -> Dict[str, Any]:
if state_path:
state_obj = json.loads(open(state_path, "r").read())
else:
@@ -83,17 +83,19 @@ class Integration(object):
Returns the spec for this integration. The spec is a JSON-Schema object describing the required configurations (e.g: username and password)
required to run this integration.
"""
raw_spec = pkgutil.get_data(self.__class__.__module__.split(".")[0], "spec.json")
raw_spec: Optional[bytes] = pkgutil.get_data(self.__class__.__module__.split(".")[0], "spec.json")
if not raw_spec:
raise ValueError("Unable to find spec.json.")
return ConnectorSpecification.parse_obj(json.loads(raw_spec))
def check(self, logger: AirbyteLogger, config: json) -> AirbyteConnectionStatus:
def check(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the integration e.g: if a provided Stripe API token can be used to connect
to the Stripe API.
"""
raise Exception("Not Implemented")
def discover(self, logger: AirbyteLogger, config: json) -> AirbyteCatalog:
def discover(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteCatalog:
"""
Returns an AirbyteCatalog representing the available streams and fields in this integration. For example, given valid credentials to a
Postgres database, returns an Airbyte catalog where each postgres table is a stream, and each table column is a field.
@@ -106,8 +108,8 @@ class Source(Integration):
super().__init__()
def read(
self, logger: AirbyteLogger, config: json, catalog: ConfiguredAirbyteCatalog, state_path: Dict[str, any]
) -> Generator[AirbyteMessage, None, None]:
self, logger: AirbyteLogger, config: Mapping[str, Any], catalog: ConfiguredAirbyteCatalog, state: MutableMapping[str, Any] = None
) -> Iterable[AirbyteMessage]:
"""
Returns a generator of the AirbyteMessages generated by reading the source with the given configuration, catalog, and state.
"""

10
airbyte-cdk/python/airbyte_cdk/base_python/source.py
View File

@@ -23,7 +23,7 @@
import copy
from datetime import datetime
from typing import Any, Iterator, Mapping, MutableMapping, Type
from typing import Any, Iterable, Mapping, MutableMapping, Type
from airbyte_cdk.models import (
AirbyteCatalog,
@@ -46,7 +46,7 @@ from .logger import AirbyteLogger
class BaseSource(Source):
"""Base source that designed to work with clients derived from BaseClient"""
client_class: Type[BaseClient] = None
client_class: Type[BaseClient]
@property
def name(self) -> str:
@@ -55,9 +55,7 @@ class BaseSource(Source):
def _get_client(self, config: Mapping):
"""Construct client"""
client = self.client_class(**config)
return client
return self.client_class(**config)
def discover(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteCatalog:
"""Discover streams"""
@@ -76,7 +74,7 @@ class BaseSource(Source):
def read(
self, logger: AirbyteLogger, config: Mapping[str, Any], catalog: ConfiguredAirbyteCatalog, state: MutableMapping[str, Any] = None
) -> Iterator[AirbyteMessage]:
) -> Iterable[AirbyteMessage]:
state = state or {}
client = self._get_client(config)