dagster-airbyte 0.25.4__tar.gz → 0.25.5__tar.gz

{dagster-airbyte-0.25.4/dagster_airbyte.egg-info → dagster-airbyte-0.25.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dagster-airbyte
-Version: 0.25.4
+Version: 0.25.5
 Summary: Package for integrating Airbyte with Dagster.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-airbyte
 Author: Dagster Labs

{dagster-airbyte-0.25.4 → dagster-airbyte-0.25.5}/dagster_airbyte/__init__.py

@@ -21,11 +21,14 @@ from dagster_airbyte.asset_defs import (
 from dagster_airbyte.ops import airbyte_sync_op as airbyte_sync_op
 from dagster_airbyte.resources import (
     AirbyteCloudResource as AirbyteCloudResource,
+    AirbyteCloudWorkspace as AirbyteCloudWorkspace,
     AirbyteResource as AirbyteResource,
     AirbyteState as AirbyteState,
     airbyte_cloud_resource as airbyte_cloud_resource,
     airbyte_resource as airbyte_resource,
+    load_airbyte_cloud_asset_specs as load_airbyte_cloud_asset_specs,
 )
+from dagster_airbyte.translator import DagsterAirbyteTranslator as DagsterAirbyteTranslator
 from dagster_airbyte.types import AirbyteOutput as AirbyteOutput
 from dagster_airbyte.version import __version__ as __version__
 

{dagster-airbyte-0.25.4 → dagster-airbyte-0.25.5}/dagster_airbyte/resources.py

@@ -6,32 +6,52 @@ import time
 from abc import abstractmethod
 from contextlib import contextmanager
 from datetime import datetime, timedelta
-from typing import Any, Dict, List, Mapping, Optional, cast
+from typing import Any, Dict, List, Mapping, Optional, Sequence, cast
 
 import requests
 from dagster import (
     ConfigurableResource,
+    Definitions,
     Failure,
     InitResourceContext,
     _check as check,
     get_dagster_logger,
     resource,
 )
+from dagster._annotations import experimental
 from dagster._config.pythonic_config import infer_schema_from_config_class
+from dagster._core.definitions.asset_spec import AssetSpec
+from dagster._core.definitions.definitions_load_context import StateBackedDefinitionsLoader
 from dagster._core.definitions.resource_definition import dagster_maintained_resource
+from dagster._model import DagsterModel
+from dagster._record import record
 from dagster._utils.cached_method import cached_method
 from dagster._utils.merger import deep_merge_dicts
 from pydantic import Field, PrivateAttr
 from requests.exceptions import RequestException
 
+from dagster_airbyte.translator import (
+    AirbyteConnection,
+    AirbyteDestination,
+    AirbyteWorkspaceData,
+    DagsterAirbyteTranslator,
+)
 from dagster_airbyte.types import AirbyteOutput
 
+AIRBYTE_REST_API_BASE = "https://api.airbyte.com"
+AIRBYTE_REST_API_VERSION = "v1"
+
+AIRBYTE_CONFIGURATION_API_BASE = "https://cloud.airbyte.com/api"
+AIRBYTE_CONFIGURATION_API_VERSION = "v1"
+
 DEFAULT_POLL_INTERVAL_SECONDS = 10
 
 # The access token expire every 3 minutes in Airbyte Cloud.
 # Refresh is needed after 2.5 minutes to avoid the "token expired" error message.
 AIRBYTE_CLOUD_REFRESH_TIMEDELTA_SECONDS = 150
 
+AIRBYTE_RECONSTRUCTION_METADATA_KEY_PREFIX = "dagster-airbyte/reconstruction_metadata"
+
 
 class AirbyteState:
     RUNNING = "running"
@@ -791,3 +811,326 @@ def airbyte_cloud_resource(context) -> AirbyteCloudResource:
 
     """
     return AirbyteCloudResource.from_resource_context(context)
+
+
+# -------------
+# Resources v2
+# -------------
+
+
+@experimental
+class AirbyteCloudClient(DagsterModel):
+    """This class exposes methods on top of the Airbyte APIs for Airbyte Cloud."""
+
+    workspace_id: str = Field(..., description="The Airbyte workspace ID")
+    client_id: str = Field(..., description="The Airbyte client ID.")
+    client_secret: str = Field(..., description="The Airbyte client secret.")
+    request_max_retries: int = Field(
+        ...,
+        description=(
+            "The maximum number of times requests to the Airbyte API should be retried "
+            "before failing."
+        ),
+    )
+    request_retry_delay: float = Field(
+        ...,
+        description="Time (in seconds) to wait between each request retry.",
+    )
+    request_timeout: int = Field(
+        ...,
+        description="Time (in seconds) after which the requests to Airbyte are declared timed out.",
+    )
+
+    _access_token_value: Optional[str] = PrivateAttr(default=None)
+    _access_token_timestamp: Optional[float] = PrivateAttr(default=None)
+
+    @property
+    @cached_method
+    def _log(self) -> logging.Logger:
+        return get_dagster_logger()
+
+    @property
+    def rest_api_base_url(self) -> str:
+        return f"{AIRBYTE_REST_API_BASE}/{AIRBYTE_REST_API_VERSION}"
+
+    @property
+    def configuration_api_base_url(self) -> str:
+        return f"{AIRBYTE_CONFIGURATION_API_BASE}/{AIRBYTE_CONFIGURATION_API_VERSION}"
+
+    @property
+    def all_additional_request_params(self) -> Mapping[str, Any]:
+        return {**self.authorization_request_params, **self.user_agent_request_params}
+
+    @property
+    def authorization_request_params(self) -> Mapping[str, Any]:
+        # Make sure the access token is refreshed before using it when calling the API.
+        if self._needs_refreshed_access_token():
+            self._refresh_access_token()
+        return {
+            "Authorization": f"Bearer {self._access_token_value}",
+        }
+
+    @property
+    def user_agent_request_params(self) -> Mapping[str, Any]:
+        return {
+            "User-Agent": "dagster",
+        }
+
+    def _refresh_access_token(self) -> None:
+        response = check.not_none(
+            self._make_request(
+                method="POST",
+                endpoint="applications/token",
+                base_url=self.rest_api_base_url,
+                data={
+                    "client_id": self.client_id,
+                    "client_secret": self.client_secret,
+                },
+                # Must not pass the bearer access token when refreshing it.
+                include_additional_request_params=False,
+            )
+        )
+        self._access_token_value = str(response["access_token"])
+        self._access_token_timestamp = datetime.now().timestamp()
+
+    def _needs_refreshed_access_token(self) -> bool:
+        return (
+            not self._access_token_value
+            or not self._access_token_timestamp
+            or self._access_token_timestamp
+            <= (
+                datetime.now() - timedelta(seconds=AIRBYTE_CLOUD_REFRESH_TIMEDELTA_SECONDS)
+            ).timestamp()
+        )
+
+    def _get_session(self, include_additional_request_params: bool) -> requests.Session:
+        headers = {"accept": "application/json"}
+        if include_additional_request_params:
+            headers = {
+                **headers,
+                **self.all_additional_request_params,
+            }
+        session = requests.Session()
+        session.headers.update(headers)
+        return session
+
+    def _make_request(
+        self,
+        method: str,
+        endpoint: str,
+        base_url: str,
+        data: Optional[Mapping[str, Any]] = None,
+        params: Optional[Mapping[str, Any]] = None,
+        include_additional_request_params: bool = True,
+    ) -> Mapping[str, Any]:
+        """Creates and sends a request to the desired Airbyte REST API endpoint.
+
+        Args:
+            method (str): The http method to use for this request (e.g. "POST", "GET", "PATCH").
+            endpoint (str): The Airbyte API endpoint to send this request to.
+            base_url (str): The base url to the Airbyte API to use.
+            data (Optional[Dict[str, Any]]): JSON-formatted data string to be included in the request.
+            params (Optional[Dict[str, Any]]): JSON-formatted query params to be included in the request.
+            include_additional_request_params (bool): Whether to include authorization and user-agent headers
+            to the request parameters. Defaults to True.
+
+        Returns:
+            Dict[str, Any]: Parsed json data from the response to this request
+        """
+        url = f"{base_url}/{endpoint}"
+
+        num_retries = 0
+        while True:
+            try:
+                session = self._get_session(
+                    include_additional_request_params=include_additional_request_params
+                )
+                response = session.request(
+                    method=method, url=url, json=data, params=params, timeout=self.request_timeout
+                )
+                response.raise_for_status()
+                return response.json()
+            except RequestException as e:
+                self._log.error(
+                    f"Request to Airbyte API failed for url {url} with method {method} : {e}"
+                )
+                if num_retries == self.request_max_retries:
+                    break
+                num_retries += 1
+                time.sleep(self.request_retry_delay)
+
+        raise Failure(f"Max retries ({self.request_max_retries}) exceeded with url: {url}.")
+
+    def get_connections(self) -> Mapping[str, Any]:
+        """Fetches all connections of an Airbyte workspace from the Airbyte REST API."""
+        return self._make_request(
+            method="GET",
+            endpoint="connections",
+            base_url=self.rest_api_base_url,
+            params={"workspaceIds": self.workspace_id},
+        )
+
+    def get_connection_details(self, connection_id) -> Mapping[str, Any]:
+        """Fetches details about a given connection from the Airbyte Configuration API.
+        The Airbyte Configuration API is an internal and may change in the future.
+        """
+        # Using the Airbyte Configuration API to get the connection details, including streams and their configs.
+        # https://airbyte-public-api-docs.s3.us-east-2.amazonaws.com/rapidoc-api-docs.html#post-/v1/connections/get
+        # https://github.com/airbytehq/airbyte-platform/blob/v1.0.0/airbyte-api/server-api/src/main/openapi/config.yaml
+        return self._make_request(
+            method="POST",
+            endpoint="connections/get",
+            base_url=self.configuration_api_base_url,
+            data={"connectionId": connection_id},
+        )
+
+    def get_destination_details(self, destination_id: str) -> Mapping[str, Any]:
+        """Fetches details about a given destination from the Airbyte REST API."""
+        return self._make_request(
+            method="GET",
+            endpoint=f"destinations/{destination_id}",
+            base_url=self.rest_api_base_url,
+        )
+
+
+@experimental
+class AirbyteCloudWorkspace(ConfigurableResource):
+    """This class represents a Airbyte Cloud workspace and provides utilities
+    to interact with Airbyte APIs.
+    """
+
+    workspace_id: str = Field(..., description="The Airbyte Cloud workspace ID")
+    client_id: str = Field(..., description="The Airbyte Cloud client ID.")
+    client_secret: str = Field(..., description="The Airbyte Cloud client secret.")
+    request_max_retries: int = Field(
+        default=3,
+        description=(
+            "The maximum number of times requests to the Airbyte API should be retried "
+            "before failing."
+        ),
+    )
+    request_retry_delay: float = Field(
+        default=0.25,
+        description="Time (in seconds) to wait between each request retry.",
+    )
+    request_timeout: int = Field(
+        default=15,
+        description="Time (in seconds) after which the requests to Airbyte are declared timed out.",
+    )
+
+    _client: AirbyteCloudClient = PrivateAttr(default=None)
+
+    @cached_method
+    def get_client(self) -> AirbyteCloudClient:
+        return AirbyteCloudClient(
+            workspace_id=self.workspace_id,
+            client_id=self.client_id,
+            client_secret=self.client_secret,
+            request_max_retries=self.request_max_retries,
+            request_retry_delay=self.request_retry_delay,
+            request_timeout=self.request_timeout,
+        )
+
+    def fetch_airbyte_workspace_data(
+        self,
+    ) -> AirbyteWorkspaceData:
+        """Retrieves all Airbyte content from the workspace and returns it as a AirbyteWorkspaceData object.
+
+        Returns:
+            AirbyteWorkspaceData: A snapshot of the Airbyte workspace's content.
+        """
+        connections_by_id = {}
+        destinations_by_id = {}
+
+        client = self.get_client()
+        connections = client.get_connections()["data"]
+
+        for partial_connection_details in connections:
+            full_connection_details = client.get_connection_details(
+                connection_id=partial_connection_details["connectionId"]
+            )
+            connection = AirbyteConnection.from_connection_details(
+                connection_details=full_connection_details
+            )
+            connections_by_id[connection.id] = connection
+
+            destination_details = client.get_destination_details(
+                destination_id=connection.destination_id
+            )
+            destination = AirbyteDestination.from_destination_details(
+                destination_details=destination_details
+            )
+            destinations_by_id[destination.id] = destination
+
+        return AirbyteWorkspaceData(
+            connections_by_id=connections_by_id,
+            destinations_by_id=destinations_by_id,
+        )
+
+
+@experimental
+def load_airbyte_cloud_asset_specs(
+    workspace: AirbyteCloudWorkspace,
+    dagster_airbyte_translator: Optional[DagsterAirbyteTranslator] = None,
+) -> Sequence[AssetSpec]:
+    """Returns a list of AssetSpecs representing the Airbyte content in the workspace.
+
+    Args:
+        workspace (AirbyteCloudWorkspace): The Airbyte Cloud workspace to fetch assets from.
+        dagster_airbyte_translator (Optional[DagsterAirbyteTranslator], optional): The translator to use
+            to convert Airbyte content into :py:class:`dagster.AssetSpec`.
+            Defaults to :py:class:`DagsterAirbyteTranslator`.
+
+    Returns:
+        List[AssetSpec]: The set of assets representing the Airbyte content in the workspace.
+
+    Examples:
+        Loading the asset specs for a given Airbyte Cloud workspace:
+
+        .. code-block:: python
+
+            from dagster_airbyte import AirbyteCloudWorkspace, load_airbyte_cloud_asset_specs
+
+            import dagster as dg
+
+            airbyte_cloud_workspace = AirbyteCloudWorkspace(
+                workspace_id=dg.EnvVar("AIRBYTE_CLOUD_WORKSPACE_ID"),
+                client_id=dg.EnvVar("AIRBYTE_CLOUD_CLIENT_ID"),
+                client_secret=dg.EnvVar("AIRBYTE_CLOUD_CLIENT_SECRET"),
+            )
+
+
+            airbyte_cloud_specs = load_airbyte_cloud_asset_specs(airbyte_cloud_workspace)
+            defs = dg.Definitions(assets=airbyte_cloud_specs)
+    """
+    with workspace.process_config_and_initialize_cm() as initialized_workspace:
+        return check.is_list(
+            AirbyteCloudWorkspaceDefsLoader(
+                workspace=initialized_workspace,
+                translator=dagster_airbyte_translator or DagsterAirbyteTranslator(),
+            )
+            .build_defs()
+            .assets,
+            AssetSpec,
+        )
+
+
+@record
+class AirbyteCloudWorkspaceDefsLoader(StateBackedDefinitionsLoader[Mapping[str, Any]]):
+    workspace: AirbyteCloudWorkspace
+    translator: DagsterAirbyteTranslator
+
+    @property
+    def defs_key(self) -> str:
+        return f"{AIRBYTE_RECONSTRUCTION_METADATA_KEY_PREFIX}/{self.workspace.workspace_id}"
+
+    def fetch_state(self) -> AirbyteWorkspaceData:
+        return self.workspace.fetch_airbyte_workspace_data()
+
+    def defs_from_state(self, state: AirbyteWorkspaceData) -> Definitions:
+        all_asset_specs = [
+            self.translator.get_asset_spec(props)
+            for props in state.to_airbyte_connection_table_props_data()
+        ]
+
+        return Definitions(assets=all_asset_specs)

dagster-airbyte-0.25.5/dagster_airbyte/translator.py

@@ -0,0 +1,190 @@
+from typing import Any, List, Mapping, Optional, Sequence
+
+from dagster._annotations import experimental
+from dagster._core.definitions.asset_key import AssetKey
+from dagster._core.definitions.asset_spec import AssetSpec
+from dagster._core.definitions.metadata.metadata_set import NamespacedMetadataSet, TableMetadataSet
+from dagster._record import record
+from dagster._serdes.serdes import whitelist_for_serdes
+from dagster._utils.cached_method import cached_method
+
+from dagster_airbyte.utils import generate_table_schema, get_airbyte_connection_table_name
+
+
+@record
+class AirbyteConnectionTableProps:
+    table_name: str
+    stream_prefix: Optional[str]
+    stream_name: str
+    json_schema: Mapping[str, Any]
+    connection_id: str
+    connection_name: str
+    database: Optional[str]
+    schema: Optional[str]
+
+    @property
+    def fully_qualified_table_name(self) -> Optional[str]:
+        return (
+            f"{self.database}.{self.schema}.{self.stream_name}"
+            if self.database and self.schema
+            else None
+        )
+
+
+@whitelist_for_serdes
+@record
+class AirbyteConnection:
+    """Represents an Airbyte connection, based on data as returned from the API."""
+
+    id: str
+    name: str
+    stream_prefix: Optional[str]
+    streams: Mapping[str, "AirbyteStream"]
+    destination_id: str
+
+    @classmethod
+    def from_connection_details(
+        cls,
+        connection_details: Mapping[str, Any],
+    ) -> "AirbyteConnection":
+        return cls(
+            id=connection_details["connectionId"],
+            name=connection_details["name"],
+            stream_prefix=connection_details.get("prefix"),
+            streams={
+                stream_details["stream"]["name"]: AirbyteStream.from_stream_details(
+                    stream_details=stream_details
+                )
+                for stream_details in connection_details.get("syncCatalog", {}).get("streams", [])
+            },
+            destination_id=connection_details["destinationId"],
+        )
+
+
+@whitelist_for_serdes
+@record
+class AirbyteDestination:
+    """Represents an Airbyte destination, based on data as returned from the API."""
+
+    id: str
+    database: Optional[str]
+    schema: Optional[str]
+
+    @classmethod
+    def from_destination_details(
+        cls,
+        destination_details: Mapping[str, Any],
+    ) -> "AirbyteDestination":
+        return cls(
+            id=destination_details["destinationId"],
+            database=destination_details["configuration"].get("database"),
+            schema=destination_details["configuration"].get("schema"),
+        )
+
+
+@whitelist_for_serdes
+@record
+class AirbyteStream:
+    """Represents an Airbyte stream, based on data as returned from the API.
+    A stream in Airbyte corresponds to a table.
+    """
+
+    name: str
+    selected: bool
+    json_schema: Mapping[str, Any]
+
+    @classmethod
+    def from_stream_details(
+        cls,
+        stream_details: Mapping[str, Any],
+    ) -> "AirbyteStream":
+        return cls(
+            name=stream_details["stream"]["name"],
+            selected=stream_details["config"].get("selected", False),
+            json_schema=stream_details["stream"].get("jsonSchema", {}),
+        )
+
+
+@whitelist_for_serdes
+@record
+class AirbyteWorkspaceData:
+    """A record representing all content in an Airbyte workspace.
+    This applies to both Airbyte OSS and Cloud.
+    """
+
+    connections_by_id: Mapping[str, AirbyteConnection]
+    destinations_by_id: Mapping[str, AirbyteDestination]
+
+    @cached_method
+    def to_airbyte_connection_table_props_data(self) -> Sequence[AirbyteConnectionTableProps]:
+        """Method that converts a `AirbyteWorkspaceData` object
+        to a collection of `AirbyteConnectionTableProps` objects.
+        """
+        data: List[AirbyteConnectionTableProps] = []
+
+        for connection in self.connections_by_id.values():
+            destination = self.destinations_by_id[connection.destination_id]
+
+            for stream in connection.streams.values():
+                if stream.selected:
+                    data.append(
+                        AirbyteConnectionTableProps(
+                            table_name=get_airbyte_connection_table_name(
+                                stream_prefix=connection.stream_prefix,
+                                stream_name=stream.name,
+                            ),
+                            stream_prefix=connection.stream_prefix,
+                            stream_name=stream.name,
+                            json_schema=stream.json_schema,
+                            connection_id=connection.id,
+                            connection_name=connection.name,
+                            database=destination.database,
+                            schema=destination.schema,
+                        )
+                    )
+
+        return data
+
+
+class AirbyteMetadataSet(NamespacedMetadataSet):
+    connection_id: str
+    connection_name: str
+    stream_prefix: Optional[str] = None
+
+    @classmethod
+    def namespace(cls) -> str:
+        return "dagster-airbyte"
+
+
+@experimental
+class DagsterAirbyteTranslator:
+    """Translator class which converts a `AirbyteConnectionTableProps` object into AssetSpecs.
+    Subclass this class to implement custom logic how to translate Airbyte content into asset spec.
+    """
+
+    def get_asset_spec(self, props: AirbyteConnectionTableProps) -> AssetSpec:
+        """Get the AssetSpec for a table synced by an Airbyte connection."""
+        table_schema_props = (
+            props.json_schema.get("properties")
+            or props.json_schema.get("items", {}).get("properties")
+            or {}
+        )
+        column_schema = generate_table_schema(table_schema_props)
+
+        metadata = {
+            **TableMetadataSet(
+                column_schema=column_schema,
+                table_name=props.fully_qualified_table_name,
+            ),
+            **AirbyteMetadataSet(
+                connection_id=props.connection_id,
+                connection_name=props.connection_name,
+                stream_prefix=props.stream_prefix,
+            ),
+        }
+
+        return AssetSpec(
+            key=AssetKey(props.table_name),
+            metadata=metadata,
+            kinds={"airbyte"},
+        )

{dagster-airbyte-0.25.4 → dagster-airbyte-0.25.5}/dagster_airbyte/utils.py

@@ -6,6 +6,10 @@ from dagster._core.definitions.metadata.table import TableColumn, TableSchema
 from dagster_airbyte.types import AirbyteOutput
 
 
+def get_airbyte_connection_table_name(stream_prefix: Optional[str], stream_name: str) -> str:
+    return f"{stream_prefix if stream_prefix else ''}{stream_name}"
+
+
 def generate_table_schema(stream_schema_props: Mapping[str, Any]) -> TableSchema:
     return TableSchema(
         columns=sorted(

dagster-airbyte-0.25.5/dagster_airbyte/version.py

@@ -0,0 +1 @@
+__version__ = "0.25.5"

{dagster-airbyte-0.25.4 → dagster-airbyte-0.25.5/dagster_airbyte.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dagster-airbyte
-Version: 0.25.4
+Version: 0.25.5
 Summary: Package for integrating Airbyte with Dagster.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-airbyte
 Author: Dagster Labs

{dagster-airbyte-0.25.4 → dagster-airbyte-0.25.5}/dagster_airbyte.egg-info/SOURCES.txt

@@ -9,6 +9,7 @@ dagster_airbyte/cli.py
 dagster_airbyte/ops.py
 dagster_airbyte/py.typed
 dagster_airbyte/resources.py
+dagster_airbyte/translator.py
 dagster_airbyte/types.py
 dagster_airbyte/utils.py
 dagster_airbyte/version.py

{dagster-airbyte-0.25.4 → dagster-airbyte-0.25.5}/dagster_airbyte.egg-info/requires.txt

@@ -1,8 +1,8 @@
-dagster==1.9.4
+dagster==1.9.5
 requests
 
 [managed]
-dagster-managed-elements==0.25.4
+dagster-managed-elements==0.25.5
 
 [test]
 requests-mock

{dagster-airbyte-0.25.4 → dagster-airbyte-0.25.5}/setup.py

@@ -37,7 +37,7 @@ setup(
     include_package_data=True,
     python_requires=">=3.9,<3.13",
     install_requires=[
-        "dagster==1.9.4",
+        "dagster==1.9.5",
         "requests",
     ],
     zip_safe=False,
@@ -52,7 +52,7 @@ setup(
             "flaky",
         ],
         "managed": [
-            "dagster-managed-elements==0.25.4",
+            "dagster-managed-elements==0.25.5",
         ],
     },
 )

dagster-airbyte-0.25.4/dagster_airbyte/version.py

@@ -1 +0,0 @@
-__version__ = "0.25.4"