This is an automated email from the ASF dual-hosted git repository. mobuchowski pushed a commit to branch dbt-cloud-cause in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 4bac5ec5aab8e7ae31cdfecb05d71baa18abc059 Author: Maciej Obuchowski <[email protected]> AuthorDate: Mon Jul 20 15:01:45 2026 +0200 dbt cloud: add passing OL ParentRunFacet as dbt Cloud cause Signed-off-by: Maciej Obuchowski <[email protected]> --- providers/dbt/cloud/docs/operators.rst | 14 ++++ .../airflow/providers/dbt/cloud/operators/dbt.py | 22 ++++++- .../providers/dbt/cloud/utils/openlineage.py | 58 +++++++++++++++++ .../tests/unit/dbt/cloud/operators/test_dbt.py | 74 ++++++++++++++++++++++ .../tests/unit/dbt/cloud/utils/test_openlineage.py | 69 +++++++++++++++++++- 5 files changed, 235 insertions(+), 2 deletions(-) diff --git a/providers/dbt/cloud/docs/operators.rst b/providers/dbt/cloud/docs/operators.rst index a765035ef9f..80ed3a1b1a7 100644 --- a/providers/dbt/cloud/docs/operators.rst +++ b/providers/dbt/cloud/docs/operators.rst @@ -87,6 +87,20 @@ by providing the ``project_name``, ``environment_name``, and ``job_name``. Please note that it will only work if the above three parameters uniquely identify a job in your account (i.e. you cannot have two jobs with the same name in the same project and environment). +Linking dbt Cloud runs to the triggering Airflow task +"""""""""""""""""""""""""""""""""""""""""""""""""""""" + +Setting ``openlineage_inject_parent_job_info`` to True replaces the triggered run's ``cause`` field +with the Airflow task's OpenLineage parent (and root) run identifiers as JSON. A consumer that reads +dbt Cloud runs can then parse this JSON and attach a ``ParentRunFacet``, linking the dbt Cloud run back +to the Airflow task that triggered it. This overwrites the ``trigger_reason``. + +The parameter defaults to the ``openlineage.spark_inject_parent_job_info`` configuration value. Because +the dbt Cloud ``cause`` has a length limit, the injected payload is trimmed to fit: the essential +``parent`` link is always kept, the ``root`` block is dropped if the two together do not fit, and if +even the parent identifiers alone do not fit, injection is skipped and the original ``cause`` is left +unchanged. + .. exampleinclude:: /../tests/system/dbt/cloud/example_dbt_cloud.py :language: python :dedent: 4 diff --git a/providers/dbt/cloud/src/airflow/providers/dbt/cloud/operators/dbt.py b/providers/dbt/cloud/src/airflow/providers/dbt/cloud/operators/dbt.py index 6547fadb0f9..b5fa9c8807e 100644 --- a/providers/dbt/cloud/src/airflow/providers/dbt/cloud/operators/dbt.py +++ b/providers/dbt/cloud/src/airflow/providers/dbt/cloud/operators/dbt.py @@ -31,7 +31,10 @@ from airflow.providers.dbt.cloud.hooks.dbt import ( JobRunInfo, ) from airflow.providers.dbt.cloud.triggers.dbt import DbtCloudRunJobTrigger -from airflow.providers.dbt.cloud.utils.openlineage import generate_openlineage_events_from_dbt_cloud_run +from airflow.providers.dbt.cloud.utils.openlineage import ( + generate_openlineage_events_from_dbt_cloud_run, + inject_parent_job_information_into_dbt_cloud_cause, +) if TYPE_CHECKING: from airflow.providers.openlineage.extractors import OperatorLineage @@ -85,6 +88,11 @@ class DbtCloudRunJobOperator(BaseOperator): https://docs.getdbt.com/dbt-cloud/api-v2#/operations/Retry%20Failed%20Job :param deferrable: Run operator in the deferrable mode :param hook_params: Extra arguments passed to the DbtCloudHook constructor. + :param openlineage_inject_parent_job_info: If True, the triggered dbt Cloud run's ``cause`` field + is replaced with the OpenLineage parent job information of the Airflow task as JSON. A consumer + reading dbt Cloud runs can parse this to link the dbt Cloud run back to the Airflow task that + triggered it. Note this overwrites the ``trigger_reason``. Defaults to the + ``openlineage.spark_inject_parent_job_info`` config value. :param execution_timeout: Maximum time allowed for the task to run. If exceeded, the dbt Cloud job will be cancelled and the task will fail. When both ``execution_timeout`` and ``timeout`` are set, the earlier deadline takes precedence. @@ -126,6 +134,11 @@ class DbtCloudRunJobOperator(BaseOperator): retry_from_failure: bool = False, deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), hook_params: dict[str, Any] | None = None, + # Defaults to the Spark flag until OpenLineage exposes a transport-agnostic + # ``openlineage.inject_parent_job_info`` config to migrate to. + openlineage_inject_parent_job_info: bool = conf.getboolean( + "openlineage", "spark_inject_parent_job_info", fallback=False + ), **kwargs, ) -> None: super().__init__(**kwargs) @@ -147,6 +160,7 @@ class DbtCloudRunJobOperator(BaseOperator): self.retry_from_failure = retry_from_failure self.deferrable = deferrable self.hook_params = hook_params or {} + self.openlineage_inject_parent_job_info = openlineage_inject_parent_job_info def _resolve_job_id(self) -> int: if self.job_id is not None: @@ -219,6 +233,12 @@ class DbtCloudRunJobOperator(BaseOperator): f"Triggered via Apache Airflow by task {self.task_id!r} in the {self.dag.dag_id} DAG." ) + if self.openlineage_inject_parent_job_info: + self.log.info("Injecting OpenLineage parent job information into dbt Cloud run cause.") + self.trigger_reason = inject_parent_job_information_into_dbt_cloud_cause( + self.trigger_reason, context["ti"] + ) + self.job_id = self._resolve_job_id() self.run_id, job_run_url = self._get_or_trigger_run(context) diff --git a/providers/dbt/cloud/src/airflow/providers/dbt/cloud/utils/openlineage.py b/providers/dbt/cloud/src/airflow/providers/dbt/cloud/utils/openlineage.py index 188ca1255f6..d9336d07189 100644 --- a/providers/dbt/cloud/src/airflow/providers/dbt/cloud/utils/openlineage.py +++ b/providers/dbt/cloud/src/airflow/providers/dbt/cloud/utils/openlineage.py @@ -17,11 +17,13 @@ from __future__ import annotations import asyncio +import json import logging import re from typing import TYPE_CHECKING from airflow.providers.common.compat.openlineage.check import require_openlineage_version +from airflow.providers.dbt.cloud.hooks.dbt import DBT_CAUSE_MAX_LENGTH from airflow.providers.dbt.cloud.version_compat import AIRFLOW_V_3_0_PLUS if TYPE_CHECKING: @@ -97,6 +99,62 @@ def _get_parent_run_metadata(task_instance): ) +def _get_openlineage_parent_facet_dict(task_instance) -> dict: + """Build the OpenLineage ``ParentRunFacet`` payload (parent and root) for a task instance.""" + metadata = _get_parent_run_metadata(task_instance) + return { + "parent": { + "run": {"runId": metadata.run_id}, + "job": {"namespace": metadata.job_namespace, "name": metadata.job_name}, + }, + "root": { + "run": {"runId": metadata.root_parent_run_id}, + "job": { + "namespace": metadata.root_parent_job_namespace, + "name": metadata.root_parent_job_name, + }, + }, + } + + +def inject_parent_job_information_into_dbt_cloud_cause(cause: str | None, task_instance) -> str | None: + """ + Replace the dbt Cloud run ``cause`` with OpenLineage parent job information as JSON. + + This serializes the Airflow task's OpenLineage parent (and root) run identifiers into the + triggered run's ``cause`` field. A consumer that reads dbt Cloud runs can parse this JSON and + attach a ``ParentRunFacet``, linking the dbt Cloud run back to the Airflow task that triggered it. + + The dbt Cloud ``cause`` is limited to ``DBT_CAUSE_MAX_LENGTH`` characters. The essential ``parent`` + link is always kept; the ``root`` block is dropped if the two together do not fit. If even the + parent-only payload does not fit, the original ``cause`` is returned unchanged. + """ + try: + facet = _get_openlineage_parent_facet_dict(task_instance) + except ImportError: + log.warning( + "Could not import OpenLineage provider. Skipping the injection of OpenLineage " + "parent job information into the dbt Cloud run cause." + ) + return cause + + candidates = [ + {"parent": facet["parent"], "root": facet["root"]}, + {"parent": facet["parent"]}, + ] + for payload in candidates: + serialized = json.dumps(payload, separators=(",", ":")) + if len(serialized) <= DBT_CAUSE_MAX_LENGTH: + return serialized + + log.warning( + "OpenLineage parent job information exceeds the dbt Cloud cause limit of %s characters. " + "Skipping the injection of OpenLineage parent job information into the dbt Cloud run cause.", + DBT_CAUSE_MAX_LENGTH, + ) + return cause + + @require_openlineage_version(provider_min_version="2.5.0") def generate_openlineage_events_from_dbt_cloud_run( operator: DbtCloudRunJobOperator | DbtCloudJobRunSensor, task_instance: TaskInstance diff --git a/providers/dbt/cloud/tests/unit/dbt/cloud/operators/test_dbt.py b/providers/dbt/cloud/tests/unit/dbt/cloud/operators/test_dbt.py index bb5988acdea..083fa049bbb 100644 --- a/providers/dbt/cloud/tests/unit/dbt/cloud/operators/test_dbt.py +++ b/providers/dbt/cloud/tests/unit/dbt/cloud/operators/test_dbt.py @@ -568,6 +568,80 @@ class TestDbtCloudRunJobOperator: max_number_of_calls = timeout // self.config["check_interval"] + 1 assert mock_get_job_run.call_count <= max_number_of_calls + @patch( + "airflow.providers.dbt.cloud.operators.dbt.inject_parent_job_information_into_dbt_cloud_cause", + return_value='{"parent": "info"}', + ) + @patch.object( + DbtCloudHook, "trigger_job_run", return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE) + ) + def test_execute_injects_openlineage_parent_job_info_into_cause(self, mock_run_job, mock_inject): + operator = DbtCloudRunJobOperator( + task_id=TASK_ID, + dbt_cloud_conn_id=ACCOUNT_ID_CONN, + job_id=JOB_ID, + dag=self.dag, + wait_for_termination=False, + openlineage_inject_parent_job_info=True, + ) + + operator.execute(context=self.mock_context) + + mock_inject.assert_called_once_with( + f"Triggered via Apache Airflow by task {TASK_ID!r} in the {self.dag.dag_id} DAG.", + self.mock_ti, + ) + assert mock_run_job.call_args.kwargs["cause"] == '{"parent": "info"}' + + @patch( + "airflow.providers.dbt.cloud.operators.dbt.inject_parent_job_information_into_dbt_cloud_cause", + side_effect=lambda cause, ti: cause, + ) + @patch.object( + DbtCloudHook, "trigger_job_run", return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE) + ) + def test_execute_injection_skipped_keeps_original_cause(self, mock_run_job, mock_inject): + # When the identifiers do not fit, the helper returns the cause unchanged. + operator = DbtCloudRunJobOperator( + task_id=TASK_ID, + dbt_cloud_conn_id=ACCOUNT_ID_CONN, + job_id=JOB_ID, + dag=self.dag, + wait_for_termination=False, + openlineage_inject_parent_job_info=True, + ) + + operator.execute(context=self.mock_context) + + assert ( + mock_run_job.call_args.kwargs["cause"] + == f"Triggered via Apache Airflow by task {TASK_ID!r} in the {self.dag.dag_id} DAG." + ) + + @patch( + "airflow.providers.dbt.cloud.operators.dbt.inject_parent_job_information_into_dbt_cloud_cause", + ) + @patch.object( + DbtCloudHook, "trigger_job_run", return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE) + ) + def test_execute_does_not_inject_parent_job_info_by_default(self, mock_run_job, mock_inject): + operator = DbtCloudRunJobOperator( + task_id=TASK_ID, + dbt_cloud_conn_id=ACCOUNT_ID_CONN, + job_id=JOB_ID, + dag=self.dag, + wait_for_termination=False, + ) + + assert operator.openlineage_inject_parent_job_info is False + operator.execute(context=self.mock_context) + + mock_inject.assert_not_called() + assert ( + mock_run_job.call_args.kwargs["cause"] + == f"Triggered via Apache Airflow by task {TASK_ID!r} in the {self.dag.dag_id} DAG." + ) + @patch.object(DbtCloudHook, "trigger_job_run") @pytest.mark.parametrize( ("conn_id", "account_id"), diff --git a/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py b/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py index 32e7342c3e5..cd911109c0f 100644 --- a/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py +++ b/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py @@ -18,6 +18,7 @@ from __future__ import annotations import json from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -25,11 +26,12 @@ from openlineage.client.constants import __version__ from packaging.version import parse from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException -from airflow.providers.dbt.cloud.hooks.dbt import DbtCloudHook +from airflow.providers.dbt.cloud.hooks.dbt import DBT_CAUSE_MAX_LENGTH, DbtCloudHook from airflow.providers.dbt.cloud.operators.dbt import DbtCloudRunJobOperator from airflow.providers.dbt.cloud.utils.openlineage import ( _get_parent_run_metadata, generate_openlineage_events_from_dbt_cloud_run, + inject_parent_job_information_into_dbt_cloud_cause, ) from airflow.providers.openlineage.conf import namespace from airflow.providers.openlineage.extractors import OperatorLineage @@ -223,3 +225,68 @@ class TestGenerateOpenLineageEventsFromDbtCloudRun: operator = DbtCloudRunJobOperator(task_id="dbt-job-runid-taskid", job_id=1500) assert operator.run_id is None assert operator.get_openlineage_facets_on_complete(MagicMock()) == OperatorLineage() + + +PARENT_RUN_ID = "01941f29-7c00-7087-8906-40e512c257bd" +ROOT_RUN_ID = "01941f29-7c00-743e-b109-28b18d0a19c5" + + +def _fake_metadata(namespace="default", job_name="dag.task", root_job_name="dag"): + return SimpleNamespace( + run_id=PARENT_RUN_ID, + job_namespace=namespace, + job_name=job_name, + root_parent_run_id=ROOT_RUN_ID, + root_parent_job_namespace=namespace, + root_parent_job_name=root_job_name, + ) + + +class TestInjectParentJobInformationIntoDbtCloudCause: + @patch("airflow.providers.dbt.cloud.utils.openlineage._get_parent_run_metadata") + def test_parent_and_root_when_it_fits(self, mock_metadata): + mock_metadata.return_value = _fake_metadata() + + result = inject_parent_job_information_into_dbt_cloud_cause("my reason", MagicMock()) + + payload = json.loads(result) + assert payload == { + "parent": { + "run": {"runId": PARENT_RUN_ID}, + "job": {"namespace": "default", "name": "dag.task"}, + }, + "root": { + "run": {"runId": ROOT_RUN_ID}, + "job": {"namespace": "default", "name": "dag"}, + }, + } + # The human-readable cause is never embedded, only the lineage identifiers. + assert "reason" not in payload + assert "my reason" not in result + + @patch("airflow.providers.dbt.cloud.utils.openlineage._get_parent_run_metadata") + def test_root_dropped_when_parent_and_root_too_long(self, mock_metadata): + mock_metadata.return_value = _fake_metadata(namespace="N" * 90) + + result = inject_parent_job_information_into_dbt_cloud_cause("original cause", MagicMock()) + + payload = json.loads(result) + assert len(result) <= DBT_CAUSE_MAX_LENGTH + assert "root" not in payload + assert payload["parent"]["run"]["runId"] == PARENT_RUN_ID + + @patch("airflow.providers.dbt.cloud.utils.openlineage._get_parent_run_metadata") + def test_injection_skipped_when_parent_alone_too_long(self, mock_metadata): + mock_metadata.return_value = _fake_metadata(namespace="N" * 250) + + result = inject_parent_job_information_into_dbt_cloud_cause("original cause", MagicMock()) + + assert result == "original cause" + + @patch("airflow.providers.dbt.cloud.utils.openlineage._get_parent_run_metadata") + def test_injection_skipped_when_openlineage_unavailable(self, mock_metadata): + mock_metadata.side_effect = ImportError("no openlineage") + + result = inject_parent_job_information_into_dbt_cloud_cause("original cause", MagicMock()) + + assert result == "original cause"
