dagster-dbt 0.23.3__py3-none-any.whl → 0.28.4__py3-none-any.whl

dagster_dbt/dagster_dbt_translator.py

@@ -1,33 +1,54 @@
+import os
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
-from typing import Any, Mapping, Optional, Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
 
 from dagster import (
+    AssetCheckSpec,
+    AssetDep,
     AssetKey,
+    AssetSpec,
     AutoMaterializePolicy,
-    FreshnessPolicy,
+    AutomationCondition,
+    DagsterInvalidDefinitionError,
     PartitionMapping,
-    _check as check,
 )
-from dagster._annotations import public
-from dagster._core.definitions.asset_key import (
-    CoercibleToAssetKeyPrefix,
-    check_opt_coercible_to_asset_key_prefix_param,
+from dagster._annotations import beta, public
+from dagster._core.definitions.metadata.source_code import (
+    CodeReferencesMetadataSet,
+    CodeReferencesMetadataValue,
+    LocalFileCodeReference,
 )
-from dagster._core.definitions.utils import is_valid_definition_tag_key
-
-from .asset_utils import (
+from dagster._core.definitions.partitions.definition import PartitionsDefinition
+from dagster._utils.tags import is_valid_tag_key
+from dagster.components.resolved.base import Resolvable
+from dagster_shared import check
+
+from dagster_dbt.asset_utils import (
+    DAGSTER_DBT_MANIFEST_METADATA_KEY,
+    DAGSTER_DBT_PROJECT_METADATA_KEY,
+    DAGSTER_DBT_TRANSLATOR_METADATA_KEY,
+    DAGSTER_DBT_UNIQUE_ID_METADATA_KEY,
+    default_asset_check_fn,
     default_asset_key_fn,
     default_auto_materialize_policy_fn,
+    default_code_version_fn,
     default_description_fn,
-    default_freshness_policy_fn,
     default_group_from_dbt_resource_props,
     default_metadata_from_dbt_resource_props,
     default_owners_from_dbt_resource_props,
+    get_node,
+    get_upstream_unique_ids,
+    has_self_dependency,
 )
 
+if TYPE_CHECKING:
+    from dagster_dbt.dbt_project import DbtProject
+
 
 @dataclass(frozen=True)
-class DagsterDbtTranslatorSettings:
+class DagsterDbtTranslatorSettings(Resolvable):
     """Settings to enable Dagster features for your dbt project.
 
     Args:
@@ -35,10 +56,19 @@ class DagsterDbtTranslatorSettings:
             Defaults to True.
         enable_duplicate_source_asset_keys (bool): Whether to allow dbt sources with duplicate
             Dagster asset keys. Defaults to False.
+        enable_code_references (bool): Whether to enable Dagster code references for dbt resources.
+            Defaults to False.
+        enable_dbt_selection_by_name (bool): Whether to enable selecting dbt resources by name,
+            rather than fully qualified name. Defaults to False.
+        enable_source_tests_as_checks (bool): Whether to load dbt source tests as Dagster asset checks.
+            Defaults to False. If False, asset observations will be emitted for source tests.
     """
 
     enable_asset_checks: bool = True
     enable_duplicate_source_asset_keys: bool = False
+    enable_code_references: bool = False
+    enable_dbt_selection_by_name: bool = False
+    enable_source_tests_as_checks: bool = False
 
 
 class DagsterDbtTranslator:
@@ -64,6 +94,145 @@ class DagsterDbtTranslator:
 
         return self._settings
 
+    def get_resource_props(self, manifest: Mapping[str, Any], unique_id: str) -> Mapping[str, Any]:
+        """Given a parsed manifest and a dbt unique_id, returns the dictionary of properties
+        for the corresponding dbt resource (e.g. model, seed, snapshot, source) as defined
+        in your dbt project. This can be used as a convenience method when overriding the
+        `get_asset_spec` method.
+
+        Args:
+            manifest (Mapping[str, Any]): The parsed manifest of the dbt project.
+            unique_id (str): The unique_id of the dbt resource.
+
+        Returns:
+            Mapping[str, Any]: The dictionary of properties for the corresponding dbt resource.
+
+        Examples:
+            .. code-block:: python
+
+                class CustomDagsterDbtTranslator(DagsterDbtTranslator):
+
+                    def get_asset_spec(self, manifest: Mapping[str, Any], unique_id: str, project: Optional[DbtProject]) -> dg.AssetSpec:
+                        base_spec = super().get_asset_spec(manifest, unique_id, project)
+                        resource_props = self.get_resource_props(manifest, unique_id)
+                        if resource_props["meta"].get("use_custom_group"):
+                            return base_spec.replace_attributes(group_name="custom_group")
+                        else:
+                            return base_spec
+        """
+        return get_node(manifest, unique_id)
+
+    def get_asset_spec(
+        self,
+        manifest: Mapping[str, Any],
+        unique_id: str,
+        project: Optional["DbtProject"],
+    ) -> AssetSpec:
+        """Returns an AssetSpec representing a specific dbt resource."""
+        # memoize resolution for a given manifest & unique_id
+        # since we recursively call get_asset_spec for dependencies
+        memo_id = (id(manifest), unique_id, id(project))
+
+        # Don't initialize this in the constructor in case a subclass does not call __init__
+        if not hasattr(self, "_resolved_specs"):
+            self._resolved_specs = {}
+
+        if memo_id in self._resolved_specs:
+            return self._resolved_specs[memo_id]
+
+        group_props = {group["name"]: group for group in manifest.get("groups", {}).values()}
+        resource_props = self.get_resource_props(manifest, unique_id)
+
+        # calculate the dependencies for the asset
+        upstream_ids = get_upstream_unique_ids(manifest, resource_props)
+        deps = [
+            AssetDep(
+                asset=self.get_asset_spec(manifest, upstream_id, project).key,
+                partition_mapping=self.get_partition_mapping(
+                    resource_props, self.get_resource_props(manifest, upstream_id)
+                ),
+            )
+            for upstream_id in upstream_ids
+        ]
+        self_partition_mapping = self.get_partition_mapping(resource_props, resource_props)
+        if self_partition_mapping and has_self_dependency(resource_props):
+            deps.append(
+                AssetDep(
+                    asset=self.get_asset_key(resource_props),
+                    partition_mapping=self_partition_mapping,
+                )
+            )
+
+        resource_group_props = group_props.get(resource_props.get("group") or "")
+        if resource_group_props:
+            owners_resource_props = {
+                **resource_props,
+                # this overrides the group key in resource_props, which is bad as
+                # this key is not always empty and this dictionary generally differs
+                # in structure from other inputs, but this is necessary for backcompat
+                **({"group": resource_group_props} if resource_group_props else {}),
+            }
+        else:
+            owners_resource_props = resource_props
+
+        spec = AssetSpec(
+            key=self.get_asset_key(resource_props),
+            deps=deps,
+            description=self.get_description(resource_props),
+            metadata=self.get_metadata(resource_props),
+            skippable=True,
+            group_name=self.get_group_name(resource_props),
+            code_version=self.get_code_version(resource_props),
+            automation_condition=self.get_automation_condition(resource_props),
+            owners=self.get_owners(owners_resource_props),
+            tags=self.get_tags(resource_props),
+            kinds={"dbt", manifest.get("metadata", {}).get("adapter_type", "dbt")},
+            partitions_def=self.get_partitions_def(resource_props),
+        )
+
+        # add integration-specific metadata to the spec
+        spec = spec.merge_attributes(
+            metadata={
+                DAGSTER_DBT_MANIFEST_METADATA_KEY: DbtManifestWrapper(manifest=manifest),
+                DAGSTER_DBT_TRANSLATOR_METADATA_KEY: self,
+                DAGSTER_DBT_UNIQUE_ID_METADATA_KEY: resource_props["unique_id"],
+                **({DAGSTER_DBT_PROJECT_METADATA_KEY: project} if project else {}),
+            }
+        )
+        if self.settings.enable_code_references:
+            if not project:
+                raise DagsterInvalidDefinitionError(
+                    "enable_code_references requires a DbtProject to be supplied"
+                    " to the @dbt_assets decorator."
+                )
+
+            spec = spec.replace_attributes(
+                metadata=_attach_sql_model_code_reference(
+                    existing_metadata=spec.metadata,
+                    dbt_resource_props=resource_props,
+                    project=project,
+                )
+            )
+
+        self._resolved_specs[memo_id] = spec
+
+        return self._resolved_specs[memo_id]
+
+    def get_asset_check_spec(
+        self,
+        asset_spec: AssetSpec,
+        manifest: Mapping[str, Any],
+        unique_id: str,
+        project: Optional["DbtProject"],
+    ) -> Optional[AssetCheckSpec]:
+        return default_asset_check_fn(
+            manifest=manifest,
+            dagster_dbt_translator=self,
+            asset_key=asset_spec.key,
+            test_unique_id=unique_id,
+            project=project,
+        )
+
     @public
     def get_asset_key(self, dbt_resource_props: Mapping[str, Any]) -> AssetKey:
         """A function that takes a dictionary representing properties of a dbt resource, and
@@ -119,6 +288,7 @@ class DagsterDbtTranslator:
         return default_asset_key_fn(dbt_resource_props)
 
     @public
+    @beta(emit_runtime_warning=False)
     def get_partition_mapping(
         self,
         dbt_resource_props: Mapping[str, Any],
@@ -223,8 +393,8 @@ class DagsterDbtTranslator:
         https://docs.getdbt.com/reference/artifacts/manifest-json#resource-details
 
         dbt tags are strings, but Dagster tags are key-value pairs. To bridge this divide, the dbt
-        tag string is used as the Dagster tag key, and the Dagster tag value is set to special
-        sentinel value `"__dagster_no_value"`.
+        tag string is used as the Dagster tag key, and the Dagster tag value is set to the empty
+        string, "".
 
         Any dbt tags that don't match Dagster's supported tag key format (e.g. they contain
         unsupported characters) will be ignored.
@@ -250,7 +420,7 @@ class DagsterDbtTranslator:
                         return {"custom": "tag"}
         """
         tags = dbt_resource_props.get("tags", [])
-        return {tag: "" for tag in tags if is_valid_definition_tag_key(tag)}
+        return {tag: "" for tag in tags if is_valid_tag_key(tag)}
 
     @public
     def get_group_name(self, dbt_resource_props: Mapping[str, Any]) -> Optional[str]:
@@ -284,6 +454,38 @@ class DagsterDbtTranslator:
         """
         return default_group_from_dbt_resource_props(dbt_resource_props)
 
+    @public
+    def get_code_version(self, dbt_resource_props: Mapping[str, Any]) -> Optional[str]:
+        """A function that takes a dictionary representing properties of a dbt resource, and
+        returns the Dagster code version for that resource.
+
+        Note that a dbt resource is unrelated to Dagster's resource concept, and simply represents
+        a model, seed, snapshot or source in a given dbt project. You can learn more about dbt
+        resources and the properties available in this dictionary here:
+        https://docs.getdbt.com/reference/artifacts/manifest-json#resource-details
+
+        This method can be overridden to provide a custom code version for a dbt resource.
+
+        Args:
+            dbt_resource_props (Mapping[str, Any]): A dictionary representing the dbt resource.
+
+        Returns:
+            Optional[str]: A Dagster code version.
+
+        Examples:
+            .. code-block:: python
+
+                from typing import Any, Mapping
+
+                from dagster_dbt import DagsterDbtTranslator
+
+
+                class CustomDagsterDbtTranslator(DagsterDbtTranslator):
+                    def get_code_version(self, dbt_resource_props: Mapping[str, Any]) -> Optional[str]:
+                        return dbt_resource_props["checksum"]["checksum"]
+        """
+        return default_code_version_fn(dbt_resource_props)
+
     @public
     def get_owners(self, dbt_resource_props: Mapping[str, Any]) -> Optional[Sequence[str]]:
         """A function that takes a dictionary representing properties of a dbt resource, and
@@ -317,27 +519,28 @@ class DagsterDbtTranslator:
         return default_owners_from_dbt_resource_props(dbt_resource_props)
 
     @public
-    def get_freshness_policy(
+    @beta(emit_runtime_warning=False)
+    def get_auto_materialize_policy(
         self, dbt_resource_props: Mapping[str, Any]
-    ) -> Optional[FreshnessPolicy]:
+    ) -> Optional[AutoMaterializePolicy]:
         """A function that takes a dictionary representing properties of a dbt resource, and
-        returns the Dagster :py:class:`dagster.FreshnessPolicy` for that resource.
+        returns the Dagster :py:class:`dagster.AutoMaterializePolicy` for that resource.
 
         Note that a dbt resource is unrelated to Dagster's resource concept, and simply represents
         a model, seed, snapshot or source in a given dbt project. You can learn more about dbt
         resources and the properties available in this dictionary here:
         https://docs.getdbt.com/reference/artifacts/manifest-json#resource-details
 
-        This method can be overridden to provide a custom freshness policy for a dbt resource.
+        This method can be overridden to provide a custom auto-materialize policy for a dbt resource.
 
         Args:
             dbt_resource_props (Mapping[str, Any]): A dictionary representing the dbt resource.
 
         Returns:
-            Optional[FreshnessPolicy]: A Dagster freshness policy.
+            Optional[AutoMaterializePolicy]: A Dagster auto-materialize policy.
 
         Examples:
-            Set a custom freshness policy for all dbt resources:
+            Set a custom auto-materialize policy for all dbt resources:
 
             .. code-block:: python
 
@@ -347,10 +550,10 @@ class DagsterDbtTranslator:
 
 
                 class CustomDagsterDbtTranslator(DagsterDbtTranslator):
-                    def get_freshness_policy(self, dbt_resource_props: Mapping[str, Any]) -> Optional[FreshnessPolicy]:
-                        return FreshnessPolicy(maximum_lag_minutes=60)
+                    def get_auto_materialize_policy(self, dbt_resource_props: Mapping[str, Any]) -> Optional[AutoMaterializePolicy]:
+                        return AutoMaterializePolicy.eager()
 
-            Set a custom freshness policy for dbt resources with a specific tag:
+            Set a custom auto-materialize policy for dbt resources with a specific tag:
 
             .. code-block:: python
 
@@ -360,19 +563,21 @@ class DagsterDbtTranslator:
 
 
                 class CustomDagsterDbtTranslator(DagsterDbtTranslator):
-                    def get_freshness_policy(self, dbt_resource_props: Mapping[str, Any]) -> Optional[FreshnessPolicy]:
-                        freshness_policy = None
+                    def get_auto_materialize_policy(self, dbt_resource_props: Mapping[str, Any]) -> Optional[AutoMaterializePolicy]:
+                        auto_materialize_policy = None
                         if "my_custom_tag" in dbt_resource_props.get("tags", []):
-                            freshness_policy = FreshnessPolicy(maximum_lag_minutes=60)
+                            auto_materialize_policy = AutoMaterializePolicy.eager()
+
+                        return auto_materialize_policy
 
-                        return freshness_policy
         """
-        return default_freshness_policy_fn(dbt_resource_props)
+        return default_auto_materialize_policy_fn(dbt_resource_props)
 
     @public
-    def get_auto_materialize_policy(
+    @beta(emit_runtime_warning=False)
+    def get_automation_condition(
         self, dbt_resource_props: Mapping[str, Any]
-    ) -> Optional[AutoMaterializePolicy]:
+    ) -> Optional[AutomationCondition]:
         """A function that takes a dictionary representing properties of a dbt resource, and
         returns the Dagster :py:class:`dagster.AutoMaterializePolicy` for that resource.
 
@@ -381,7 +586,7 @@ class DagsterDbtTranslator:
         resources and the properties available in this dictionary here:
         https://docs.getdbt.com/reference/artifacts/manifest-json#resource-details
 
-        This method can be overridden to provide a custom auto-materialize policy for a dbt resource.
+        This method can be overridden to provide a custom AutomationCondition for a dbt resource.
 
         Args:
             dbt_resource_props (Mapping[str, Any]): A dictionary representing the dbt resource.
@@ -390,7 +595,7 @@ class DagsterDbtTranslator:
             Optional[AutoMaterializePolicy]: A Dagster auto-materialize policy.
 
         Examples:
-            Set a custom auto-materialize policy for all dbt resources:
+            Set a custom AutomationCondition for all dbt resources:
 
             .. code-block:: python
 
@@ -400,10 +605,10 @@ class DagsterDbtTranslator:
 
 
                 class CustomDagsterDbtTranslator(DagsterDbtTranslator):
-                    def get_auto_materialize_policy(self, dbt_resource_props: Mapping[str, Any]) -> Optional[AutoMaterializePolicy]:
-                        return AutoMaterializePolicy.eager()
+                    def get_automation_condition(self, dbt_resource_props: Mapping[str, Any]) -> Optional[AutomationCondition]:
+                        return AutomationCondition.eager()
 
-            Set a custom auto-materialize policy for dbt resources with a specific tag:
+            Set a custom AutomationCondition for dbt resources with a specific tag:
 
             .. code-block:: python
 
@@ -413,58 +618,58 @@ class DagsterDbtTranslator:
 
 
                 class CustomDagsterDbtTranslator(DagsterDbtTranslator):
-                    def get_auto_materialize_policy(self, dbt_resource_props: Mapping[str, Any]) -> Optional[AutoMaterializePolicy]:
-                        auto_materialize_policy = None
+                    def get_automation_condition(self, dbt_resource_props: Mapping[str, Any]) -> Optional[AutomationCondition]:
+                        automation_condition = None
                         if "my_custom_tag" in dbt_resource_props.get("tags", []):
-                            auto_materialize_policy = AutoMaterializePolicy.eager()
+                            automation_condition = AutomationCondition.eager()
 
-                        return auto_materialize_policy
+                        return automation_condition
 
         """
-        return default_auto_materialize_policy_fn(dbt_resource_props)
+        auto_materialize_policy = self.get_auto_materialize_policy(dbt_resource_props)
+        return (
+            auto_materialize_policy.to_automation_condition() if auto_materialize_policy else None
+        )
 
+    def get_partitions_def(
+        self, dbt_resource_props: Mapping[str, Any]
+    ) -> Optional[PartitionsDefinition]:
+        """[INTERNAL] A function that takes a dictionary representing properties of a dbt resource, and
+        returns the Dagster :py:class:`dagster.PartitionsDefinition` for that resource.
 
-class KeyPrefixDagsterDbtTranslator(DagsterDbtTranslator):
-    """A DagsterDbtTranslator that applies prefixes to the asset keys generated from dbt resources.
+        This method can be overridden to provide a custom PartitionsDefinition for a dbt resource.
 
-    Attributes:
-        asset_key_prefix (Optional[Union[str, Sequence[str]]]): A prefix to apply to all dbt models,
-            seeds, snapshots, etc. This will *not* apply to dbt sources.
-        source_asset_key_prefix (Optional[Union[str, Sequence[str]]]): A prefix to apply to all dbt
-            sources.
-    """
+        Args:
+            dbt_resource_props (Mapping[str, Any]): A dictionary representing the dbt resource.
 
-    def __init__(
-        self,
-        asset_key_prefix: Optional[CoercibleToAssetKeyPrefix] = None,
-        source_asset_key_prefix: Optional[CoercibleToAssetKeyPrefix] = None,
-        *args,
-        **kwargs,
-    ):
-        self._asset_key_prefix = (
-            check_opt_coercible_to_asset_key_prefix_param(asset_key_prefix, "asset_key_prefix")
-            or []
-        )
-        self._source_asset_key_prefix = (
-            check_opt_coercible_to_asset_key_prefix_param(
-                source_asset_key_prefix, "source_asset_key_prefix"
-            )
-            or []
-        )
+        Returns:
+            Optional[PartitionsDefinition]: A Dagster partitions definition.
 
-        super().__init__(*args, **kwargs)
+        Examples:
+            Set a custom AutomationCondition for dbt resources with a specific tag:
 
-    @public
-    def get_asset_key(self, dbt_resource_props: Mapping[str, Any]) -> AssetKey:
-        base_key = default_asset_key_fn(dbt_resource_props)
-        if dbt_resource_props["resource_type"] == "source":
-            return base_key.with_prefix(self._source_asset_key_prefix)
-        else:
-            return base_key.with_prefix(self._asset_key_prefix)
+            .. code-block:: python
+
+                from typing import Any, Mapping
+
+                from dagster import DailyPartitionsDefinition
+                from dagster_dbt import DagsterDbtTranslator
+
+
+                class CustomDagsterDbtTranslator(DagsterDbtTranslator):
+                    def get_partitions_def(self, dbt_resource_props: Mapping[str, Any]) -> Optional[PartitionsDefinition]:
+                        if "my_custom_tag" in dbt_resource_props.get("tags", []):
+                            return DailyPartitionsDefinition(start_date="2022-01-01")
+                        else:
+                            return None
+        """
+        return None
 
 
 @dataclass
 class DbtManifestWrapper:
+    """Wrapper around parsed DBT manifest json to provide convenient and efficient access."""
+
     manifest: Mapping[str, Any]
 
 
@@ -492,3 +697,41 @@ def validate_opt_translator(
             " DagsterDbtTranslator."
         ),
     )
+
+
+def _attach_sql_model_code_reference(
+    existing_metadata: Mapping[str, Any],
+    dbt_resource_props: Mapping[str, Any],
+    project: "DbtProject",
+) -> Mapping[str, Any]:
+    """Pulls the SQL model location for a dbt resource and attaches it as a code reference to the
+    existing metadata.
+    """
+    existing_references_meta = CodeReferencesMetadataSet.extract(existing_metadata)
+    references = (
+        existing_references_meta.code_references.code_references
+        if existing_references_meta.code_references
+        else []
+    )
+
+    if "original_file_path" not in dbt_resource_props:
+        raise DagsterInvalidDefinitionError(
+            "Cannot attach SQL model code reference because 'original_file_path' is not present"
+            " in the dbt resource properties."
+        )
+
+    # attempt to get root_path, which is removed from manifests in newer dbt versions
+    relative_path = Path(dbt_resource_props["original_file_path"])
+    abs_path = project.project_dir.joinpath(relative_path).resolve()
+
+    return {
+        **existing_metadata,
+        **CodeReferencesMetadataSet(
+            code_references=CodeReferencesMetadataValue(
+                code_references=[
+                    *references,
+                    LocalFileCodeReference(file_path=os.fspath(abs_path)),
+                ],
+            )
+        ),
+    }

dagster_dbt/dbt_core_version.py

@@ -0,0 +1 @@
+DBT_CORE_VERSION_UPPER_BOUND = "1.11"

dagster_dbt/dbt_manifest.py

@@ -1,16 +1,17 @@
-from functools import lru_cache
+from collections.abc import Mapping
+from functools import cache
 from pathlib import Path
-from typing import Any, Mapping, Union, cast
+from typing import Any, Union, cast
 
 import dagster._check as check
 import orjson
 
-from .errors import DagsterDbtManifestNotFoundError
+from dagster_dbt.errors import DagsterDbtManifestNotFoundError
 
 DbtManifestParam = Union[Mapping[str, Any], str, Path]
 
 
-@lru_cache(maxsize=None)
+@cache
 def read_manifest_path(manifest_path: Path) -> Mapping[str, Any]:
     """Reads a dbt manifest path and returns the parsed JSON as a dict.
 
@@ -23,7 +24,7 @@ def read_manifest_path(manifest_path: Path) -> Mapping[str, Any]:
     if not manifest_path.exists():
         raise DagsterDbtManifestNotFoundError(f"{manifest_path} does not exist.")
 
-    return cast(Mapping[str, Any], orjson.loads(manifest_path.read_bytes()))
+    return cast("Mapping[str, Any]", orjson.loads(manifest_path.read_bytes()))
 
 
 def validate_manifest(manifest: DbtManifestParam) -> Mapping[str, Any]: