dagster-airbyte 0.25.1__py3-none-any.whl → 0.25.10__py3-none-any.whl

dagster_airbyte/__init__.py

@@ -14,17 +14,26 @@ 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
 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 (
+    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/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/asset_defs.py

@@ -1,28 +1,15 @@
 import hashlib
 import inspect
 import os
-import re
 from abc import abstractmethod
+from collections.abc import Iterable, Mapping, Sequence
 from functools import partial
 from itertools import chain
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Mapping,
-    NamedTuple,
-    Optional,
-    Sequence,
-    Set,
-    Tuple,
-    Union,
-    cast,
-)
+from typing import Any, Callable, NamedTuple, Optional, Union, cast
 
 import yaml
 from dagster import (
+    AssetExecutionContext,
     AssetKey,
     AssetOut,
     AutoMaterializePolicy,
@@ -33,6 +20,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 +33,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,
@@ -66,7 +62,7 @@ def _build_airbyte_asset_defn_metadata(
     destination_schema: Optional[str],
     table_to_asset_key_fn: Callable[[str], AssetKey],
     asset_key_prefix: Optional[Sequence[str]] = None,
-    normalization_tables: Optional[Mapping[str, Set[str]]] = None,
+    normalization_tables: Optional[Mapping[str, set[str]]] = None,
     normalization_raw_table_names_by_table: Optional[Mapping[str, str]] = None,
     upstream_assets: Optional[Iterable[AssetKey]] = None,
     group_name: Optional[str] = None,
@@ -96,7 +92,7 @@ def _build_airbyte_asset_defn_metadata(
         for table in tables
     }
 
-    internal_deps: Dict[str, Set[AssetKey]] = {}
+    internal_deps: dict[str, set[AssetKey]] = {}
 
     metadata_encodable_normalization_tables = (
         {k: list(v) for k, v in normalization_tables.items()} if normalization_tables else {}
@@ -115,7 +111,7 @@ def _build_airbyte_asset_defn_metadata(
     for table in destination_tables:
         internal_deps[table] = set(upstream_assets or [])
 
-    table_names: Dict[str, str] = {}
+    table_names: dict[str, str] = {}
     for table in destination_tables:
         if destination_database and destination_schema and table:
             # Use the destination raw table name to create the table name
@@ -186,8 +182,8 @@ def _build_airbyte_assets_from_metadata(
     metadata = cast(Mapping[str, Any], assets_defn_meta.extra_metadata)
     connection_id = cast(str, metadata["connection_id"])
     group_name = cast(Optional[str], metadata["group_name"])
-    destination_tables = cast(List[str], metadata["destination_tables"])
-    normalization_tables = cast(Mapping[str, List[str]], metadata["normalization_tables"])
+    destination_tables = cast(list[str], metadata["destination_tables"])
+    normalization_tables = cast(Mapping[str, list[str]], metadata["normalization_tables"])
     io_manager_key = cast(Optional[str], metadata["io_manager_key"])
 
     @multi_asset(
@@ -256,9 +252,9 @@ def build_airbyte_assets(
     destination_schema: Optional[str] = None,
     asset_key_prefix: Optional[Sequence[str]] = None,
     group_name: Optional[str] = None,
-    normalization_tables: Optional[Mapping[str, Set[str]]] = None,
+    normalization_tables: Optional[Mapping[str, set[str]]] = None,
     deps: Optional[Iterable[Union[CoercibleToAssetKey, AssetsDefinition, SourceAsset]]] = None,
-    upstream_assets: Optional[Set[AssetKey]] = None,
+    upstream_assets: Optional[set[AssetKey]] = None,
     schema_by_table_name: Optional[Mapping[str, TableSchema]] = None,
     freshness_policy: Optional[FreshnessPolicy] = None,
     stream_to_asset_map: Optional[Mapping[str, str]] = None,
@@ -301,7 +297,7 @@ def build_airbyte_assets(
         chain([destination_tables], normalization_tables.values() if normalization_tables else [])
     )
 
-    table_names: Dict[str, str] = {}
+    table_names: dict[str, str] = {}
     for table in destination_tables:
         if destination_database and destination_schema and table:
             table_names[table] = ".".join([destination_database, destination_schema, table])
@@ -429,7 +425,7 @@ def _get_normalization_tables_for_schema(
     For more information on Airbyte's normalization process, see:
     https://docs.airbyte.com/understanding-airbyte/basic-normalization/#nesting
     """
-    out: Dict[str, AirbyteTableMetadata] = {}
+    out: dict[str, AirbyteTableMetadata] = {}
     # Object types are broken into a new table, as long as they have children
 
     sub_schemas = _get_sub_schemas(schema)
@@ -462,11 +458,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",
@@ -474,7 +465,7 @@ class AirbyteConnectionMetadata(
             ("name", str),
             ("stream_prefix", str),
             ("has_basic_normalization", bool),
-            ("stream_data", List[Mapping[str, Any]]),
+            ("stream_data", list[Mapping[str, Any]]),
             ("destination", Mapping[str, Any]),
         ],
     )
@@ -533,7 +524,7 @@ class AirbyteConnectionMetadata(
         tables associated with each enabled stream and values representing any affiliated
         tables created by Airbyte's normalization process, if enabled.
         """
-        tables: Dict[str, AirbyteTableMetadata] = {}
+        tables: dict[str, AirbyteTableMetadata] = {}
 
         enabled_streams = [
             stream for stream in self.stream_data if stream.get("config", {}).get("selected", False)
@@ -548,7 +539,7 @@ class AirbyteConnectionMetadata(
                 if "json_schema" in stream["stream"]
                 else stream["stream"]["jsonSchema"]
             )
-            normalization_tables: Dict[str, AirbyteTableMetadata] = {}
+            normalization_tables: dict[str, AirbyteTableMetadata] = {}
             schema_props = schema.get("properties", schema.get("items", {}).get("properties", {}))
             if self.has_basic_normalization and return_normalization_tables:
                 for k, v in schema_props.items():
@@ -576,7 +567,7 @@ def _get_schema_by_table_name(
                 [
                     (k, v.schema)
                     for k, v in cast(
-                        Dict[str, AirbyteTableMetadata], meta.normalization_tables
+                        dict[str, AirbyteTableMetadata], meta.normalization_tables
                     ).items()
                 ]
                 for meta in stream_table_metadata.values()
@@ -627,11 +618,11 @@ class AirbyteCoreCacheableAssetsDefinition(CacheableAssetsDefinition):
         super().__init__(unique_id=f"airbyte-{contents.hexdigest()}")
 
     @abstractmethod
-    def _get_connections(self) -> Sequence[Tuple[str, AirbyteConnectionMetadata]]:
+    def _get_connections(self) -> Sequence[tuple[str, AirbyteConnectionMetadata]]:
         pass
 
     def compute_cacheable_data(self) -> Sequence[AssetsDefinitionCacheableData]:
-        asset_defn_data: List[AssetsDefinitionCacheableData] = []
+        asset_defn_data: list[AssetsDefinitionCacheableData] = []
         for connection_id, connection in self._get_connections():
             stream_table_metadata = connection.parse_stream_tables(
                 self._create_assets_for_normalization_tables
@@ -745,11 +736,11 @@ class AirbyteInstanceCacheableAssetsDefinition(AirbyteCoreCacheableAssetsDefinit
             )
             self._airbyte_instance: AirbyteResource = self._partially_initialized_airbyte_instance
 
-    def _get_connections(self) -> Sequence[Tuple[str, AirbyteConnectionMetadata]]:
+    def _get_connections(self) -> Sequence[tuple[str, AirbyteConnectionMetadata]]:
         workspace_id = self._workspace_id
         if not workspace_id:
             workspaces = cast(
-                List[Dict[str, Any]],
+                list[dict[str, Any]],
                 check.not_none(
                     self._airbyte_instance.make_request(endpoint="/workspaces/list", data={})
                 ).get("workspaces", []),
@@ -761,7 +752,7 @@ class AirbyteInstanceCacheableAssetsDefinition(AirbyteCoreCacheableAssetsDefinit
             workspace_id = workspaces[0].get("workspaceId")
 
         connections = cast(
-            List[Dict[str, Any]],
+            list[dict[str, Any]],
             check.not_none(
                 self._airbyte_instance.make_request(
                     endpoint="/connections/list", data={"workspaceId": workspace_id}
@@ -769,12 +760,12 @@ class AirbyteInstanceCacheableAssetsDefinition(AirbyteCoreCacheableAssetsDefinit
             ).get("connections", []),
         )
 
-        output_connections: List[Tuple[str, AirbyteConnectionMetadata]] = []
+        output_connections: list[tuple[str, AirbyteConnectionMetadata]] = []
         for connection_json in connections:
             connection_id = cast(str, connection_json.get("connectionId"))
 
             operations_json = cast(
-                Dict[str, Any],
+                dict[str, Any],
                 check.not_none(
                     self._airbyte_instance.make_request(
                         endpoint="/operations/list",
@@ -785,7 +776,7 @@ class AirbyteInstanceCacheableAssetsDefinition(AirbyteCoreCacheableAssetsDefinit
 
             destination_id = cast(str, connection_json.get("destinationId"))
             destination_json = cast(
-                Dict[str, Any],
+                dict[str, Any],
                 check.not_none(
                     self._airbyte_instance.make_request(
                         endpoint="/destinations/get",
@@ -847,10 +838,10 @@ class AirbyteYAMLCacheableAssetsDefinition(AirbyteCoreCacheableAssetsDefinition)
         self._project_dir = project_dir
         self._connection_directories = connection_directories
 
-    def _get_connections(self) -> Sequence[Tuple[str, AirbyteConnectionMetadata]]:
+    def _get_connections(self) -> Sequence[tuple[str, AirbyteConnectionMetadata]]:
         connections_dir = os.path.join(self._project_dir, "connections")
 
-        output_connections: List[Tuple[str, AirbyteConnectionMetadata]] = []
+        output_connections: list[tuple[str, AirbyteConnectionMetadata]] = []
 
         connection_directories = self._connection_directories or os.listdir(connections_dir)
         for connection_name in connection_directories:
@@ -908,7 +899,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 +1004,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 +1023,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