dagster-airbyte 0.25.6__tar.gz → 0.25.7__tar.gz

{dagster-airbyte-0.25.6/dagster_airbyte.egg-info → dagster-airbyte-0.25.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dagster-airbyte
-Version: 0.25.6
+Version: 0.25.7
 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.6 → dagster-airbyte-0.25.7}/dagster_airbyte/__init__.py

@@ -14,8 +14,10 @@ try:
 except ImportError:
     pass
 
+from dagster_airbyte.asset_decorator import airbyte_assets as airbyte_assets
 from dagster_airbyte.asset_defs import (
     build_airbyte_assets as build_airbyte_assets,
+    build_airbyte_assets_definitions as build_airbyte_assets_definitions,
     load_assets_from_airbyte_instance as load_assets_from_airbyte_instance,
 )
 from dagster_airbyte.ops import airbyte_sync_op as airbyte_sync_op
@@ -23,12 +25,16 @@ 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.translator import (
+    AirbyteConnectionTableProps as AirbyteConnectionTableProps,
+    AirbyteJobStatusType as AirbyteJobStatusType,
+    AirbyteState as AirbyteState,
+    DagsterAirbyteTranslator as DagsterAirbyteTranslator,
+)
 from dagster_airbyte.types import AirbyteOutput as AirbyteOutput
 from dagster_airbyte.version import __version__ as __version__
 

dagster-airbyte-0.25.7/dagster_airbyte/asset_decorator.py

@@ -0,0 +1,113 @@
+from typing import Any, Callable, Optional
+
+from dagster import AssetsDefinition, multi_asset
+from dagster._annotations import experimental
+
+from dagster_airbyte.resources import AirbyteCloudWorkspace
+from dagster_airbyte.translator import AirbyteMetadataSet, DagsterAirbyteTranslator
+
+
+@experimental
+def airbyte_assets(
+    *,
+    connection_id: str,
+    workspace: AirbyteCloudWorkspace,
+    name: Optional[str] = None,
+    group_name: Optional[str] = None,
+    dagster_airbyte_translator: Optional[DagsterAirbyteTranslator] = None,
+) -> Callable[[Callable[..., Any]], AssetsDefinition]:
+    """Create a definition for how to sync the tables of a given Airbyte connection.
+
+    Args:
+        connection_id (str): The Airbyte Connection ID.
+        workspace (AirbyteCloudWorkspace): The Airbyte workspace to fetch assets from.
+        name (Optional[str], optional): The name of the op.
+        group_name (Optional[str], optional): The name of the asset group.
+        dagster_airbyte_translator (Optional[DagsterAirbyteTranslator], optional): The translator to use
+            to convert Airbyte content into :py:class:`dagster.AssetSpec`.
+            Defaults to :py:class:`DagsterAirbyteTranslator`.
+
+    Examples:
+        Sync the tables of an Airbyte connection:
+
+        .. code-block:: python
+
+            from dagster_airbyte import AirbyteCloudWorkspace, airbyte_assets
+
+            import dagster as dg
+
+            airbyte_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_assets(
+                connection_id="airbyte_connection_id",
+                workspace=airbyte_workspace,
+            )
+            def airbyte_connection_assets(context: dg.AssetExecutionContext, airbyte: AirbyteCloudWorkspace):
+                yield from airbyte.sync_and_poll(context=context)
+
+
+            defs = dg.Definitions(
+                assets=[airbyte_connection_assets],
+                resources={"airbyte": airbyte_workspace},
+            )
+
+        Sync the tables of an Airbyte connection with a custom translator:
+
+        .. code-block:: python
+
+            from dagster_airbyte import (
+                DagsterAirbyteTranslator,
+                AirbyteConnectionTableProps,
+                AirbyteCloudWorkspace,
+                airbyte_assets
+            )
+
+            import dagster as dg
+
+            class CustomDagsterAirbyteTranslator(DagsterAirbyteTranslator):
+                def get_asset_spec(self, props: AirbyteConnectionTableProps) -> dg.AssetSpec:
+                    default_spec = super().get_asset_spec(props)
+                    return default_spec.merge_attributes(
+                        metadata={"custom": "metadata"},
+                    )
+
+            airbyte_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_assets(
+                connection_id="airbyte_connection_id",
+                workspace=airbyte_workspace,
+                dagster_airbyte_translator=CustomDagsterAirbyteTranslator()
+            )
+            def airbyte_connection_assets(context: dg.AssetExecutionContext, airbyte: AirbyteCloudWorkspace):
+                yield from airbyte.sync_and_poll(context=context)
+
+
+            defs = dg.Definitions(
+                assets=[airbyte_connection_assets],
+                resources={"airbyte": airbyte_workspace},
+            )
+    """
+    dagster_airbyte_translator = dagster_airbyte_translator or DagsterAirbyteTranslator()
+
+    return multi_asset(
+        name=name,
+        group_name=group_name,
+        can_subset=True,
+        specs=[
+            spec
+            for spec in workspace.load_asset_specs(
+                dagster_airbyte_translator=dagster_airbyte_translator
+            )
+            if AirbyteMetadataSet.extract(spec.metadata).connection_id == connection_id
+        ],
+    )

