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

dagster_airbyte/__init__.py

@@ -1,5 +1,10 @@
 from dagster_shared.libraries import DagsterLibraryRegistry
 
+from dagster_airbyte.components.workspace_component.component import (
+    AirbyteCloudWorkspaceComponent as AirbyteCloudWorkspaceComponent,
+    AirbyteWorkspaceComponent as AirbyteWorkspaceComponent,
+)
+
 try:
     from dagster_airbyte.managed import (
         AirbyteConnection as AirbyteConnection,
@@ -20,13 +25,17 @@ from dagster_airbyte.asset_defs import (
     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 (
+from dagster_airbyte.legacy_resources import (
     AirbyteCloudResource as AirbyteCloudResource,
-    AirbyteCloudWorkspace as AirbyteCloudWorkspace,
     AirbyteResource as AirbyteResource,
     airbyte_cloud_resource as airbyte_cloud_resource,
     airbyte_resource as airbyte_resource,
+)
+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 (

dagster_airbyte/asset_decorator.py

@@ -1,10 +1,11 @@
-from typing import Any, Callable, Optional
+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
+from dagster_airbyte.resources import AirbyteCloudWorkspace, AirbyteWorkspace
 from dagster_airbyte.translator import AirbyteMetadataSet, DagsterAirbyteTranslator
 
 
@@ -12,7 +13,7 @@ from dagster_airbyte.translator import AirbyteMetadataSet, DagsterAirbyteTransla
 def airbyte_assets(
     *,
     connection_id: str,
-    workspace: AirbyteCloudWorkspace,
+    workspace: Union[AirbyteWorkspace, AirbyteCloudWorkspace],
     name: Optional[str] = None,
     group_name: Optional[str] = None,
     dagster_airbyte_translator: Optional[DagsterAirbyteTranslator] = None,
@@ -21,7 +22,7 @@ def airbyte_assets(
 
     Args:
         connection_id (str): The Airbyte Connection ID.
-        workspace (AirbyteCloudWorkspace): The Airbyte workspace to fetch assets from.
+        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

dagster_airbyte/asset_defs.py

@@ -2,10 +2,10 @@ import hashlib
 import inspect
 import os
 from abc import abstractmethod
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from functools import partial
 from itertools import chain
-from typing import Any, Callable, NamedTuple, Optional, Union, cast
+from typing import Any, NamedTuple, Optional, Union, cast
 
 import yaml
 from dagster import (
@@ -13,16 +13,16 @@ from dagster import (
     AssetKey,
     AssetOut,
     AutoMaterializePolicy,
-    FreshnessPolicy,
+    LegacyFreshnessPolicy,
     Nothing,
     Output,
     ResourceDefinition,
     SourceAsset,
     _check as check,
 )
-from dagster._annotations import beta
+from dagster._annotations import beta, hidden_param, only_allow_hidden_params_in_kwargs, superseded
 from dagster._core.definitions import AssetsDefinition, multi_asset
-from dagster._core.definitions.cacheable_assets import (
+from dagster._core.definitions.assets.definition.cacheable_assets_definition import (
     AssetsDefinitionCacheableData,
     CacheableAssetsDefinition,
 )
@@ -34,13 +34,17 @@ from dagster._core.execution.context.init import build_init_resource_context
 from dagster._utils.merger import merge_dicts
 
 from dagster_airbyte.asset_decorator import airbyte_assets
-from dagster_airbyte.resources import (
+from dagster_airbyte.legacy_resources import (
     AirbyteCloudResource,
-    AirbyteCloudWorkspace,
     AirbyteResource,
     BaseAirbyteResource,
 )
-from dagster_airbyte.translator import AirbyteMetadataSet, DagsterAirbyteTranslator
+from dagster_airbyte.resources import AirbyteCloudWorkspace, AirbyteWorkspace, BaseAirbyteWorkspace
+from dagster_airbyte.translator import (
+    AirbyteConnection,
+    AirbyteMetadataSet,
+    DagsterAirbyteTranslator,
+)
 from dagster_airbyte.types import AirbyteTableMetadata
 from dagster_airbyte.utils import (
     clean_name,
@@ -68,7 +72,7 @@ def _build_airbyte_asset_defn_metadata(
     group_name: Optional[str] = None,
     io_manager_key: Optional[str] = None,
     schema_by_table_name: Optional[Mapping[str, TableSchema]] = None,
-    freshness_policy: Optional[FreshnessPolicy] = None,
+    legacy_freshness_policy: Optional[LegacyFreshnessPolicy] = None,
     auto_materialize_policy: Optional[AutoMaterializePolicy] = None,
 ) -> AssetsDefinitionCacheableData:
     asset_key_prefix = (
@@ -157,8 +161,10 @@ def _build_airbyte_asset_defn_metadata(
                 for table in tables
             }
         ),
-        freshness_policies_by_output_name=(
-            {output: freshness_policy for output in outputs} if freshness_policy else None
+        legacy_freshness_policies_by_output_name=(
+            {output: legacy_freshness_policy for output in outputs}
+            if legacy_freshness_policy
+            else None
         ),
         auto_materialize_policies_by_output_name=(
             {output: auto_materialize_policy for output in outputs}
@@ -198,9 +204,9 @@ def _build_airbyte_assets_from_metadata(
                     else None
                 ),
                 io_manager_key=io_manager_key,
-                freshness_policy=(
-                    assets_defn_meta.freshness_policies_by_output_name.get(k)
-                    if assets_defn_meta.freshness_policies_by_output_name
+                legacy_freshness_policy=(
+                    assets_defn_meta.legacy_freshness_policies_by_output_name.get(k)
+                    if assets_defn_meta.legacy_freshness_policies_by_output_name
                     else None
                 ),
                 dagster_type=Nothing,
@@ -245,6 +251,14 @@ def _build_airbyte_assets_from_metadata(
     return _assets
 
 
+@hidden_param(
+    param="legacy_freshness_policy",
+    breaking_version="1.13.0",
+)
+@hidden_param(
+    param="auto_materialize_policy",
+    breaking_version="1.10.0",
+)
 def build_airbyte_assets(
     connection_id: str,
     destination_tables: Sequence[str],
@@ -256,9 +270,8 @@ def build_airbyte_assets(
     deps: Optional[Iterable[Union[CoercibleToAssetKey, AssetsDefinition, SourceAsset]]] = 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,
-    auto_materialize_policy: Optional[AutoMaterializePolicy] = None,
+    **kwargs,
 ) -> Sequence[AssetsDefinition]:
     """Builds a set of assets representing the tables created by an Airbyte sync operation.
 
@@ -278,11 +291,13 @@ def build_airbyte_assets(
         deps (Optional[Sequence[Union[AssetsDefinition, SourceAsset, str, AssetKey]]]):
             A list of assets to add as sources.
         upstream_assets (Optional[Set[AssetKey]]): Deprecated, use deps instead. A list of assets to add as sources.
-        freshness_policy (Optional[FreshnessPolicy]): A freshness policy to apply to the assets
         stream_to_asset_map (Optional[Mapping[str, str]]): A mapping of an Airbyte stream name to a Dagster asset.
             This allows the use of the "prefix" setting in Airbyte with special characters that aren't valid asset names.
-        auto_materialize_policy (Optional[AutoMaterializePolicy]): An auto materialization policy to apply to the assets.
     """
+    only_allow_hidden_params_in_kwargs(build_airbyte_assets, kwargs)
+    legacy_freshness_policy = kwargs.get("legacy_freshness_policy")
+    auto_materialize_policy = kwargs.get("auto_materialize_policy")
+
     if upstream_assets is not None and deps is not None:
         raise DagsterInvalidDefinitionError(
             "Cannot specify both deps and upstream_assets to build_airbyte_assets. Use only deps"
@@ -325,7 +340,7 @@ def build_airbyte_assets(
                     ),
                 }
             ),
-            freshness_policy=freshness_policy,
+            legacy_freshness_policy=legacy_freshness_policy,
             auto_materialize_policy=auto_materialize_policy,
         )
         for table in tables
@@ -588,7 +603,7 @@ class AirbyteCoreCacheableAssetsDefinition(CacheableAssetsDefinition):
         connection_filter: Optional[Callable[[AirbyteConnectionMetadata], bool]],
         connection_to_asset_key_fn: Optional[Callable[[AirbyteConnectionMetadata, str], AssetKey]],
         connection_to_freshness_policy_fn: Optional[
-            Callable[[AirbyteConnectionMetadata], Optional[FreshnessPolicy]]
+            Callable[[AirbyteConnectionMetadata], Optional[LegacyFreshnessPolicy]]
         ],
         connection_to_auto_materialize_policy_fn: Optional[
             Callable[[AirbyteConnectionMetadata], Optional[AutoMaterializePolicy]]
@@ -671,7 +686,7 @@ class AirbyteCoreCacheableAssetsDefinition(CacheableAssetsDefinition):
                 ),
                 schema_by_table_name=schema_by_table_name,
                 table_to_asset_key_fn=table_to_asset_key,
-                freshness_policy=self._connection_to_freshness_policy_fn(connection),
+                legacy_freshness_policy=self._connection_to_freshness_policy_fn(connection),
                 auto_materialize_policy=self._connection_to_auto_materialize_policy_fn(connection),
             )
 
@@ -704,7 +719,7 @@ class AirbyteInstanceCacheableAssetsDefinition(AirbyteCoreCacheableAssetsDefinit
         connection_filter: Optional[Callable[[AirbyteConnectionMetadata], bool]],
         connection_to_asset_key_fn: Optional[Callable[[AirbyteConnectionMetadata, str], AssetKey]],
         connection_to_freshness_policy_fn: Optional[
-            Callable[[AirbyteConnectionMetadata], Optional[FreshnessPolicy]]
+            Callable[[AirbyteConnectionMetadata], Optional[LegacyFreshnessPolicy]]
         ],
         connection_to_auto_materialize_policy_fn: Optional[
             Callable[[AirbyteConnectionMetadata], Optional[AutoMaterializePolicy]]
@@ -818,7 +833,7 @@ class AirbyteYAMLCacheableAssetsDefinition(AirbyteCoreCacheableAssetsDefinition)
         connection_directories: Optional[Sequence[str]],
         connection_to_asset_key_fn: Optional[Callable[[AirbyteConnectionMetadata, str], AssetKey]],
         connection_to_freshness_policy_fn: Optional[
-            Callable[[AirbyteConnectionMetadata], Optional[FreshnessPolicy]]
+            Callable[[AirbyteConnectionMetadata], Optional[LegacyFreshnessPolicy]]
         ],
         connection_to_auto_materialize_policy_fn: Optional[
             Callable[[AirbyteConnectionMetadata], Optional[AutoMaterializePolicy]]
@@ -894,6 +909,11 @@ class AirbyteYAMLCacheableAssetsDefinition(AirbyteCoreCacheableAssetsDefinition)
         return output_connections
 
 
+@superseded(
+    additional_warn_text=(
+        "If you are using Airbyte 1.6.0 or higher, please see the migration guide: https://docs.dagster.io/integrations/libraries/airbyte/migration-guide"
+    )
+)
 def load_assets_from_airbyte_instance(
     airbyte: Union[AirbyteResource, ResourceDefinition],
     workspace_id: Optional[str] = None,
@@ -910,7 +930,7 @@ def load_assets_from_airbyte_instance(
         Callable[[AirbyteConnectionMetadata, str], AssetKey]
     ] = None,
     connection_to_freshness_policy_fn: Optional[
-        Callable[[AirbyteConnectionMetadata], Optional[FreshnessPolicy]]
+        Callable[[AirbyteConnectionMetadata], Optional[LegacyFreshnessPolicy]]
     ] = None,
     connection_to_auto_materialize_policy_fn: Optional[
         Callable[[AirbyteConnectionMetadata], Optional[AutoMaterializePolicy]]
@@ -1033,16 +1053,19 @@ def load_assets_from_airbyte_instance(
 @beta
 def build_airbyte_assets_definitions(
     *,
-    workspace: AirbyteCloudWorkspace,
+    workspace: Union[AirbyteWorkspace, AirbyteCloudWorkspace],
     dagster_airbyte_translator: Optional[DagsterAirbyteTranslator] = None,
+    connection_selector_fn: Optional[Callable[[AirbyteConnection], bool]] = None,
 ) -> Sequence[AssetsDefinition]:
     """The list of AssetsDefinition for all connections in the Airbyte workspace.
 
     Args:
-        workspace (AirbyteCloudWorkspace): The Airbyte workspace to fetch assets from.
+        workspace (Union[AirbyteWorkspace, 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`.
+        connection_selector_fn (Optional[Callable[[AirbyteConnection], bool]]): A function that allows for filtering
+            which Airbyte connection assets are created for.
 
     Returns:
         List[AssetsDefinition]: The list of AssetsDefinition for all connections in the Airbyte workspace.
@@ -1062,7 +1085,6 @@ def build_airbyte_assets_definitions(
                 client_secret=dg.EnvVar("AIRBYTE_CLOUD_CLIENT_SECRET"),
             )
 
-
             airbyte_assets = build_airbyte_assets_definitions(workspace=workspace)
 
             defs = dg.Definitions(
@@ -1096,21 +1118,46 @@ def build_airbyte_assets_definitions(
                 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},
+            )
+
+        Filter connections by name:
+
+        .. 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,
+                connection_selector_fn=lambda connection: connection.name in ["connection1", "connection2"]
+            )
+
             defs = dg.Definitions(
                 assets=airbyte_assets,
                 resources={"airbyte": airbyte_workspace},
             )
     """
     dagster_airbyte_translator = dagster_airbyte_translator or DagsterAirbyteTranslator()
+    connection_selector_fn = connection_selector_fn or (lambda connection: True)
 
     all_asset_specs = workspace.load_asset_specs(
-        dagster_airbyte_translator=dagster_airbyte_translator
+        dagster_airbyte_translator=dagster_airbyte_translator,
+        connection_selector_fn=connection_selector_fn,
     )
 
     connections = {
@@ -1127,10 +1174,10 @@ def build_airbyte_assets_definitions(
         @airbyte_assets(
             connection_id=connection_id,
             workspace=workspace,
-            name=clean_name(connection_name),
+            name=f"airbyte_{clean_name(connection_name)}",
             dagster_airbyte_translator=dagster_airbyte_translator,
         )
-        def _asset_fn(context: AssetExecutionContext, airbyte: AirbyteCloudWorkspace):
+        def _asset_fn(context: AssetExecutionContext, airbyte: BaseAirbyteWorkspace):
             yield from airbyte.sync_and_poll(context=context)
 
         _asset_fns.append(_asset_fn)