dagster-sling 0.28.8__py3-none-any.whl

dagster_sling/__init__.py

@@ -0,0 +1,24 @@
+from dagster_shared.libraries import DagsterLibraryRegistry
+
+from dagster_sling.asset_decorator import sling_assets
+from dagster_sling.components.sling_replication_collection.component import (
+    SlingReplicationCollectionComponent,
+    SlingReplicationSpecModel,
+)
+from dagster_sling.dagster_sling_translator import DagsterSlingTranslator
+from dagster_sling.resources import SlingConnectionResource, SlingMode, SlingResource
+from dagster_sling.sling_replication import SlingReplicationParam
+from dagster_sling.version import __version__
+
+__all__ = [
+    "DagsterSlingTranslator",
+    "SlingConnectionResource",
+    "SlingMode",
+    "SlingReplicationCollectionComponent",
+    "SlingReplicationParam",
+    "SlingReplicationSpecModel",
+    "SlingResource",
+    "sling_assets",
+]
+
+DagsterLibraryRegistry.register("dagster-sling", __version__)

dagster_sling/asset_decorator.py

@@ -0,0 +1,159 @@
+from collections.abc import Callable, Iterable, Mapping
+from copy import deepcopy
+from typing import Any, Optional
+
+from dagster import (
+    AssetsDefinition,
+    BackfillPolicy,
+    PartitionsDefinition,
+    _check as check,
+    multi_asset,
+)
+from dagster._core.definitions.assets.definition.asset_spec import AssetSpec
+from dagster._utils.merger import deep_merge_dicts
+from dagster._utils.security import non_secure_md5_hash_str
+
+from dagster_sling.dagster_sling_translator import DagsterSlingTranslator
+from dagster_sling.sling_replication import SlingReplicationParam, validate_replication
+
+METADATA_KEY_TRANSLATOR = "dagster_sling/dagster_sling_translator"
+METADATA_KEY_REPLICATION_CONFIG = "dagster_sling/sling_replication_config"
+
+
+def get_streams_from_replication(
+    replication_config: Mapping[str, Any],
+) -> Iterable[Mapping[str, Any]]:
+    """Returns a list of streams and their configs from a Sling replication config."""
+    for stream, config in replication_config.get("streams", {}).items():
+        if config and config.get("disabled", False):
+            continue
+        yield {"name": stream, "config": config}
+
+
+def streams_with_default_dagster_meta(
+    streams: Iterable[Mapping[str, Any]], replication_config: Mapping[str, Any]
+) -> Iterable[Mapping[str, Any]]:
+    """Ensures dagster meta configs in the `defaults` block of the replication_config are passed to
+    the assets definition object.
+    """
+    default_dagster_meta = replication_config.get("defaults", {}).get("meta", {}).get("dagster", {})
+    if not default_dagster_meta:
+        yield from streams
+    else:
+        for stream in streams:
+            name = stream["name"]
+            config = deepcopy(stream["config"])
+            if not config:
+                config = {"meta": {"dagster": default_dagster_meta}}
+            else:
+                config["meta"] = deep_merge_dicts(
+                    {"dagster": default_dagster_meta}, config.get("meta", {})
+                )
+            yield {"name": name, "config": config}
+
+
+def sling_assets(
+    *,
+    replication_config: SlingReplicationParam,
+    dagster_sling_translator: Optional[DagsterSlingTranslator] = None,
+    name: Optional[str] = None,
+    partitions_def: Optional[PartitionsDefinition] = None,
+    backfill_policy: Optional[BackfillPolicy] = None,
+    op_tags: Optional[Mapping[str, Any]] = None,
+    pool: Optional[str] = None,
+) -> Callable[[Callable[..., Any]], AssetsDefinition]:
+    """Create a definition for how to materialize a set of Sling replication streams as Dagster assets, as
+    described by a Sling replication config. This will create on Asset for every Sling target stream.
+
+    A Sling Replication config is a configuration that maps sources to destinations. For the full
+    spec and descriptions, see `Sling's Documentation <https://docs.slingdata.io/sling-cli/run/configuration>`_.
+
+    Args:
+        replication_config (Union[Mapping[str, Any], str, Path]): A path to a Sling replication config, or a dictionary
+            of a replication config.
+        dagster_sling_translator: (DagsterSlingTranslator): Allows customization of how to map a Sling stream to a Dagster
+          AssetKey.
+        name (Optional[str]: The name of the op.
+        partitions_def (Optional[PartitionsDefinition]): The partitions definition for this asset.
+        backfill_policy (Optional[BackfillPolicy]): The backfill policy for this asset.
+        op_tags (Optional[Mapping[str, Any]]): The tags for the underlying op.
+        pool (Optional[str]): A string that identifies the concurrency pool that governs the sling assets' execution.
+
+    Examples:
+        Running a sync by providing a path to a Sling Replication config:
+
+        .. code-block:: python
+
+            from dagster_sling import sling_assets, SlingResource, SlingConnectionResource
+
+            sling_resource = SlingResource(
+                connections=[
+                    SlingConnectionResource(
+                        name="MY_POSTGRES", type="postgres", connection_string=EnvVar("POSTGRES_URL")
+                    ),
+                    SlingConnectionResource(
+                        name="MY_DUCKDB",
+                        type="duckdb",
+                        connection_string="duckdb:///var/tmp/duckdb.db",
+                    ),
+                ]
+            )
+
+            config_path = "/path/to/replication.yaml"
+            @sling_assets(replication_config=config_path)
+            def my_assets(context, sling: SlingResource):
+                yield from sling.replicate(context=context)
+    """
+    return multi_asset(
+        name=name,
+        partitions_def=partitions_def,
+        can_subset=True,
+        op_tags=op_tags,
+        backfill_policy=backfill_policy,
+        specs=get_sling_asset_specs(replication_config, dagster_sling_translator, partitions_def),
+        pool=pool,
+    )
+
+
+def get_sling_asset_specs(
+    replication_config: SlingReplicationParam,
+    dagster_sling_translator: Optional[DagsterSlingTranslator] = None,
+    partitions_def: Optional[PartitionsDefinition] = None,
+) -> list[AssetSpec]:
+    replication_config = validate_replication(replication_config)
+
+    raw_streams = get_streams_from_replication(replication_config)
+
+    streams = streams_with_default_dagster_meta(raw_streams, replication_config)
+
+    code_version = non_secure_md5_hash_str(str(replication_config).encode())
+
+    dagster_sling_translator = (
+        check.opt_inst_param(
+            dagster_sling_translator, "dagster_sling_translator", DagsterSlingTranslator
+        )
+        or DagsterSlingTranslator()
+    )
+
+    def update_code_version_if_unset_by_translator(asset_spec: AssetSpec) -> AssetSpec:
+        if asset_spec.code_version is None:
+            return asset_spec.replace_attributes(code_version=code_version)
+        return asset_spec
+
+    base_specs = [
+        update_code_version_if_unset_by_translator(
+            dagster_sling_translator.get_asset_spec(stream).merge_attributes(
+                metadata={
+                    METADATA_KEY_TRANSLATOR: dagster_sling_translator,
+                    METADATA_KEY_REPLICATION_CONFIG: replication_config,
+                }
+            )
+        )
+        for stream in streams
+    ]
+    return [
+        spec.replace_attributes(
+            partitions_def=partitions_def or spec.partitions_def, skippable=True
+        )
+        for spec in base_specs
+    ]