{dagster-airbyte-0.25.6 → dagster-airbyte-0.25.7}/dagster_airbyte/asset_defs.py

@@ -1,7 +1,6 @@
 import hashlib
 import inspect
 import os
-import re
 from abc import abstractmethod
 from functools import partial
 from itertools import chain
@@ -23,6 +22,7 @@ from typing import (
 
 import yaml
 from dagster import (
+    AssetExecutionContext,
     AssetKey,
     AssetOut,
     AutoMaterializePolicy,
@@ -33,6 +33,7 @@ from dagster import (
     SourceAsset,
     _check as check,
 )
+from dagster._annotations import experimental
 from dagster._core.definitions import AssetsDefinition, multi_asset
 from dagster._core.definitions.cacheable_assets import (
     AssetsDefinitionCacheableData,
@@ -45,9 +46,17 @@ from dagster._core.errors import DagsterInvalidDefinitionError, DagsterInvalidIn
 from dagster._core.execution.context.init import build_init_resource_context
 from dagster._utils.merger import merge_dicts
 
-from dagster_airbyte.resources import AirbyteCloudResource, AirbyteResource, BaseAirbyteResource
+from dagster_airbyte.asset_decorator import airbyte_assets
+from dagster_airbyte.resources import (
+    AirbyteCloudResource,
+    AirbyteCloudWorkspace,
+    AirbyteResource,
+    BaseAirbyteResource,
+)
+from dagster_airbyte.translator import AirbyteMetadataSet, DagsterAirbyteTranslator
 from dagster_airbyte.types import AirbyteTableMetadata
 from dagster_airbyte.utils import (
+    clean_name,
     generate_materializations,
     generate_table_schema,
     is_basic_normalization_operation,
@@ -462,11 +471,6 @@ def _get_normalization_tables_for_schema(
     return out
 
 
-def _clean_name(name: str) -> str:
-    """Cleans an input to be a valid Dagster asset name."""
-    return re.sub(r"[^a-z0-9]+", "_", name.lower())
-
-
 class AirbyteConnectionMetadata(
     NamedTuple(
         "_AirbyteConnectionMetadata",
@@ -908,7 +912,7 @@ def load_assets_from_airbyte_instance(
     workspace_id: Optional[str] = None,
     key_prefix: Optional[CoercibleToAssetKeyPrefix] = None,
     create_assets_for_normalization_tables: bool = True,
-    connection_to_group_fn: Optional[Callable[[str], Optional[str]]] = _clean_name,
+    connection_to_group_fn: Optional[Callable[[str], Optional[str]]] = clean_name,
     connection_meta_to_group_fn: Optional[
         Callable[[AirbyteConnectionMetadata], Optional[str]]
     ] = None,
@@ -1013,7 +1017,7 @@ def load_assets_from_airbyte_instance(
     check.invariant(
         not connection_meta_to_group_fn
         or not connection_to_group_fn
-        or connection_to_group_fn == _clean_name,
+        or connection_to_group_fn == clean_name,
         "Cannot specify both connection_meta_to_group_fn and connection_to_group_fn",
     )
 
@@ -1032,3 +1036,117 @@ def load_assets_from_airbyte_instance(
         connection_to_freshness_policy_fn=connection_to_freshness_policy_fn,
         connection_to_auto_materialize_policy_fn=connection_to_auto_materialize_policy_fn,
     )
+
+
+# -----------------------
+# Reworked assets factory
+# -----------------------
+
+
+@experimental
+def build_airbyte_assets_definitions(
+    *,
+    workspace: AirbyteCloudWorkspace,
+    dagster_airbyte_translator: Optional[DagsterAirbyteTranslator] = None,
+) -> Sequence[AssetsDefinition]:
+    """The list of AssetsDefinition for all connections in the Airbyte workspace.
+
+    Args:
+        workspace (AirbyteCloudWorkspace): The Airbyte 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[AssetsDefinition]: The list of AssetsDefinition for all connections in the Airbyte workspace.
+
+    Examples:
+        Sync the tables of a Airbyte connection:
+
+        .. code-block:: python
+
+            from dagster_airbyte import AirbyteCloudWorkspace, build_airbyte_assets_definitions
+
+            import dagster as dg
+
+            airbyte_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_assets = build_airbyte_assets_definitions(workspace=workspace)
+
+            defs = dg.Definitions(
+                assets=airbyte_assets,
+                resources={"airbyte": airbyte_workspace},
+            )
+
+        Sync the tables of a Airbyte connection with a custom translator:
+
+        .. code-block:: python
+
+            from dagster_airbyte import (
+                DagsterAirbyteTranslator,
+                AirbyteConnectionTableProps,
+                AirbyteCloudWorkspace,
+                build_airbyte_assets_definitions
+            )
+
+            import dagster as dg
+
+            class CustomDagsterAirbyteTranslator(DagsterAirbyteTranslator):
+                def get_asset_spec(self, props: AirbyteConnectionTableProps) -> dg.AssetSpec:
+                    default_spec = super().get_asset_spec(props)
+                    return default_spec.merge_attributes(
+                        metadata={"custom": "metadata"},
+                    )
+
+            airbyte_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_assets = build_airbyte_assets_definitions(
+                workspace=workspace,
+                dagster_airbyte_translator=CustomDagsterAirbyteTranslator()
+            )
+
+            defs = dg.Definitions(
+                assets=airbyte_assets,
+                resources={"airbyte": airbyte_workspace},
+            )
+    """
+    dagster_airbyte_translator = dagster_airbyte_translator or DagsterAirbyteTranslator()
+
+    all_asset_specs = workspace.load_asset_specs(
+        dagster_airbyte_translator=dagster_airbyte_translator
+    )
+
+    connections = {
+        (
+            check.not_none(AirbyteMetadataSet.extract(spec.metadata).connection_id),
+            check.not_none(AirbyteMetadataSet.extract(spec.metadata).connection_name),
+        )
+        for spec in all_asset_specs
+    }
+
+    _asset_fns = []
+    for connection_id, connection_name in connections:
+
+        @airbyte_assets(
+            connection_id=connection_id,
+            workspace=workspace,
+            name=clean_name(connection_name),
+            group_name=clean_name(connection_name),
+            dagster_airbyte_translator=dagster_airbyte_translator,
+        )
+        def _asset_fn(context: AssetExecutionContext, airbyte: AirbyteCloudWorkspace):
+            yield from airbyte.sync_and_poll(context=context)
+
+        _asset_fns.append(_asset_fn)
+
+    return _asset_fns

{dagster-airbyte-0.25.6 → dagster-airbyte-0.25.7}/dagster_airbyte/managed/reconciliation.py

@@ -36,7 +36,6 @@ from dagster_managed_elements.utils import UNSET, diff_dicts
 from dagster_airbyte.asset_defs import (
     AirbyteConnectionMetadata,
     AirbyteInstanceCacheableAssetsDefinition,
-    _clean_name,
 )
 from dagster_airbyte.managed.types import (
     MANAGED_ELEMENTS_DEPRECATION_MSG,
@@ -50,7 +49,7 @@ from dagster_airbyte.managed.types import (
     InitializedAirbyteSource,
 )
 from dagster_airbyte.resources import AirbyteResource
-from dagster_airbyte.utils import is_basic_normalization_operation
+from dagster_airbyte.utils import clean_name, is_basic_normalization_operation
 
 
 def gen_configured_stream_json(
@@ -746,7 +745,7 @@ def load_assets_from_connections(
     connections: Iterable[AirbyteConnection],
     key_prefix: Optional[CoercibleToAssetKeyPrefix] = None,
     create_assets_for_normalization_tables: bool = True,
-    connection_to_group_fn: Optional[Callable[[str], Optional[str]]] = _clean_name,
+    connection_to_group_fn: Optional[Callable[[str], Optional[str]]] = clean_name,
     connection_meta_to_group_fn: Optional[
         Callable[[AirbyteConnectionMetadata], Optional[str]]
     ] = None,
@@ -821,7 +820,7 @@ def load_assets_from_connections(
     check.invariant(
         not connection_meta_to_group_fn
         or not connection_to_group_fn
-        or connection_to_group_fn == _clean_name,
+        or connection_to_group_fn == clean_name,
         "Cannot specify both connection_meta_to_group_fn and connection_to_group_fn",
     )
 

{dagster-airbyte-0.25.6 → dagster-airbyte-0.25.7}/dagster_airbyte/resources.py

@@ -10,15 +10,18 @@ from typing import Any, Dict, List, Mapping, Optional, Sequence, cast
 
 import requests
 from dagster import (
+    AssetExecutionContext,
+    AssetMaterialization,
     ConfigurableResource,
     Definitions,
     Failure,
     InitResourceContext,
+    MaterializeResult,
     _check as check,
     get_dagster_logger,
     resource,
 )
-from dagster._annotations import experimental
+from dagster._annotations import experimental, public
 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
@@ -32,11 +35,20 @@ from requests.exceptions import RequestException
 
 from dagster_airbyte.translator import (
     AirbyteConnection,
+    AirbyteConnectionTableProps,
     AirbyteDestination,
+    AirbyteJob,
+    AirbyteJobStatusType,
+    AirbyteMetadataSet,
     AirbyteWorkspaceData,
     DagsterAirbyteTranslator,
 )
 from dagster_airbyte.types import AirbyteOutput
+from dagster_airbyte.utils import (
+    DAGSTER_AIRBYTE_TRANSLATOR_METADATA_KEY,
+    get_airbyte_connection_table_name,
+    get_translator_from_airbyte_assets,
+)
 
 AIRBYTE_REST_API_BASE = "https://api.airbyte.com"
 AIRBYTE_REST_API_VERSION = "v1"
@@ -53,16 +65,6 @@ AIRBYTE_CLOUD_REFRESH_TIMEDELTA_SECONDS = 150
 AIRBYTE_RECONSTRUCTION_METADATA_KEY_PREFIX = "dagster-airbyte/reconstruction_metadata"
 
 
-class AirbyteState:
-    RUNNING = "running"
-    SUCCEEDED = "succeeded"
-    CANCELLED = "cancelled"
-    PENDING = "pending"
-    FAILED = "failed"
-    ERROR = "error"
-    INCOMPLETE = "incomplete"
-
-
 class AirbyteResourceState:
     def __init__(self) -> None:
         self.request_cache: Dict[str, Optional[Mapping[str, object]]] = {}
@@ -252,13 +254,17 @@ class BaseAirbyteResource(ConfigurableResource):
                 job_info = cast(Dict[str, object], job_details.get("job", {}))
                 state = job_info.get("status")
 
-                if state in (AirbyteState.RUNNING, AirbyteState.PENDING, AirbyteState.INCOMPLETE):
+                if state in (
+                    AirbyteJobStatusType.RUNNING,
+                    AirbyteJobStatusType.PENDING,
+                    AirbyteJobStatusType.INCOMPLETE,
+                ):
                     continue
-                elif state == AirbyteState.SUCCEEDED:
+                elif state == AirbyteJobStatusType.SUCCEEDED:
                     break
-                elif state == AirbyteState.ERROR:
+                elif state == AirbyteJobStatusType.ERROR:
                     raise Failure(f"Job failed: {job_id}")
-                elif state == AirbyteState.CANCELLED:
+                elif state == AirbyteJobStatusType.CANCELLED:
                     raise Failure(f"Job was cancelled: {job_id}")
                 else:
                     raise Failure(f"Encountered unexpected state `{state}` for job_id {job_id}")
@@ -266,7 +272,12 @@ class BaseAirbyteResource(ConfigurableResource):
             # if Airbyte sync has not completed, make sure to cancel it so that it doesn't outlive
             # the python process
             if (
-                state not in (AirbyteState.SUCCEEDED, AirbyteState.ERROR, AirbyteState.CANCELLED)
+                state
+                not in (
+                    AirbyteJobStatusType.SUCCEEDED,
+                    AirbyteJobStatusType.ERROR,
+                    AirbyteJobStatusType.CANCELLED,
+                )
                 and self.cancel_sync_on_run_termination
             ):
                 self.cancel_job(job_id)
@@ -742,13 +753,17 @@ class AirbyteResource(BaseAirbyteResource):
                 job_info = cast(Dict[str, object], job_details.get("job", {}))
                 state = job_info.get("status")
 
-                if state in (AirbyteState.RUNNING, AirbyteState.PENDING, AirbyteState.INCOMPLETE):
+                if state in (
+                    AirbyteJobStatusType.RUNNING,
+                    AirbyteJobStatusType.PENDING,
+                    AirbyteJobStatusType.INCOMPLETE,
+                ):
                     continue
-                elif state == AirbyteState.SUCCEEDED:
+                elif state == AirbyteJobStatusType.SUCCEEDED:
                     break
-                elif state == AirbyteState.ERROR:
+                elif state == AirbyteJobStatusType.ERROR:
                     raise Failure(f"Job failed: {job_id}")
-                elif state == AirbyteState.CANCELLED:
+                elif state == AirbyteJobStatusType.CANCELLED:
                     raise Failure(f"Job was cancelled: {job_id}")
                 else:
                     raise Failure(f"Encountered unexpected state `{state}` for job_id {job_id}")
@@ -756,7 +771,12 @@ class AirbyteResource(BaseAirbyteResource):
             # if Airbyte sync has not completed, make sure to cancel it so that it doesn't outlive
             # the python process
             if (
-                state not in (AirbyteState.SUCCEEDED, AirbyteState.ERROR, AirbyteState.CANCELLED)
+                state
+                not in (
+                    AirbyteJobStatusType.SUCCEEDED,
+                    AirbyteJobStatusType.ERROR,
+                    AirbyteJobStatusType.CANCELLED,
+                )
                 and self.cancel_sync_on_run_termination
             ):
                 self.cancel_job(job_id)
@@ -992,6 +1012,101 @@ class AirbyteCloudClient(DagsterModel):
             base_url=self.rest_api_base_url,
         )
 
+    def start_sync_job(self, connection_id: str) -> Mapping[str, Any]:
+        return self._make_request(
+            method="POST",
+            endpoint="jobs",
+            base_url=self.rest_api_base_url,
+            data={
+                "connectionId": connection_id,
+                "jobType": "sync",
+            },
+        )
+
+    def get_job_details(self, job_id: int) -> Mapping[str, Any]:
+        return self._make_request(
+            method="GET", endpoint=f"jobs/{job_id}", base_url=self.rest_api_base_url
+        )
+
+    def cancel_job(self, job_id: int) -> Mapping[str, Any]:
+        return self._make_request(
+            method="DELETE", endpoint=f"jobs/{job_id}", base_url=self.rest_api_base_url
+        )
+
+    def sync_and_poll(
+        self,
+        connection_id: str,
+        poll_interval: Optional[float] = None,
+        poll_timeout: Optional[float] = None,
+        cancel_on_termination: bool = True,
+    ) -> AirbyteOutput:
+        """Initializes a sync operation for the given connection, and polls until it completes.
+
+        Args:
+            connection_id (str): The Airbyte Connection ID. You can retrieve this value from the
+                "Connection" tab of a given connection in the Airbyte UI.
+            poll_interval (float): The time (in seconds) that will be waited between successive polls.
+            poll_timeout (float): The maximum time that will wait before this operation is timed
+                out. By default, this will never time out.
+            cancel_on_termination (bool): Whether to cancel a sync in Airbyte if the Dagster runner is terminated.
+                This may be useful to disable if using Airbyte sources that cannot be cancelled and
+                resumed easily, or if your Dagster deployment may experience runner interruptions
+                that do not impact your Airbyte deployment.
+
+        Returns:
+            :py:class:`~AirbyteOutput`:
+                Details of the sync job.
+        """
+        connection_details = self.get_connection_details(connection_id)
+        start_job_details = self.start_sync_job(connection_id)
+        job = AirbyteJob.from_job_details(job_details=start_job_details)
+
+        self._log.info(f"Job {job.id} initialized for connection_id={connection_id}.")
+        poll_start = datetime.now()
+        poll_interval = (
+            poll_interval if poll_interval is not None else DEFAULT_POLL_INTERVAL_SECONDS
+        )
+        try:
+            while True:
+                if poll_timeout and datetime.now() > poll_start + timedelta(seconds=poll_timeout):
+                    raise Failure(
+                        f"Timeout: Airbyte job {job.id} is not ready after the timeout"
+                        f" {poll_timeout} seconds"
+                    )
+
+                time.sleep(poll_interval)
+                # We return these job details in the AirbyteOutput when the job succeeds
+                poll_job_details = self.get_job_details(job.id)
+                job = AirbyteJob.from_job_details(job_details=poll_job_details)
+                if job.status in (
+                    AirbyteJobStatusType.RUNNING,
+                    AirbyteJobStatusType.PENDING,
+                    AirbyteJobStatusType.INCOMPLETE,
+                ):
+                    continue
+                elif job.status == AirbyteJobStatusType.SUCCEEDED:
+                    break
+                elif job.status in [AirbyteJobStatusType.ERROR, AirbyteJobStatusType.FAILED]:
+                    raise Failure(f"Job failed: {job.id}")
+                elif job.status == AirbyteJobStatusType.CANCELLED:
+                    raise Failure(f"Job was cancelled: {job.id}")
+                else:
+                    raise Failure(
+                        f"Encountered unexpected state `{job.status}` for job_id {job.id}"
+                    )
+        finally:
+            # if Airbyte sync has not completed, make sure to cancel it so that it doesn't outlive
+            # the python process
+            if cancel_on_termination and job.status not in (
+                AirbyteJobStatusType.SUCCEEDED,
+                AirbyteJobStatusType.ERROR,
+                AirbyteJobStatusType.CANCELLED,
+                AirbyteJobStatusType.FAILED,
+            ):
+                self.cancel_job(job.id)
+
+        return AirbyteOutput(job_details=poll_job_details, connection_details=connection_details)
+
 
 @experimental
 class AirbyteCloudWorkspace(ConfigurableResource):
@@ -1067,6 +1182,130 @@ class AirbyteCloudWorkspace(ConfigurableResource):
             destinations_by_id=destinations_by_id,
         )
 
+    @cached_method
+    def load_asset_specs(
+        self,
+        dagster_airbyte_translator: Optional[DagsterAirbyteTranslator] = None,
+    ) -> Sequence[AssetSpec]:
+        """Returns a list of AssetSpecs representing the Airbyte content in the workspace.
+
+        Args:
+            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 workspace:
+            .. code-block:: python
+
+                from dagster_airbyte import AirbyteCloudWorkspace
+
+                import dagster as dg
+
+                airbyte_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_specs = airbyte_workspace.load_asset_specs()
+                defs = dg.Definitions(assets=airbyte_specs, resources={"airbyte": airbyte_workspace}
+        """
+        dagster_airbyte_translator = dagster_airbyte_translator or DagsterAirbyteTranslator()
+
+        return load_airbyte_cloud_asset_specs(
+            workspace=self, dagster_airbyte_translator=dagster_airbyte_translator
+        )
+
+    def _generate_materialization(
+        self,
+        airbyte_output: AirbyteOutput,
+        dagster_airbyte_translator: DagsterAirbyteTranslator,
+    ):
+        connection = AirbyteConnection.from_connection_details(
+            connection_details=airbyte_output.connection_details
+        )
+
+        for stream in connection.streams.values():
+            if stream.selected:
+                connection_table_name = get_airbyte_connection_table_name(
+                    stream_prefix=connection.stream_prefix,
+                    stream_name=stream.name,
+                )
+                stream_asset_spec = dagster_airbyte_translator.get_asset_spec(
+                    props=AirbyteConnectionTableProps(
+                        table_name=connection_table_name,
+                        stream_prefix=connection.stream_prefix,
+                        stream_name=stream.name,
+                        json_schema=stream.json_schema,
+                        connection_id=connection.id,
+                        connection_name=connection.name,
+                        destination_type=None,
+                        database=None,
+                        schema=None,
+                    )
+                )
+
+                yield AssetMaterialization(
+                    asset_key=stream_asset_spec.key,
+                    description=(
+                        f"Table generated via Airbyte Cloud sync "
+                        f"for connection {connection.name}: {connection_table_name}"
+                    ),
+                    metadata=stream_asset_spec.metadata,
+                )
+
+    @public
+    @experimental
+    def sync_and_poll(self, context: AssetExecutionContext):
+        """Executes a sync and poll process to materialize Airbyte Cloud assets.
+            This method can only be used in the context of an asset execution.
+
+        Args:
+            context (AssetExecutionContext): The execution context
+                from within `@airbyte_assets`.
+
+        Returns:
+            Iterator[Union[AssetMaterialization, MaterializeResult]]: An iterator of MaterializeResult
+                or AssetMaterialization.
+        """
+        assets_def = context.assets_def
+        dagster_airbyte_translator = get_translator_from_airbyte_assets(assets_def)
+        connection_id = next(
+            check.not_none(AirbyteMetadataSet.extract(spec.metadata).connection_id)
+            for spec in assets_def.specs
+        )
+
+        client = self.get_client()
+        airbyte_output = client.sync_and_poll(
+            connection_id=connection_id,
+        )
+
+        materialized_asset_keys = set()
+        for materialization in self._generate_materialization(
+            airbyte_output=airbyte_output, dagster_airbyte_translator=dagster_airbyte_translator
+        ):
+            # Scan through all tables actually created, if it was expected then emit a MaterializeResult.
+            # Otherwise, emit a runtime AssetMaterialization.
+            if materialization.asset_key in context.selected_asset_keys:
+                yield MaterializeResult(
+                    asset_key=materialization.asset_key, metadata=materialization.metadata
+                )
+                materialized_asset_keys.add(materialization.asset_key)
+            else:
+                context.log.warning(
+                    f"An unexpected asset was materialized: {materialization.asset_key}. "
+                    f"Yielding a materialization event."
+                )
+                yield materialization
+
+        unmaterialized_asset_keys = context.selected_asset_keys - materialized_asset_keys
+        if unmaterialized_asset_keys:
+            context.log.warning(f"Assets were not materialized: {unmaterialized_asset_keys}")
+
 
 @experimental
 def load_airbyte_cloud_asset_specs(
@@ -1103,16 +1342,23 @@ def load_airbyte_cloud_asset_specs(
             airbyte_cloud_specs = load_airbyte_cloud_asset_specs(airbyte_cloud_workspace)
             defs = dg.Definitions(assets=airbyte_cloud_specs)
     """
+    dagster_airbyte_translator = dagster_airbyte_translator or DagsterAirbyteTranslator()
+
     with workspace.process_config_and_initialize_cm() as initialized_workspace:
-        return check.is_list(
-            AirbyteCloudWorkspaceDefsLoader(
-                workspace=initialized_workspace,
-                translator=dagster_airbyte_translator or DagsterAirbyteTranslator(),
+        return [
+            spec.merge_attributes(
+                metadata={DAGSTER_AIRBYTE_TRANSLATOR_METADATA_KEY: dagster_airbyte_translator}
             )
-            .build_defs()
-            .assets,
-            AssetSpec,
-        )
+            for spec in check.is_list(
+                AirbyteCloudWorkspaceDefsLoader(
+                    workspace=initialized_workspace,
+                    translator=dagster_airbyte_translator,
+                )
+                .build_defs()
+                .assets,
+                AssetSpec,
+            )
+        ]
 
 
 @record

{dagster-airbyte-0.25.6 → dagster-airbyte-0.25.7}/dagster_airbyte/translator.py

@@ -1,6 +1,7 @@
+from enum import Enum
 from typing import Any, List, Mapping, Optional, Sequence
 
-from dagster._annotations import experimental
+from dagster._annotations import deprecated, 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
@@ -11,6 +12,27 @@ from dagster._utils.cached_method import cached_method
 from dagster_airbyte.utils import generate_table_schema, get_airbyte_connection_table_name
 
 
+class AirbyteJobStatusType(str, Enum):
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    CANCELLED = "cancelled"
+    PENDING = "pending"
+    FAILED = "failed"
+    ERROR = "error"
+    INCOMPLETE = "incomplete"
+
+
+@deprecated(breaking_version="1.10", additional_warn_text="Use `AirbyteJobStatusType` instead.")
+class AirbyteState:
+    RUNNING = AirbyteJobStatusType.RUNNING
+    SUCCEEDED = AirbyteJobStatusType.SUCCEEDED
+    CANCELLED = AirbyteJobStatusType.CANCELLED
+    PENDING = AirbyteJobStatusType.PENDING
+    FAILED = AirbyteJobStatusType.FAILED
+    ERROR = AirbyteJobStatusType.ERROR
+    INCOMPLETE = AirbyteJobStatusType.INCOMPLETE
+
+
 @record
 class AirbyteConnectionTableProps:
     table_name: str
@@ -19,7 +41,7 @@ class AirbyteConnectionTableProps:
     json_schema: Mapping[str, Any]
     connection_id: str
     connection_name: str
-    destination_type: str
+    destination_type: Optional[str]
     database: Optional[str]
     schema: Optional[str]
 
@@ -108,6 +130,25 @@ class AirbyteStream:
         )
 
 
+@whitelist_for_serdes
+@record
+class AirbyteJob:
+    """Represents an Airbyte job, based on data as returned from the API."""
+
+    id: int
+    status: str
+
+    @classmethod
+    def from_job_details(
+        cls,
+        job_details: Mapping[str, Any],
+    ) -> "AirbyteJob":
+        return cls(
+            id=job_details["jobId"],
+            status=job_details["status"],
+        )
+
+
 @whitelist_for_serdes
 @record
 class AirbyteWorkspaceData:
@@ -190,5 +231,5 @@ class DagsterAirbyteTranslator:
         return AssetSpec(
             key=AssetKey(props.table_name),
             metadata=metadata,
-            kinds={"airbyte", props.destination_type},
+            kinds={"airbyte", *({props.destination_type} if props.destination_type else set())},
         )

{dagster-airbyte-0.25.6 → dagster-airbyte-0.25.7}/dagster_airbyte/utils.py

@@ -1,10 +1,26 @@
-from typing import Any, Iterator, Mapping, Optional, Sequence
+import re
+from typing import TYPE_CHECKING, Any, Iterator, Mapping, Optional, Sequence
 
-from dagster import AssetMaterialization, MetadataValue
+from dagster import (
+    AssetMaterialization,
+    AssetsDefinition,
+    DagsterInvariantViolationError,
+    MetadataValue,
+)
 from dagster._core.definitions.metadata.table import TableColumn, TableSchema
 
 from dagster_airbyte.types import AirbyteOutput
 
+if TYPE_CHECKING:
+    from dagster_airbyte import DagsterAirbyteTranslator
+
+DAGSTER_AIRBYTE_TRANSLATOR_METADATA_KEY = "dagster-airbyte/dagster_airbyte_translator"
+
+
+def clean_name(name: str) -> str:
+    """Cleans an input to be a valid Dagster asset name."""
+    return re.sub(r"[^a-z0-9]+", "_", name.lower())
+
 
 def get_airbyte_connection_table_name(stream_prefix: Optional[str], stream_name: str) -> str:
     return f"{stream_prefix if stream_prefix else ''}{stream_name}"
@@ -78,3 +94,18 @@ def generate_materializations(
             all_stream_stats.get(stream_name, {}),
             asset_key_prefix=asset_key_prefix,
         )
+
+
+def get_translator_from_airbyte_assets(
+    airbyte_assets: AssetsDefinition,
+) -> "DagsterAirbyteTranslator":
+    metadata_by_key = airbyte_assets.metadata_by_key or {}
+    first_asset_key = next(iter(airbyte_assets.metadata_by_key.keys()))
+    first_metadata = metadata_by_key.get(first_asset_key, {})
+    dagster_airbyte_translator = first_metadata.get(DAGSTER_AIRBYTE_TRANSLATOR_METADATA_KEY)
+    if dagster_airbyte_translator is None:
+        raise DagsterInvariantViolationError(
+            f"Expected to find airbyte translator metadata on asset {first_asset_key.to_user_string()},"
+            " but did not. Did you pass in assets that weren't generated by @airbyte_assets?"
+        )
+    return dagster_airbyte_translator

dagster-airbyte-0.25.7/dagster_airbyte/version.py

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

{dagster-airbyte-0.25.6 → dagster-airbyte-0.25.7/dagster_airbyte.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dagster-airbyte
-Version: 0.25.6
+Version: 0.25.7
 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.6 → dagster-airbyte-0.25.7}/dagster_airbyte.egg-info/SOURCES.txt

@@ -4,6 +4,7 @@ README.md
 setup.cfg
 setup.py
 dagster_airbyte/__init__.py
+dagster_airbyte/asset_decorator.py
 dagster_airbyte/asset_defs.py
 dagster_airbyte/cli.py
 dagster_airbyte/ops.py

{dagster-airbyte-0.25.6 → dagster-airbyte-0.25.7}/dagster_airbyte.egg-info/requires.txt

@@ -1,8 +1,8 @@
-dagster==1.9.6
+dagster==1.9.7
 requests
 
 [managed]
-dagster-managed-elements==0.25.6
+dagster-managed-elements==0.25.7
 
 [test]
 requests-mock

{dagster-airbyte-0.25.6 → dagster-airbyte-0.25.7}/setup.py

@@ -37,7 +37,7 @@ setup(
     include_package_data=True,
     python_requires=">=3.9,<3.13",
     install_requires=[
-        "dagster==1.9.6",
+        "dagster==1.9.7",
         "requests",
     ],
     zip_safe=False,
@@ -52,7 +52,7 @@ setup(
             "flaky",
         ],
         "managed": [
-            "dagster-managed-elements==0.25.6",
+            "dagster-managed-elements==0.25.7",
         ],
     },
 )

dagster-airbyte-0.25.6/dagster_airbyte/version.py

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