diff --git a/records_mover/creds/base_creds.py b/records_mover/creds/base_creds.py index 2bd18c0f4..72d54a102 100644 --- a/records_mover/creds/base_creds.py +++ b/records_mover/creds/base_creds.py @@ -55,7 +55,7 @@ def __init__(self, str, None] = PleaseInfer.token, default_airbyte_creds: Union[PleaseInfer, - str, + Dict[str, Any], None] = PleaseInfer.token) -> None: self._default_db_creds_name = default_db_creds_name self._default_aws_creds_name = default_aws_creds_name @@ -271,10 +271,10 @@ def default_scratch_gcs_url(self) -> Optional[str]: self._scratch_gcs_url = self._infer_scratch_gcs_url() return self._scratch_gcs_url - def _infer_airbyte_creds(self) -> Optional[str]: + def _infer_airbyte_creds(self) -> Dict[str, Any]: raise NotImplementedError - def airbyte(self) -> dict: + def airbyte(self) -> Optional[Dict[str, Any]]: if self._default_airbyte_creds is PleaseInfer.token: self._default_airbyte_creds = self._infer_airbyte_creds() return self._default_airbyte_creds diff --git a/records_mover/creds/creds_via_env.py b/records_mover/creds/creds_via_env.py index 891980ba0..629c2e78a 100644 --- a/records_mover/creds/creds_via_env.py +++ b/records_mover/creds/creds_via_env.py @@ -2,7 +2,7 @@ import os import base64 import json -from typing import Iterable, Optional +from typing import Iterable, Optional, Any, Dict from .base_creds import BaseCreds from typing import TYPE_CHECKING from db_facts.db_facts_types import DBFacts @@ -13,9 +13,9 @@ class CredsViaEnv(BaseCreds): - def _infer_airbyte_creds(self) -> Optional[str]: + def _infer_airbyte_creds(self) -> Dict[str, Any]: if 'AIRBYTE_CONNECTION' not in os.environ: - return None + return {} return { 'user': 'username', 'host': 'host', diff --git a/records_mover/records/airbyte/airbyte.py b/records_mover/records/airbyte/airbyte.py index 21d8aa26d..fdfa873ad 100644 --- a/records_mover/records/airbyte/airbyte.py +++ b/records_mover/records/airbyte/airbyte.py @@ -12,8 +12,9 @@ class AirbyteEngine: General thoughts, maybe this should encapsulate making a request to airbyte Thus we'd have a method for making a request which'd know how to authenticate """ + session: Session - def __init__(self, session: Session = None): + def __init__(self, session: Session): """ Args: session: Optional records mover Session, exposed for testing. diff --git a/records_mover/records/cli.py b/records_mover/records/cli.py index 963f6a146..25707f3d7 100644 --- a/records_mover/records/cli.py +++ b/records_mover/records/cli.py @@ -49,7 +49,7 @@ def job_fn(raw_config: Dict[str, Any]) -> None: return job_fn -def build_parser() -> argparse.ArgumentParser: +def build_parser(bootstrap_session: Session) -> argparse.ArgumentParser: # skip in-memory sources/targets like dataframes that don't make # sense from the command-line source_method_name_by_cli_name = { @@ -85,8 +85,6 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument('-hc', '--healthcheck', action='store_true', required=False, help='Returns health of the ' 'configured airbyte instance') subparsers = parser.add_subparsers(help='subcommand_help') - from records_mover import Session - bootstrap_session = Session() for source in sources: for target in targets: @@ -112,12 +110,14 @@ def main() -> None: warnings.filterwarnings("ignore", "Your application has authenticated using end user credentials") - parser = build_parser() + from records_mover import Session + session = Session() + parser = build_parser(session) args = parser.parse_args() raw_config = vars(args) func = getattr(args, 'func', None) if args.healthcheck: - engine = AirbyteEngine() + engine = AirbyteEngine(session) result = engine.healthcheck() if result: print("Airbyte Status: OK!") diff --git a/types/stubs/airbyte/__init__.pyi b/types/stubs/airbyte/__init__.pyi new file mode 100644 index 000000000..e69de29bb diff --git a/types/stubs/requests/__init__.pyi b/types/stubs/requests/__init__.pyi new file mode 100644 index 000000000..e5bc06be3 --- /dev/null +++ b/types/stubs/requests/__init__.pyi @@ -0,0 +1,2 @@ +def get(url, params=None, **kwargs): + ... \ No newline at end of file