dagster-airbyte 0.24.3__py3-none-any.whl → 0.28.3__py3-none-any.whl

dagster_airbyte/__init__.py

@@ -1,7 +1,12 @@
-from dagster._core.libraries import DagsterLibraryRegistry
+from dagster_shared.libraries import DagsterLibraryRegistry
+
+from dagster_airbyte.components.workspace_component.component import (
+    AirbyteCloudWorkspaceComponent as AirbyteCloudWorkspaceComponent,
+    AirbyteWorkspaceComponent as AirbyteWorkspaceComponent,
+)
 
 try:
-    from .managed import (
+    from dagster_airbyte.managed import (
         AirbyteConnection as AirbyteConnection,
         AirbyteDestination as AirbyteDestination,
         AirbyteDestinationNamespace as AirbyteDestinationNamespace,
@@ -14,20 +19,32 @@ try:
 except ImportError:
     pass
 
-from .asset_defs import (
+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,
-    load_assets_from_airbyte_project as load_assets_from_airbyte_project,
 )
-from .ops import airbyte_sync_op as airbyte_sync_op
-from .resources import (
+from dagster_airbyte.legacy_resources import (
     AirbyteCloudResource as AirbyteCloudResource,
     AirbyteResource as AirbyteResource,
-    AirbyteState as AirbyteState,
     airbyte_cloud_resource as airbyte_cloud_resource,
     airbyte_resource as airbyte_resource,
 )
-from .types import AirbyteOutput as AirbyteOutput
-from .version import __version__ as __version__
+from dagster_airbyte.ops import airbyte_sync_op as airbyte_sync_op
+from dagster_airbyte.resources import (
+    AirbyteCloudWorkspace as AirbyteCloudWorkspace,
+    AirbyteWorkspace as AirbyteWorkspace,
+    load_airbyte_asset_specs as load_airbyte_asset_specs,
+    load_airbyte_cloud_asset_specs as load_airbyte_cloud_asset_specs,
+)
+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__
 
 DagsterLibraryRegistry.register("dagster-airbyte", __version__)

dagster_airbyte/asset_decorator.py

@@ -0,0 +1,123 @@
+from collections.abc import Callable
+from typing import Any, Optional, Union
+
+from dagster import AssetsDefinition, multi_asset
+from dagster._annotations import beta
+from dagster._core.errors import DagsterInvariantViolationError
+
+from dagster_airbyte.resources import AirbyteCloudWorkspace, AirbyteWorkspace
+from dagster_airbyte.translator import AirbyteMetadataSet, DagsterAirbyteTranslator
+
+
+@beta
+def airbyte_assets(
+    *,
+    connection_id: str,
+    workspace: Union[AirbyteWorkspace, 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 (Union[AirbyteWorkspace, 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()
+
+    specs = [
+        spec
+        for spec in workspace.load_asset_specs(
+            dagster_airbyte_translator=dagster_airbyte_translator
+        )
+        if AirbyteMetadataSet.extract(spec.metadata).connection_id == connection_id
+    ]
+
+    if any([spec for spec in specs if spec.group_name]) and group_name:
+        raise DagsterInvariantViolationError(
+            f"Cannot set group_name parameter on airbyte_assets with connection ID {connection_id} - "
+            f"one or more of the Airbyte asset specs have a group_name defined."
+        )
+
+    return multi_asset(
+        name=name,
+        group_name=group_name,
+        can_subset=True,
+        specs=specs,
+    )