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

dagster_dbt/cloud_v2/sensor_builder.py

@@ -0,0 +1,254 @@
+from collections.abc import Iterator, Sequence
+from datetime import timedelta
+from typing import Optional, Union
+
+from dagster import (
+    AssetCheckEvaluation,
+    AssetKey,
+    AssetMaterialization,
+    AssetObservation,
+    DefaultSensorStatus,
+    SensorDefinition,
+    SensorEvaluationContext,
+    SensorResult,
+    _check as check,
+    sensor,
+)
+from dagster._core.definitions.repository_definition.repository_definition import (
+    RepositoryDefinition,
+)
+from dagster._grpc.client import DEFAULT_SENSOR_GRPC_TIMEOUT
+from dagster._record import record
+from dagster._serdes import deserialize_value, serialize_value
+from dagster._time import datetime_from_timestamp, get_current_datetime
+from dagster_shared.serdes import whitelist_for_serdes
+
+from dagster_dbt.cloud_v2.resources import DbtCloudWorkspace
+from dagster_dbt.cloud_v2.run_handler import (
+    COMPLETED_AT_TIMESTAMP_METADATA_KEY,
+    DbtCloudJobRunResults,
+)
+from dagster_dbt.cloud_v2.types import DbtCloudRun
+from dagster_dbt.dagster_dbt_translator import DagsterDbtTranslator
+from dagster_dbt.utils import clean_name
+
+MAIN_LOOP_TIMEOUT_SECONDS = DEFAULT_SENSOR_GRPC_TIMEOUT - 20
+DEFAULT_DBT_CLOUD_SENSOR_INTERVAL_SECONDS = 30
+START_LOOKBACK_SECONDS = 60  # Lookback one minute in time for the initial setting of the cursor.
+
+
+@record
+class BatchResult:
+    idx: int
+    asset_events: Sequence[AssetMaterialization]
+    all_asset_keys_materialized: set[AssetKey]
+
+
+@whitelist_for_serdes
+@record
+class DbtCloudPollingSensorCursor:
+    """A cursor that stores the last effective timestamp and offset."""
+
+    finished_at_lower_bound: Optional[float] = None
+    finished_at_upper_bound: Optional[float] = None
+    offset: Optional[int] = None
+
+
+def materializations_from_batch_iter(
+    context: SensorEvaluationContext,
+    finished_at_lower_bound: float,
+    finished_at_upper_bound: float,
+    offset: int,
+    workspace: DbtCloudWorkspace,
+    dagster_dbt_translator: DagsterDbtTranslator,
+) -> Iterator[Optional[BatchResult]]:
+    client = workspace.get_client()
+    workspace_data = workspace.get_or_fetch_workspace_data()
+
+    total_processed_runs = 0
+    while True:
+        latest_offset = total_processed_runs + offset
+        runs, total_runs = client.get_runs_batch(
+            project_id=workspace.project_id,
+            environment_id=workspace.environment_id,
+            finished_at_lower_bound=datetime_from_timestamp(finished_at_lower_bound),
+            finished_at_upper_bound=datetime_from_timestamp(finished_at_upper_bound),
+            offset=latest_offset,
+        )
+        if len(runs) == 0:
+            yield None
+            context.log.info("Received no runs. Breaking.")
+            break
+        context.log.info(
+            f"Processing {len(runs)}/{total_runs} runs for dbt Cloud workspace "
+            f"for project {workspace.project_name} and environment {workspace.environment_name}..."
+        )
+        for i, run_details in enumerate(runs):
+            run = DbtCloudRun.from_run_details(run_details=run_details)
+
+            if run.job_definition_id == workspace_data.adhoc_job_id:
+                context.log.info(f"Run {run.id} was triggered by Dagster. Continuing.")
+                continue
+
+            run_artifacts = client.list_run_artifacts(run_id=run.id)
+            if "run_results.json" not in run_artifacts:
+                context.log.info(
+                    f"Run {run.id} does not have a run_results.json artifact. Continuing."
+                )
+                continue
+
+            run_results = DbtCloudJobRunResults.from_run_results_json(
+                run_results_json=client.get_run_results_json(run_id=run.id)
+            )
+            events = run_results.to_default_asset_events(
+                client=workspace.get_client(),
+                manifest=workspace_data.manifest,
+                dagster_dbt_translator=dagster_dbt_translator,
+            )
+            # Currently, only materializations are tracked
+            mats = [event for event in events if isinstance(event, AssetMaterialization)]
+            context.log.info(f"Found {len(mats)} materializations for {run.id}")
+
+            all_asset_keys_materialized = {mat.asset_key for mat in mats}
+            yield (
+                BatchResult(
+                    idx=i + latest_offset,
+                    asset_events=mats,
+                    all_asset_keys_materialized=all_asset_keys_materialized,
+                )
+                if mats
+                else None
+            )
+        total_processed_runs += len(runs)
+        context.log.info(
+            f"Processed {total_processed_runs}/{total_runs} runs for dbt Cloud workspace "
+            f"for project {workspace.project_name} and environment {workspace.environment_name}..."
+        )
+        if total_processed_runs == total_runs:
+            yield None
+            context.log.info("Processed all runs. Breaking.")
+            break
+
+
+def sorted_asset_events(
+    asset_events: Sequence[Union[AssetMaterialization, AssetObservation, AssetCheckEvaluation]],
+    repository_def: RepositoryDefinition,
+) -> list[Union[AssetMaterialization, AssetObservation, AssetCheckEvaluation]]:
+    """Sort asset events by end date and toposort order."""
+    topo_aks = repository_def.asset_graph.toposorted_asset_keys
+    materializations_and_timestamps = [
+        (mat.metadata[COMPLETED_AT_TIMESTAMP_METADATA_KEY].value, mat) for mat in asset_events
+    ]
+    return [
+        sorted_event[1]
+        for sorted_event in sorted(
+            materializations_and_timestamps, key=lambda x: (x[0], topo_aks.index(x[1].asset_key))
+        )
+    ]
+
+
+def build_dbt_cloud_polling_sensor(
+    *,
+    workspace: DbtCloudWorkspace,
+    dagster_dbt_translator: Optional[DagsterDbtTranslator] = None,
+    minimum_interval_seconds: int = DEFAULT_DBT_CLOUD_SENSOR_INTERVAL_SECONDS,
+    default_sensor_status: Optional[DefaultSensorStatus] = None,
+) -> SensorDefinition:
+    """The constructed sensor polls the dbt Cloud Workspace for activity, and inserts asset events into Dagster's event log.
+
+    Args:
+        workspace (DbtCloudWorkspace): The dbt Cloud workspace to poll for runs.
+        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`.
+        minimum_interval_seconds (int, optional): The minimum interval in seconds between sensor runs. Defaults to 30.
+        default_sensor_status (Optional[DefaultSensorStatus], optional): The default status of the sensor.
+
+    Returns:
+        Definitions: A `SensorDefinitions` object.
+    """
+    dagster_dbt_translator = dagster_dbt_translator or DagsterDbtTranslator()
+
+    @sensor(
+        name=clean_name(
+            f"{workspace.account_name}_{workspace.project_name}_{workspace.environment_name}__run_status_sensor"
+        ),
+        description=(
+            f"dbt Cloud polling sensor for dbt Cloud workspace for account {workspace.account_name}, "
+            f"project {workspace.project_name} and environment {workspace.environment_name}"
+        ),
+        minimum_interval_seconds=minimum_interval_seconds,
+        default_status=default_sensor_status or DefaultSensorStatus.RUNNING,
+    )
+    def dbt_cloud_run_sensor(context: SensorEvaluationContext) -> SensorResult:
+        """Sensor to report materialization events for each asset as new runs come in."""
+        context.log.info(
+            f"************"
+            f"Running sensor for dbt Cloud workspace for account {workspace.account_name}, "
+            f"project {workspace.project_name} and environment {workspace.environment_name}"
+            f"***********"
+        )
+        try:
+            cursor = (
+                deserialize_value(context.cursor, DbtCloudPollingSensorCursor)
+                if context.cursor
+                else DbtCloudPollingSensorCursor()
+            )
+        except Exception as e:
+            context.log.info(f"Failed to interpret cursor. Starting from scratch. Error: {e}")
+            cursor = DbtCloudPollingSensorCursor()
+        current_date = get_current_datetime()
+        current_offset = cursor.offset or 0
+        finished_at_lower_bound = (
+            cursor.finished_at_lower_bound
+            or (current_date - timedelta(seconds=START_LOOKBACK_SECONDS)).timestamp()
+        )
+        finished_at_upper_bound = cursor.finished_at_upper_bound or current_date.timestamp()
+        sensor_iter = materializations_from_batch_iter(
+            context=context,
+            finished_at_lower_bound=finished_at_lower_bound,
+            finished_at_upper_bound=finished_at_upper_bound,
+            offset=current_offset,
+            workspace=workspace,
+            dagster_dbt_translator=dagster_dbt_translator,
+        )
+
+        all_asset_events: list[AssetMaterialization] = []
+        latest_offset = current_offset
+        repository_def = check.not_none(context.repository_def)
+        batch_result = None
+        while get_current_datetime() - current_date < timedelta(seconds=MAIN_LOOP_TIMEOUT_SECONDS):
+            batch_result = next(sensor_iter, None)
+            if batch_result is None:
+                context.log.info("Received no batch result. Breaking.")
+                break
+            all_asset_events.extend(batch_result.asset_events)
+            latest_offset = batch_result.idx
+
+        if batch_result is not None:
+            new_cursor = DbtCloudPollingSensorCursor(
+                finished_at_lower_bound=finished_at_lower_bound,
+                finished_at_upper_bound=finished_at_upper_bound,
+                offset=latest_offset + 1,
+            )
+        else:
+            # We have completed iteration for this range
+            new_cursor = DbtCloudPollingSensorCursor(
+                finished_at_lower_bound=finished_at_upper_bound,
+                finished_at_upper_bound=None,
+                offset=0,
+            )
+
+        context.update_cursor(serialize_value(new_cursor))
+
+        context.log.info(
+            f"************"
+            f"Exiting sensor for dbt Cloud workspace for account {workspace.account_name}, "
+            f"project {workspace.project_name} and environment {workspace.environment_name}"
+            f"***********"
+        )
+        return SensorResult(
+            asset_events=sorted_asset_events(all_asset_events, repository_def),
+        )
+
+    return dbt_cloud_run_sensor

