finish dagster-airbyte types (#7562)

* finish dagster-airbyte types

examples/modern_data_stack_assets/modern_data_stack_assets/setup_airbyte.py

@@ -35,10 +35,9 @@ def _create_ab_source(client: AirbyteResource) -> str:
     ]
 
     # get latest available Postgres source definition
-    source_defs = client.make_request(
-        "/source_definitions/list_latest", data={"workspaceId": workspace_id}
+    source_defs = _safe_request(
+        client, "/source_definitions/list_latest", data={"workspaceId": workspace_id}
     )
-    assert source_defs
     postgres_definitions = [
         sd for sd in source_defs["sourceDefinitions"] if sd["name"] == "Postgres"
     ]

python_modules/libraries/dagster-airbyte/dagster_airbyte/resources.py

@@ -1,7 +1,7 @@
 import logging
 import sys
 import time
-from typing import Any, Dict, Optional
+from typing import Dict, List, Optional, cast
 
 import requests
 from dagster_airbyte.types import AirbyteOutput
@@ -53,8 +53,8 @@ def api_base_url(self) -> str:
         )
 
     def make_request(
-        self, endpoint: str, data: Optional[Dict[str, Any]]
-    ) -> Optional[Dict[str, Any]]:
+        self, endpoint: str, data: Optional[Dict[str, object]]
+    ) -> Optional[Dict[str, object]]:
         """
         Creates and sends a request to the desired Airbyte REST API endpoint.
 
@@ -91,18 +91,18 @@ def make_request(
 
         raise Failure("Exceeded max number of retries.")
 
-    def cancel_job(self, job_id: str):
+    def cancel_job(self, job_id: int):
         self.make_request(endpoint="/jobs/cancel", data={"id": job_id})
 
     def get_job_status(self, job_id: int) -> dict:
         return check.is_dict(self.make_request(endpoint="/jobs/get", data={"id": job_id}))
 
-    def start_sync(self, connection_id: str) -> dict:
+    def start_sync(self, connection_id: str) -> Dict[str, object]:
         return check.not_none(
             self.make_request(endpoint="/connections/sync", data={"connectionId": connection_id})
         )
 
-    def get_connection_details(self, connection_id: str) -> dict:
+    def get_connection_details(self, connection_id: str) -> Dict[str, object]:
         return check.not_none(
             self.make_request(endpoint="/connections/get", data={"connectionId": connection_id})
         )
@@ -112,7 +112,7 @@ def sync_and_poll(
         connection_id: str,
         poll_interval: float = DEFAULT_POLL_INTERVAL_SECONDS,
         poll_timeout: Optional[float] = None,
-    ):
+    ) -> AirbyteOutput:
         """
         Initializes a sync operation for the given connector, and polls until it completes.
 
@@ -129,7 +129,9 @@ def sync_and_poll(
         """
         connection_details = self.get_connection_details(connection_id)
         job_details = self.start_sync(connection_id)
-        job_id = job_details.get("job", {}).get("id")
+        job_info = cast(Dict[str, object], job_details.get("job", {}))
+        job_id = cast(int, job_info.get("id"))
+
         self._log.info(f"Job {job_id} initialized for connection_id={connection_id}.")
         start = time.monotonic()
         logged_attempts = 0
@@ -143,12 +145,11 @@ def sync_and_poll(
                     )
                 time.sleep(poll_interval)
                 job_details = self.get_job_status(job_id)
-                cur_attempt = len(job_details.get("attempts", []))
+                attempts = cast(List, job_details.get("attempts", []))
+                cur_attempt = len(attempts)
                 # spit out the available Airbyte log info
                 if cur_attempt:
-                    log_lines = (
-                        job_details["attempts"][logged_attempts].get("logs", {}).get("logLines", [])
-                    )
+                    log_lines = attempts[logged_attempts].get("logs", {}).get("logLines", [])
 
                     for line in log_lines[logged_lines:]:
                         sys.stdout.write(line + "\n")
@@ -160,7 +161,8 @@ def sync_and_poll(
                         logged_lines = 0
                         logged_attempts += 1
 
-                state = job_details.get("job", {}).get("status")
+                job_info = cast(Dict[str, object], job_details.get("job", {}))
+                state = job_info.get("status")
 
                 if state in (AirbyteState.RUNNING, AirbyteState.PENDING, AirbyteState.INCOMPLETE):
                     continue

python_modules/libraries/dagster-airbyte/dagster_airbyte/utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, Iterator, List
 
 from dagster_airbyte.types import AirbyteOutput
 
@@ -29,7 +29,9 @@ def _materialization_for_stream(
     )
 
 
-def generate_materializations(output: AirbyteOutput, asset_key_prefix: List[str]):
+def generate_materializations(
+    output: AirbyteOutput, asset_key_prefix: List[str]
+) -> Iterator[AssetMaterialization]:
     prefix = output.connection_details.get("prefix") or ""
     # all the streams that are set to be sync'd by this connection
     all_stream_props = {