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

dagster_dbt/cloud/types.py

@@ -1,20 +1,21 @@
+from collections.abc import Mapping
 from datetime import datetime
-from typing import Any, Mapping, Optional
+from typing import Any, Optional
 
 import dagster._check as check
-from dateutil.parser import isoparse
+from dagster._annotations import beta
+from dagster._vendored.dateutil.parser import isoparse
 
-from ..types import DbtOutput
 
-
-class DbtCloudOutput(DbtOutput):
+@beta
+class DbtCloudOutput:
     """The results of executing a dbt Cloud job, along with additional metadata produced from the
     job run.
 
     Note that users should not construct instances of this class directly. This class is intended
     to be constructed from the JSON output of dbt Cloud commands.
 
-    Attributes:
+    Args:
         run_details (Dict[str, Any]): The raw dictionary data representing the run details returned
             by the dbt Cloud API. For more info, see: https://docs.getdbt.com/dbt-cloud/api-v2#operation/getRunById
         result (Dict[str, Any]): Dictionary containing dbt-reported result information
@@ -32,7 +33,11 @@ class DbtCloudOutput(DbtOutput):
         result: Mapping[str, Any],
     ):
         self._run_details = check.mapping_param(run_details, "run_details", key_type=str)
-        super().__init__(result)
+        self._result = check.mapping_param(result, "result", key_type=str)
+
+    @property
+    def result(self) -> Mapping[str, Any]:
+        return self._result
 
     @property
     def run_details(self) -> Mapping[str, Any]:

dagster_dbt/cloud/utils.py

@@ -0,0 +1,186 @@
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from typing import Any, Optional, Union
+
+import dateutil
+from dagster import (
+    AssetKey,
+    AssetMaterialization,
+    AssetObservation,
+    MetadataValue,
+    Output,
+    _check as check,
+)
+from dagster._core.definitions.metadata import RawMetadataValue
+
+from dagster_dbt.cloud.types import DbtCloudOutput
+from dagster_dbt.utils import ASSET_RESOURCE_TYPES, default_node_info_to_asset_key
+
+
+def _resource_type(unique_id: str) -> str:
+    # returns the type of the node (e.g. model, test, snapshot)
+    return unique_id.split(".")[0]
+
+
+def _node_result_to_metadata(node_result: Mapping[str, Any]) -> Mapping[str, RawMetadataValue]:
+    return {
+        "Materialization Strategy": node_result["config"]["materialized"],
+        "Database": node_result["database"],
+        "Schema": node_result["schema"],
+        "Alias": node_result["alias"],
+        "Description": node_result["description"],
+    }
+
+
+def _timing_to_metadata(timings: Sequence[Mapping[str, Any]]) -> Mapping[str, RawMetadataValue]:
+    metadata: dict[str, RawMetadataValue] = {}
+    for timing in timings:
+        if timing["name"] == "execute":
+            desc = "Execution"
+        elif timing["name"] == "compile":
+            desc = "Compilation"
+        else:
+            continue
+
+        # dateutil does not properly expose its modules to static checkers
+        started_at = dateutil.parser.isoparse(timing["started_at"])
+        completed_at = dateutil.parser.isoparse(timing["completed_at"])
+        duration = completed_at - started_at
+        metadata.update(
+            {
+                f"{desc} Started At": started_at.isoformat(timespec="seconds"),
+                f"{desc} Completed At": started_at.isoformat(timespec="seconds"),
+                f"{desc} Duration": duration.total_seconds(),
+            }
+        )
+    return metadata
+
+
+def result_to_events(
+    result: Mapping[str, Any],
+    docs_url: Optional[str] = None,
+    node_info_to_asset_key: Optional[Callable[[Mapping[str, Any]], AssetKey]] = None,
+    manifest_json: Optional[Mapping[str, Any]] = None,
+    extra_metadata: Optional[Mapping[str, RawMetadataValue]] = None,
+    generate_asset_outputs: bool = False,
+) -> Iterator[Union[AssetMaterialization, AssetObservation, Output]]:
+    """This is a hacky solution that attempts to consolidate parsing many of the potential formats
+    that dbt can provide its results in. This is known to work for CLI Outputs for dbt versions 0.18+,
+    as well as RPC responses for a similar time period, but as the RPC response schema is not documented
+    nor enforced, this can become out of date easily.
+    """
+    node_info_to_asset_key = check.opt_callable_param(
+        node_info_to_asset_key, "node_info_to_asset_key", default=default_node_info_to_asset_key
+    )
+
+    # status comes from set of fields rather than "status"
+    if "fail" in result:
+        status = (
+            "fail"
+            if result.get("fail")
+            else "skip"
+            if result.get("skip")
+            else "error"
+            if result.get("error")
+            else "success"
+        )
+    else:
+        status = result["status"]
+
+    # all versions represent timing the same way
+    metadata = {"Status": status, "Execution Time (seconds)": result["execution_time"]}
+    metadata.update(_timing_to_metadata(result["timing"]))
+
+    # working with a response that contains the node block (RPC and CLI 0.18.x)
+    if "node" in result:
+        unique_id = result["node"]["unique_id"]
+        metadata.update(_node_result_to_metadata(result["node"]))
+    else:
+        unique_id = result["unique_id"]
+
+    if docs_url:
+        metadata["docs_url"] = MetadataValue.url(f"{docs_url}#!/model/{unique_id}")
+
+    if extra_metadata:
+        metadata.update(extra_metadata)
+
+    # if you have a manifest available, get the full node info, otherwise just populate unique_id
+    dbt_resource_props = (
+        manifest_json["nodes"][unique_id] if manifest_json else {"unique_id": unique_id}
+    )
+
+    node_resource_type = _resource_type(unique_id)
+
+    if node_resource_type in ASSET_RESOURCE_TYPES and status == "success":
+        if generate_asset_outputs:
+            yield Output(
+                value=None,
+                output_name=node_info_to_asset_key(dbt_resource_props).to_python_identifier(),
+                metadata=metadata,
+            )
+        else:
+            yield AssetMaterialization(
+                asset_key=node_info_to_asset_key(dbt_resource_props),
+                description=f"dbt node: {unique_id}",
+                metadata=metadata,
+            )
+    # can only associate tests with assets if we have manifest_json available
+    elif node_resource_type == "test" and manifest_json and status != "skipped":
+        upstream_unique_ids = manifest_json["nodes"][unique_id]["depends_on"]["nodes"]
+        # tests can apply to multiple asset keys
+        for upstream_id in upstream_unique_ids:
+            # the upstream id can reference a node or a source
+            dbt_resource_props = manifest_json["nodes"].get(upstream_id) or manifest_json[
+                "sources"
+            ].get(upstream_id)
+            if dbt_resource_props is None:
+                continue
+            upstream_asset_key = node_info_to_asset_key(dbt_resource_props)
+            yield AssetObservation(
+                asset_key=upstream_asset_key,
+                metadata={
+                    "Test ID": result["unique_id"],
+                    "Test Status": status,
+                    "Test Message": result.get("message") or "",
+                },
+            )
+
+
+def generate_events(
+    dbt_output: DbtCloudOutput,
+    node_info_to_asset_key: Optional[Callable[[Mapping[str, Any]], AssetKey]] = None,
+    manifest_json: Optional[Mapping[str, Any]] = None,
+) -> Iterator[Union[AssetMaterialization, AssetObservation]]:
+    """This function yields :py:class:`dagster.AssetMaterialization` events for each model updated by
+    a dbt command, and :py:class:`dagster.AssetObservation` events for each test run.
+
+    Information parsed from a :py:class:`~DbtOutput` object.
+    """
+    for result in dbt_output.result["results"]:
+        for event in result_to_events(
+            result,
+            docs_url=dbt_output.docs_url,
+            node_info_to_asset_key=node_info_to_asset_key,
+            manifest_json=manifest_json,
+        ):
+            yield check.inst(
+                event,
+                (AssetMaterialization, AssetObservation),
+            )
+
+
+def generate_materializations(
+    dbt_output: DbtCloudOutput,
+    asset_key_prefix: Optional[Sequence[str]] = None,
+) -> Iterator[AssetMaterialization]:
+    """This function yields :py:class:`dagster.AssetMaterialization` events for each model updated by
+    a dbt command.
+    """
+    asset_key_prefix = check.opt_sequence_param(asset_key_prefix, "asset_key_prefix", of_type=str)
+
+    for event in generate_events(
+        dbt_output,
+        node_info_to_asset_key=lambda info: AssetKey(
+            asset_key_prefix + info["unique_id"].split(".")
+        ),
+    ):
+        yield check.inst(event, AssetMaterialization)

dagster_dbt/cloud_v2/__init__.py

@@ -0,0 +1,10 @@
+from dagster_dbt.cloud_v2.asset_decorator import dbt_cloud_assets as dbt_cloud_assets
+from dagster_dbt.cloud_v2.resources import (
+    DbtCloudCredentials as DbtCloudCredentials,
+    DbtCloudWorkspace as DbtCloudWorkspace,
+    load_dbt_cloud_asset_specs as load_dbt_cloud_asset_specs,
+    load_dbt_cloud_check_specs as load_dbt_cloud_check_specs,
+)
+from dagster_dbt.cloud_v2.sensor_builder import (
+    build_dbt_cloud_polling_sensor as build_dbt_cloud_polling_sensor,
+)

dagster_dbt/cloud_v2/asset_decorator.py

@@ -0,0 +1,81 @@
+from collections.abc import Callable
+from typing import Any, Optional
+
+from dagster import AssetsDefinition, multi_asset
+from dagster._annotations import public
+from dagster._core.errors import DagsterInvariantViolationError
+
+from dagster_dbt.asset_utils import (
+    DAGSTER_DBT_EXCLUDE_METADATA_KEY,
+    DAGSTER_DBT_SELECT_METADATA_KEY,
+    DAGSTER_DBT_SELECTOR_METADATA_KEY,
+    DBT_DEFAULT_EXCLUDE,
+    DBT_DEFAULT_SELECT,
+    DBT_DEFAULT_SELECTOR,
+)
+from dagster_dbt.cloud_v2.resources import DbtCloudWorkspace
+from dagster_dbt.dagster_dbt_translator import DagsterDbtTranslator
+
+
+@public
+def dbt_cloud_assets(
+    *,
+    workspace: DbtCloudWorkspace,
+    select: str = DBT_DEFAULT_SELECT,
+    exclude: str = DBT_DEFAULT_EXCLUDE,
+    selector: str = DBT_DEFAULT_SELECTOR,
+    name: Optional[str] = None,
+    group_name: Optional[str] = None,
+    dagster_dbt_translator: Optional[DagsterDbtTranslator] = None,
+) -> Callable[[Callable[..., Any]], AssetsDefinition]:
+    """Create a definition for how to compute a set of dbt Cloud resources,
+    described by a manifest.json for a given dbt Cloud workspace.
+
+    Args:
+        workspace (DbtCloudWorkspace): The dbt Cloud workspace.
+        select (str): A dbt selection string for the models in a project that you want
+            to include. Defaults to ``fqn:*``.
+        exclude (str): A dbt selection string for the models in a project that you want
+            to exclude. Defaults to "".
+        selector (str): A dbt selector to select resources to materialize. Defaults to "".
+        name (Optional[str], optional): The name of the op.
+        group_name (Optional[str], optional): The name of the asset group.
+        dagster_dbt_translator (Optional[DagsterDbtTranslator], optional): The translator to use
+            to convert dbt Cloud content into :py:class:`dagster.AssetSpec`.
+            Defaults to :py:class:`DagsterDbtTranslator`.
+    """
+    dagster_dbt_translator = dagster_dbt_translator or DagsterDbtTranslator()
+
+    op_tags = {
+        DAGSTER_DBT_SELECT_METADATA_KEY: select,
+        DAGSTER_DBT_EXCLUDE_METADATA_KEY: exclude,
+        DAGSTER_DBT_SELECTOR_METADATA_KEY: selector,
+    }
+
+    specs = workspace.load_asset_specs(
+        dagster_dbt_translator=dagster_dbt_translator,
+        select=select,
+        exclude=exclude,
+        selector=selector,
+    )
+
+    if any([spec for spec in specs if spec.group_name]) and group_name:
+        raise DagsterInvariantViolationError(
+            f"Cannot set group_name parameter on dbt_cloud_assets for dbt Cloud workspace with account "
+            f"{workspace.account_name}, project {workspace.project_name} and environment {workspace.environment_name} -"
+            f" one or more of the dbt Cloud asset specs have a group_name defined."
+        )
+
+    return multi_asset(
+        name=name,
+        group_name=group_name,
+        can_subset=True,
+        specs=specs,
+        op_tags=op_tags,
+        check_specs=workspace.load_check_specs(
+            dagster_dbt_translator=dagster_dbt_translator,
+            select=select,
+            exclude=exclude,
+            selector=selector,
+        ),
+    )

dagster_dbt/cloud_v2/cli_invocation.py

@@ -0,0 +1,67 @@
+from collections.abc import Iterator, Mapping, Sequence
+from typing import Any, Optional, Union
+
+from dagster import (
+    AssetCheckEvaluation,
+    AssetCheckResult,
+    AssetExecutionContext,
+    AssetMaterialization,
+    Output,
+)
+from dagster._record import record
+
+from dagster_dbt.cloud_v2.client import DbtCloudWorkspaceClient
+from dagster_dbt.cloud_v2.run_handler import DbtCloudJobRunHandler, DbtCloudJobRunResults
+from dagster_dbt.dagster_dbt_translator import DagsterDbtTranslator
+
+
+@record
+class DbtCloudCliInvocation:
+    """Represents a dbt Cloud cli invocation."""
+
+    args: Sequence[str]
+    client: DbtCloudWorkspaceClient
+    manifest: Mapping[str, Any]
+    dagster_dbt_translator: DagsterDbtTranslator
+    run_handler: DbtCloudJobRunHandler
+    context: Optional[AssetExecutionContext]
+
+    @classmethod
+    def run(
+        cls,
+        job_id: int,
+        args: Sequence[str],
+        client: DbtCloudWorkspaceClient,
+        manifest: Mapping[str, Any],
+        dagster_dbt_translator: DagsterDbtTranslator,
+        context: Optional[AssetExecutionContext] = None,
+    ) -> "DbtCloudCliInvocation":
+        run_handler = DbtCloudJobRunHandler.run(
+            job_id=job_id,
+            args=args,
+            client=client,
+        )
+        return DbtCloudCliInvocation(
+            args=args,
+            client=client,
+            manifest=manifest,
+            dagster_dbt_translator=dagster_dbt_translator,
+            run_handler=run_handler,
+            context=context,
+        )
+
+    def wait(
+        self, timeout: Optional[float] = None
+    ) -> Iterator[Union[AssetCheckEvaluation, AssetCheckResult, AssetMaterialization, Output]]:
+        run = self.run_handler.wait(timeout=timeout)
+        if "run_results.json" in self.run_handler.list_run_artifacts():
+            run_results = DbtCloudJobRunResults.from_run_results_json(
+                run_results_json=self.run_handler.get_run_results()
+            )
+            yield from run_results.to_default_asset_events(
+                client=self.client,
+                manifest=self.manifest,
+                dagster_dbt_translator=self.dagster_dbt_translator,
+                context=self.context,
+            )
+        run.raise_for_status()