dagster_sling/components/sling_replication_collection/component.py

@@ -0,0 +1,274 @@
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass, field
+from functools import cached_property
+from pathlib import Path
+from typing import Annotated, Any, Literal, Optional, TypeAlias, Union
+
+from dagster import Resolvable, Resolver
+from dagster._annotations import public
+from dagster._core.definitions.assets.definition.asset_spec import AssetSpec
+from dagster._core.definitions.assets.definition.assets_definition import AssetsDefinition
+from dagster._core.definitions.definitions_class import Definitions
+from dagster._core.definitions.events import AssetMaterialization
+from dagster._core.definitions.metadata.source_code import (
+    LocalFileCodeReference,
+    merge_code_references,
+)
+from dagster._core.definitions.result import MaterializeResult
+from dagster.components.component.component import Component
+from dagster.components.core.context import ComponentLoadContext
+from dagster.components.resolved.context import ResolutionContext
+from dagster.components.resolved.core_models import OpSpec
+from dagster.components.scaffold.scaffold import scaffold_with
+from dagster.components.utils.translation import (
+    ComponentTranslator,
+    TranslationFn,
+    TranslationFnResolver,
+    create_component_translator_cls,
+)
+from dagster_shared.utils.warnings import deprecation_warning
+from pydantic import BaseModel, ConfigDict, Field
+
+from dagster_sling.asset_decorator import sling_assets
+from dagster_sling.components.sling_replication_collection.scaffolder import (
+    SlingReplicationComponentScaffolder,
+)
+from dagster_sling.dagster_sling_translator import DagsterSlingTranslator
+from dagster_sling.resources import AssetExecutionContext, SlingConnectionResource, SlingResource
+
+SlingMetadataAddons: TypeAlias = Literal["column_metadata", "row_count"]
+
+
+@dataclass
+class SlingReplicationSpecModel(Resolvable):
+    path: str
+    op: Optional[OpSpec] = None
+    translation: Optional[
+        Annotated[
+            TranslationFn[Mapping[str, Any]],
+            TranslationFnResolver(
+                template_vars_for_translation_fn=lambda data: {"stream_definition": data}
+            ),
+        ]
+    ] = None
+    include_metadata: Annotated[
+        list[SlingMetadataAddons],
+        Resolver.default(
+            description="Optionally include additional metadata in materializations generated while executing your Sling models",
+            examples=[
+                ["row_count"],
+                ["row_count", "column_metadata"],
+            ],
+        ),
+    ] = field(default_factory=list)
+
+
+class SlingConnectionResourcePropertiesModel(Resolvable, BaseModel):
+    """Properties of a Sling connection resource."""
+
+    # each connection type supports a variety of different properties
+    model_config = ConfigDict(extra="allow")
+
+    type: str = Field(
+        description="Type of the source connection, must match the Sling connection types. Use 'file' for local storage."
+    )
+    connection_string: Optional[str] = Field(
+        description="The optional connection string for the source database, if not using keyword arguments.",
+        default=None,
+    )
+
+
+def resolve_connections(
+    context: ResolutionContext,
+    connections: Mapping[str, SlingConnectionResourcePropertiesModel],
+) -> list[SlingConnectionResource]:
+    return [
+        SlingConnectionResource(
+            name=name,
+            **context.resolve_value(
+                connection if isinstance(connection, dict) else connection.model_dump()
+            ),
+        )
+        for name, connection in connections.items()
+    ]
+
+
+ResolvedSlingConnections: TypeAlias = Annotated[
+    list[SlingConnectionResource],
+    Resolver(
+        resolve_connections, model_field_type=Mapping[str, SlingConnectionResourcePropertiesModel]
+    ),
+]
+
+
+def resolve_resource(
+    context: ResolutionContext,
+    sling,
+) -> Optional[SlingResource]:
+    if sling:
+        deprecation_warning(
+            "The `sling` field is deprecated, use `connections` instead. This field will be removed in a future release.",
+            "1.11.1",
+        )
+    return SlingResource(**context.resolve_value(sling.model_dump())) if sling else None
+
+
+@public
+@scaffold_with(SlingReplicationComponentScaffolder)
+@dataclass
+class SlingReplicationCollectionComponent(Component, Resolvable):
+    """Expose one or more Sling replications to Dagster as assets.
+
+    To get started, run:
+
+    ``dg scaffold defs dagster_sling.SlingReplicationCollectionComponent {defs_path}``
+
+    This will create a defs.yaml as well as a ``replication.yaml``, which is a Sling-specific configuration
+    file. See Sling's `documentation <https://docs.slingdata.io/concepts/replication#overview>`_ on ``replication.yaml``.
+    """
+
+    connections: ResolvedSlingConnections = field(default_factory=list)
+    replications: Sequence[SlingReplicationSpecModel] = field(default_factory=list)
+    resource: Annotated[
+        Optional[SlingResource],
+        Resolver(resolve_resource, model_field_name="sling"),
+    ] = None
+
+    @cached_property
+    def sling_resource(self) -> SlingResource:
+        return self.resource or SlingResource(connections=self.connections)
+
+    @cached_property
+    def _base_translator(self) -> DagsterSlingTranslator:
+        return DagsterSlingTranslator()
+
+    @public
+    def get_asset_spec(self, stream_definition: Mapping[str, Any]) -> AssetSpec:
+        """Generates an AssetSpec for a given Sling stream definition.
+
+        This method can be overridden in a subclass to customize how Sling stream definitions
+        are converted to Dagster asset specs. By default, it delegates to the configured
+        DagsterSlingTranslator.
+
+        Args:
+            stream_definition: A dictionary representing a single stream from the Sling
+                replication config, containing source and target information
+
+        Returns:
+            An AssetSpec that represents the Sling stream as a Dagster asset
+
+        Example:
+            Override this method to add custom metadata based on stream properties:
+
+            .. code-block:: python
+
+                from dagster_sling import SlingReplicationCollectionComponent
+                from dagster import AssetSpec
+
+                class CustomSlingComponent(SlingReplicationCollectionComponent):
+                    def get_asset_spec(self, stream_definition):
+                        base_spec = super().get_asset_spec(stream_definition)
+                        return base_spec.replace_attributes(
+                            metadata={
+                                **base_spec.metadata,
+                                "source": stream_definition.get("source"),
+                                "target": stream_definition.get("target")
+                            }
+                        )
+        """
+        return self._base_translator.get_asset_spec(stream_definition)
+
+    def build_asset(
+        self, context: ComponentLoadContext, replication_spec_model: SlingReplicationSpecModel
+    ) -> AssetsDefinition:
+        op_spec = replication_spec_model.op or OpSpec()
+        translator = SlingComponentTranslator(self, replication_spec_model, context.path)
+
+        @sling_assets(
+            name=op_spec.name or Path(replication_spec_model.path).stem,
+            op_tags=op_spec.tags,
+            replication_config=context.path / replication_spec_model.path,
+            dagster_sling_translator=translator,
+            backfill_policy=op_spec.backfill_policy,
+        )
+        def _asset(context: AssetExecutionContext):
+            yield from self.execute(
+                context=context,
+                sling=self.sling_resource,
+                replication_spec_model=replication_spec_model,
+            )
+
+        return _asset
+
+    @public
+    def execute(
+        self,
+        context: AssetExecutionContext,
+        sling: SlingResource,
+        replication_spec_model: SlingReplicationSpecModel,
+    ) -> Iterator[Union[AssetMaterialization, MaterializeResult]]:
+        """Executes a Sling replication for the selected streams.
+
+        This method can be overridden in a subclass to customize the replication execution
+        behavior, such as adding custom logging, modifying metadata collection, or handling
+        results differently.
+
+        Args:
+            context: The asset execution context provided by Dagster
+            sling: The SlingResource used to execute the replication
+            replication_spec_model: The model containing replication configuration and metadata options
+
+        Yields:
+            AssetMaterialization or MaterializeResult events from the Sling replication
+
+        Example:
+            Override this method to add custom logging during replication:
+
+            .. code-block:: python
+
+                from dagster_sling import SlingReplicationCollectionComponent
+                from dagster import AssetExecutionContext
+
+                class CustomSlingComponent(SlingReplicationCollectionComponent):
+                    def execute(self, context, sling, replication_spec_model):
+                        context.log.info("Starting Sling replication")
+                        yield from super().execute(context, sling, replication_spec_model)
+                        context.log.info("Sling replication completed")
+        """
+        iterator = sling.replicate(context=context)
+        if "column_metadata" in replication_spec_model.include_metadata:
+            iterator = iterator.fetch_column_metadata()
+        if "row_count" in replication_spec_model.include_metadata:
+            iterator = iterator.fetch_row_count()
+        yield from iterator
+
+    def build_defs(self, context: ComponentLoadContext) -> Definitions:
+        return Definitions(
+            assets=[self.build_asset(context, replication) for replication in self.replications],
+        )
+
+
+class SlingComponentTranslator(
+    create_component_translator_cls(SlingReplicationCollectionComponent, DagsterSlingTranslator),
+    ComponentTranslator[SlingReplicationCollectionComponent],
+):
+    def __init__(
+        self,
+        component: SlingReplicationCollectionComponent,
+        replication_spec: SlingReplicationSpecModel,
+        base_path: Path,
+    ):
+        self._component = component
+        self._replication_spec = replication_spec
+        self._base_path = base_path
+
+    def get_asset_spec(self, stream_definition: Mapping[str, Any]) -> AssetSpec:
+        spec = super().get_asset_spec(stream_definition)
+        if self._replication_spec.translation is not None:
+            spec = self._replication_spec.translation(spec, stream_definition)
+
+        # always add code references to the replication spec
+        code_reference = LocalFileCodeReference(
+            file_path=str(self._base_path / self._replication_spec.path)
+        )
+        return merge_code_references(spec, [code_reference])

dagster_sling/components/sling_replication_collection/scaffolder.py

@@ -0,0 +1,12 @@
+import yaml
+from dagster.components.component.component_scaffolder import Scaffolder
+from dagster.components.component_scaffolding import scaffold_component
+from dagster.components.scaffold.scaffold import ScaffoldRequest
+
+
+class SlingReplicationComponentScaffolder(Scaffolder):
+    def scaffold(self, request: ScaffoldRequest) -> None:
+        scaffold_component(request, {"replications": [{"path": "replication.yaml"}]})
+        replication_path = request.target_path / "replication.yaml"
+        with open(replication_path, "w") as f:
+            yaml.dump({"source": {}, "target": {}, "streams": {}}, f)