dagster_dbt/cloud_v2/types.py

@@ -0,0 +1,143 @@
+from collections.abc import Mapping, Sequence
+from enum import Enum
+from typing import Any, Optional
+
+from dagster import Failure, MetadataValue
+from dagster._record import record
+from dagster._serdes import whitelist_for_serdes
+from dagster_shared.record import as_dict
+
+
+@record
+class DbtCloudAccount:
+    """Represents a dbt Cloud Account, based on data as returned from the API."""
+
+    id: int
+    name: Optional[str]
+
+    @classmethod
+    def from_account_details(cls, account_details: Mapping[str, Any]) -> "DbtCloudAccount":
+        return cls(
+            id=account_details["id"],
+            name=account_details.get("name"),
+        )
+
+
+@record
+class DbtCloudProject:
+    """Represents a dbt Cloud Project, based on data as returned from the API."""
+
+    id: int
+    name: Optional[str]
+
+    @classmethod
+    def from_project_details(cls, project_details: Mapping[str, Any]) -> "DbtCloudProject":
+        return cls(
+            id=project_details["id"],
+            name=project_details.get("name"),
+        )
+
+
+@record
+class DbtCloudEnvironment:
+    """Represents a dbt Cloud Environment, based on data as returned from the API."""
+
+    id: int
+    name: Optional[str]
+
+    @classmethod
+    def from_environment_details(
+        cls, environment_details: Mapping[str, Any]
+    ) -> "DbtCloudEnvironment":
+        return cls(
+            id=environment_details["id"],
+            name=environment_details.get("name"),
+        )
+
+
+@record
+class DbtCloudJob:
+    """Represents a dbt Cloud job, based on data as returned from the API."""
+
+    id: int
+    account_id: Optional[int]
+    project_id: Optional[int]
+    environment_id: Optional[int]
+    name: Optional[str]
+
+    @classmethod
+    def from_job_details(cls, job_details: Mapping[str, Any]) -> "DbtCloudJob":
+        return cls(
+            id=job_details["id"],
+            account_id=job_details.get("account_id"),
+            project_id=job_details.get("project_id"),
+            environment_id=job_details.get("environment_id"),
+            name=job_details.get("name"),
+        )
+
+
+class DbtCloudJobRunStatusType(int, Enum):
+    """Enum representing each status type for a run in dbt Cloud's ontology."""
+
+    QUEUED = 1
+    STARTING = 2
+    RUNNING = 3
+    SUCCESS = 10
+    ERROR = 20
+    CANCELLED = 30
+
+
+@record
+class DbtCloudRun:
+    """Represents a dbt Cloud run, based on data as returned from the API."""
+
+    id: int
+    trigger_id: Optional[int]
+    account_id: Optional[int]
+    project_id: Optional[int]
+    environment_id: Optional[int]
+    job_definition_id: Optional[int]
+    status: Optional[DbtCloudJobRunStatusType]
+    url: Optional[str]
+
+    @classmethod
+    def from_run_details(cls, run_details: Mapping[str, Any]) -> "DbtCloudRun":
+        return cls(
+            id=run_details["id"],
+            trigger_id=run_details.get("trigger_id"),
+            account_id=run_details.get("account_id"),
+            project_id=run_details.get("project_id"),
+            environment_id=run_details.get("environment_id"),
+            job_definition_id=run_details.get("job_definition_id"),
+            status=DbtCloudJobRunStatusType(run_details.get("status"))
+            if run_details.get("status")
+            else None,
+            url=run_details.get("href"),
+        )
+
+    def raise_for_status(self) -> None:
+        if self.status in {
+            DbtCloudJobRunStatusType.ERROR,
+            DbtCloudJobRunStatusType.CANCELLED,
+        }:
+            raise Failure(
+                f"dbt Cloud run '{self.id}' failed!",
+                metadata={
+                    "run_details": MetadataValue.json(as_dict(self)),
+                },
+            )
+
+
+@whitelist_for_serdes
+@record
+class DbtCloudWorkspaceData:
+    """Represents the data of a dbt Cloud workspace, given a project and environment."""
+
+    project_id: int
+    environment_id: int
+    # The ID of the ad hoc dbt Cloud job created by Dagster.
+    # This job is used to parse the dbt Cloud project.
+    # This job is also used to kick off cli invocation if no job ID is specified by users.
+    adhoc_job_id: int
+    manifest: Mapping[str, Any]
+    jobs: Sequence[Mapping[str, Any]]

dagster_dbt/compat.py

@@ -0,0 +1,107 @@
+import logging
+from enum import Enum
+from typing import TYPE_CHECKING, Any, TypeAlias
+
+from packaging import version
+
+# it's unclear exactly which dbt import adds a handler to the root logger, but something certainly does!
+# on this line, we keep track of the set of handlers that are on the root logger BEFORE any dbt imports
+# happen. at the end of this file, we set the root logger's handlers to the original set to ensure that
+# after this file is loaded, the root logger's handlers will be unchanged.
+existing_root_logger_handlers = [*logging.getLogger().handlers]
+
+
+try:
+    from dbt.version import __version__ as dbt_version
+
+    DBT_PYTHON_VERSION = version.parse(dbt_version)
+except ImportError:
+    DBT_PYTHON_VERSION = None
+
+# Conditionally define types for various types we use from the dbt-core package
+if TYPE_CHECKING:
+    from dbt.adapters.base.impl import (
+        BaseAdapter as _BaseAdapter,
+        BaseColumn as _BaseColumn,
+        BaseRelation as _BaseRelation,
+    )
+    from dbt.contracts.results import (
+        NodeStatus as _NodeStatus,
+        TestStatus as _TestStatus,
+    )
+    from dbt.node_types import NodeType as _NodeType
+
+    BaseAdapter: TypeAlias = _BaseAdapter
+    BaseColumn: TypeAlias = _BaseColumn
+    BaseRelation: TypeAlias = _BaseRelation
+    NodeStatus: TypeAlias = _NodeStatus
+    NodeType: TypeAlias = _NodeType
+    TestStatus: TypeAlias = _TestStatus
+    REFABLE_NODE_TYPES: list[str] = []
+else:
+    if DBT_PYTHON_VERSION is not None:
+        from dbt.adapters.base.impl import (
+            BaseAdapter as BaseAdapter,
+            BaseColumn as BaseColumn,
+            BaseRelation as BaseRelation,
+        )
+        from dbt.contracts.results import NodeStatus, TestStatus
+        from dbt.node_types import NodeType as NodeType
+
+        if DBT_PYTHON_VERSION < version.parse("1.8.0"):
+            from dbt.node_types import NodeType
+
+            REFABLE_NODE_TYPES = NodeType.refable()
+        else:
+            from dbt.node_types import REFABLE_NODE_TYPES as REFABLE_NODE_TYPES
+    else:
+        # here, we define implementations for types that will not be available if dbt-core is not
+        # installed
+        BaseAdapter = Any
+        BaseColumn = Any
+        BaseRelation = Any
+        REFABLE_NODE_TYPES = ["model", "seed", "snapshot"]
+
+        class StrEnum(str, Enum):
+            def _generate_next_value_(name, *_):
+                return name
+
+        class NodeType(StrEnum):
+            Model = "model"
+            Analysis = "analysis"
+            Test = "test"
+            Snapshot = "snapshot"
+            Operation = "operation"
+            Seed = "seed"
+            RPCCall = "rpc"
+            SqlOperation = "sql_operation"
+            Documentation = "doc"
+            Source = "source"
+            Macro = "macro"
+            Exposure = "exposure"
+            Metric = "metric"
+            Group = "group"
+            SavedQuery = "saved_query"
+            SemanticModel = "semantic_model"
+            Unit = "unit_test"
+            Fixture = "fixture"
+
+        class NodeStatus(StrEnum):
+            Success = "success"
+            Error = "error"
+            Fail = "fail"
+            Warn = "warn"
+            Skipped = "skipped"
+            PartialSuccess = "partial success"
+            Pass = "pass"
+            RuntimeErr = "runtime error"
+
+        class TestStatus(StrEnum):
+            Pass = NodeStatus.Pass
+            Error = NodeStatus.Error
+            Fail = NodeStatus.Fail
+            Warn = NodeStatus.Warn
+            Skipped = NodeStatus.Skipped
+
+
+logging.getLogger().handlers = existing_root_logger_handlers