From 57bc4ed33ec5ee27d465e5a955c473448293ff55 Mon Sep 17 00:00:00 2001 From: rahul-madaan Date: Tue, 2 Jun 2026 15:20:03 +0530 Subject: [PATCH 1/3] Add OpenLineage parent/transport info injection to DatabricksSubmitRunOperator DatabricksSubmitRunOperator did not emit any OpenLineage information. This adds optional injection of OpenLineage parent job and transport configuration into the job's ``new_cluster.spark_conf`` (single-task and multi-task forms), so the Spark job running on Databricks can correlate its lineage events with the Airflow task and send them to the same backend. The behaviour is controlled by two new operator parameters, ``openlineage_inject_parent_job_info`` and ``openlineage_inject_transport_info``, each defaulting to the corresponding ``openlineage.spark_inject_*_info`` config option, mirroring the existing Dataproc, EMR and Glue operators. Injection is skipped when the provider is unavailable, when the relevant properties are already present, or when the job has no ``new_cluster`` to modify (e.g. an existing cluster). Signed-off-by: rahul-madaan --- .../databricks/operators/databricks.py | 39 +++++ .../providers/databricks/utils/openlineage.py | 131 ++++++++++++++++ .../databricks/operators/test_databricks.py | 133 ++++++++++++++++ .../unit/databricks/utils/test_openlineage.py | 144 ++++++++++++++++++ providers/openlineage/docs/spark.rst | 1 + 5 files changed, 448 insertions(+) diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py index bf743aaf5166b..4bae6b28f0aa1 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py @@ -652,6 +652,13 @@ class DatabricksSubmitRunOperator(BaseOperator): .. seealso:: https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsRunsSubmit + :param openlineage_inject_parent_job_info: If True, injects OpenLineage parent job information + into the ``new_cluster`` ``spark_conf`` so the Spark job emits a ``parentRunFacet`` linking + back to the Airflow task. Defaults to the + ``openlineage.spark_inject_parent_job_info`` config value. + :param openlineage_inject_transport_info: If True, injects OpenLineage transport configuration + into the ``new_cluster`` ``spark_conf`` so the Spark job sends OL events to the same backend + as Airflow. Defaults to the ``openlineage.spark_inject_transport_info`` config value. .. note:: If the operator's ``params`` dict is non-empty, it is automatically forwarded into the @@ -716,6 +723,12 @@ def __init__( wait_for_termination: bool = True, git_source: dict[str, str] | None = None, deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + openlineage_inject_parent_job_info: bool = conf.getboolean( + "openlineage", "spark_inject_parent_job_info", fallback=False + ), + openlineage_inject_transport_info: bool = conf.getboolean( + "openlineage", "spark_inject_transport_info", fallback=False + ), **kwargs, ) -> None: """Create a new ``DatabricksSubmitRunOperator``.""" @@ -743,6 +756,8 @@ def __init__( self.databricks_retry_args = databricks_retry_args self.wait_for_termination = wait_for_termination self.deferrable = deferrable + self.openlineage_inject_parent_job_info = openlineage_inject_parent_job_info + self.openlineage_inject_transport_info = openlineage_inject_transport_info # This variable will be used in case our task gets killed. self.run_id: int | None = None @@ -830,6 +845,10 @@ def execute(self, context: Context): else: _inject_airflow_params_into_task(json, params_dump) + if self.openlineage_inject_parent_job_info or self.openlineage_inject_transport_info: + self.log.info("Automatic injection of OpenLineage information into Spark properties is enabled.") + json = self._inject_openlineage_properties_into_databricks_job(json, context) + normalised = cast("dict[str, Any]", normalise_json_content(json)) self.run_id = self._hook.submit_run(normalised) if self.deferrable: @@ -837,6 +856,26 @@ def execute(self, context: Context): else: _handle_databricks_operator_execution(self, self._hook, self.log, context) + def _inject_openlineage_properties_into_databricks_job(self, json: dict, context: Context) -> dict: + try: + from airflow.providers.databricks.utils.openlineage import ( + inject_openlineage_properties_into_databricks_job, + ) + + return inject_openlineage_properties_into_databricks_job( + job=json, + context=context, + inject_parent_job_info=self.openlineage_inject_parent_job_info, + inject_transport_info=self.openlineage_inject_transport_info, + ) + except Exception as e: + self.log.warning( + "An error occurred while trying to inject OpenLineage information. " + "Databricks job has not been modified by OpenLineage.", + exc_info=e, + ) + return json + def on_kill(self): if self.run_id: self._hook.cancel_run(self.run_id) diff --git a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py index 58e87a21dd7e9..1d61019b84140 100644 --- a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py +++ b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import copy import datetime import json import logging @@ -24,12 +25,17 @@ import requests from airflow.providers.common.compat.openlineage.check import require_openlineage_version +from airflow.providers.common.compat.openlineage.utils.spark import ( + inject_parent_job_information_into_spark_properties, + inject_transport_information_into_spark_properties, +) from airflow.providers.common.compat.sdk import timezone if TYPE_CHECKING: from openlineage.client.event_v2 import RunEvent from openlineage.client.facet_v2 import JobFacet + from airflow.providers.common.compat.sdk import Context from airflow.providers.databricks.hooks.databricks import DatabricksHook from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook @@ -320,3 +326,128 @@ def emit_openlineage_events_for_databricks_queries( log.info("OpenLineage has successfully finished processing information about Databricks queries.") return + + +def _is_openlineage_provider_accessible() -> bool: + """ + Check if the OpenLineage provider is accessible. + + This function attempts to import the necessary OpenLineage modules and checks if the provider + is enabled and the listener is available. + + Returns: + True if the OpenLineage provider is accessible, False otherwise. + """ + try: + from airflow.providers.openlineage.conf import is_disabled + from airflow.providers.openlineage.plugins.listener import get_openlineage_listener + except ImportError: + log.debug("OpenLineage provider could not be imported.") + return False + + if is_disabled(): + log.debug("OpenLineage provider is disabled.") + return False + + if not get_openlineage_listener(): + log.debug("OpenLineage listener could not be found.") + return False + + return True + + +def _extract_new_clusters_from_databricks_job(job: dict) -> list[dict]: + """ + Collect every ``new_cluster`` definition that can carry Spark properties in a Databricks job. + + A ``runs/submit`` payload can define a ``new_cluster`` in three places: at the top level + (single-task form), inline on each task (``tasks[].new_cluster``), or as a shared job cluster + (``job_clusters[].new_cluster``) referenced by tasks through ``job_cluster_key``. Tasks running on + an ``existing_cluster_id`` have no ``new_cluster`` to mutate and are skipped. + + Args: + job: The Databricks ``runs/submit`` job definition. + + Returns: + The list of ``new_cluster`` dicts found in the job definition. + """ + new_clusters = [] + if isinstance(job.get("new_cluster"), dict): + new_clusters.append(job["new_cluster"]) + for key in ("tasks", "job_clusters"): + if isinstance(job.get(key), list): + new_clusters.extend( + item["new_cluster"] for item in job[key] if isinstance(item.get("new_cluster"), dict) + ) + return new_clusters + + +def inject_openlineage_properties_into_databricks_job( + job: dict, context: Context, inject_parent_job_info: bool, inject_transport_info: bool +) -> dict: + """ + Inject OpenLineage properties into a Databricks job definition. + + This function does not remove existing configurations or modify the job definition in any way, + except to add the required OpenLineage properties if they are not already present. + + The entire properties injection process will be skipped if any condition is met: + - The OpenLineage provider is not accessible. + - The job has no ``new_cluster`` definition to inject Spark properties into (e.g. it only uses + an ``existing_cluster_id``, whose Spark configuration is fixed at cluster creation time). + - Both `inject_parent_job_info` and `inject_transport_info` are set to False. + + Additionally, specific information will not be injected if relevant OpenLineage properties already + exist. + + Parent job information will not be injected if: + - Any property prefixed with `spark.openlineage.parent` exists. + - `inject_parent_job_info` is False. + Transport information will not be injected if: + - Any property prefixed with `spark.openlineage.transport` exists. + - `inject_transport_info` is False. + + Args: + job: The original Databricks ``runs/submit`` job definition. + context: The Airflow context in which the job is running. + inject_parent_job_info: Flag indicating whether to inject parent job information. + inject_transport_info: Flag indicating whether to inject transport information. + + Returns: + The modified job definition with OpenLineage properties injected, if applicable. + """ + if not inject_parent_job_info and not inject_transport_info: + log.debug("Automatic injection of OpenLineage information is disabled.") + return job + + if not _is_openlineage_provider_accessible(): + log.warning( + "Could not access OpenLineage provider for automatic OpenLineage " + "properties injection. No action will be performed." + ) + return job + + job = copy.deepcopy(job) + new_clusters = _extract_new_clusters_from_databricks_job(job) + if not new_clusters: + log.warning( + "Could not find a Databricks `new_cluster` definition for automatic OpenLineage " + "properties injection. No action will be performed." + ) + return job + + for new_cluster in new_clusters: + properties = new_cluster.get("spark_conf", {}) + if inject_parent_job_info: + log.debug("Injecting OpenLineage parent job information into Spark properties.") + properties = inject_parent_job_information_into_spark_properties( + properties=properties, context=context + ) + if inject_transport_info: + log.debug("Injecting OpenLineage transport information into Spark properties.") + properties = inject_transport_information_into_spark_properties( + properties=properties, context=context + ) + new_cluster["spark_conf"] = properties + + return job diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks.py b/providers/databricks/tests/unit/databricks/operators/test_databricks.py index 3d613712caa2a..f44f1e9fbe0a1 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py @@ -1518,6 +1518,139 @@ def test_exec_with_rendered_pipeline_id_and_name_raises(self, db_mock_class): db_mock_class.assert_not_called() +class TestDatabricksSubmitRunOperatorOpenLineageInjection: + """Tests for OpenLineage parent job info and transport info injection in DatabricksSubmitRunOperator.""" + + @mock.patch( + "airflow.providers.databricks.utils.openlineage.inject_openlineage_properties_into_databricks_job" + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_inject_parent_job_info_called_when_enabled(self, db_mock_class, mock_inject): + mock_inject.side_effect = lambda job, context, inject_parent_job_info, inject_transport_info: { + **job, + "new_cluster": { + **job["new_cluster"], + "spark_conf": {"spark.openlineage.parentJobNamespace": "ns"}, + }, + } + op = DatabricksSubmitRunOperator( + task_id=TASK_ID, + new_cluster=NEW_CLUSTER, + notebook_task=NOTEBOOK_TASK, + openlineage_inject_parent_job_info=True, + ) + db_mock = db_mock_class.return_value + db_mock.submit_run.return_value = RUN_ID + db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS") + + op.execute(None) + + mock_inject.assert_called_once() + submitted = db_mock.submit_run.call_args.args[0] + assert submitted["new_cluster"]["spark_conf"]["spark.openlineage.parentJobNamespace"] == "ns" + + @mock.patch( + "airflow.providers.databricks.utils.openlineage.inject_openlineage_properties_into_databricks_job" + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_inject_parent_job_info_not_called_when_disabled(self, db_mock_class, mock_inject): + op = DatabricksSubmitRunOperator( + task_id=TASK_ID, + new_cluster=NEW_CLUSTER, + notebook_task=NOTEBOOK_TASK, + openlineage_inject_parent_job_info=False, + ) + db_mock = db_mock_class.return_value + db_mock.submit_run.return_value = RUN_ID + db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS") + + op.execute(None) + + mock_inject.assert_not_called() + + @mock.patch( + "airflow.providers.databricks.utils.openlineage.inject_openlineage_properties_into_databricks_job" + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_inject_transport_info_called_when_enabled(self, db_mock_class, mock_inject): + mock_inject.side_effect = lambda job, context, inject_parent_job_info, inject_transport_info: { + **job, + "new_cluster": {**job["new_cluster"], "spark_conf": {"spark.openlineage.transport.type": "http"}}, + } + op = DatabricksSubmitRunOperator( + task_id=TASK_ID, + new_cluster=NEW_CLUSTER, + notebook_task=NOTEBOOK_TASK, + openlineage_inject_transport_info=True, + ) + db_mock = db_mock_class.return_value + db_mock.submit_run.return_value = RUN_ID + db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS") + + op.execute(None) + + mock_inject.assert_called_once() + submitted = db_mock.submit_run.call_args.args[0] + assert submitted["new_cluster"]["spark_conf"]["spark.openlineage.transport.type"] == "http" + + @mock.patch( + "airflow.providers.databricks.utils.openlineage.inject_openlineage_properties_into_databricks_job" + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_inject_both_parent_and_transport_info(self, db_mock_class, mock_inject): + mock_inject.side_effect = lambda job, context, inject_parent_job_info, inject_transport_info: job + op = DatabricksSubmitRunOperator( + task_id=TASK_ID, + new_cluster=NEW_CLUSTER, + notebook_task=NOTEBOOK_TASK, + openlineage_inject_parent_job_info=True, + openlineage_inject_transport_info=True, + ) + db_mock = db_mock_class.return_value + db_mock.submit_run.return_value = RUN_ID + db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS") + + op.execute(None) + + mock_inject.assert_called_once() + _, call_kwargs = mock_inject.call_args + assert call_kwargs["inject_parent_job_info"] is True + assert call_kwargs["inject_transport_info"] is True + + @mock.patch( + "airflow.providers.databricks.utils.openlineage.inject_openlineage_properties_into_databricks_job" + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_inject_parent_job_info_preserves_existing_config(self, db_mock_class, mock_inject): + """Existing ``spark_conf`` entries are preserved alongside the injected OpenLineage properties.""" + new_cluster = {**NEW_CLUSTER, "spark_conf": {"spark.executor.memory": "8g"}} + mock_inject.side_effect = lambda job, context, inject_parent_job_info, inject_transport_info: { + **job, + "new_cluster": { + **job["new_cluster"], + "spark_conf": { + **job["new_cluster"]["spark_conf"], + "spark.openlineage.parentJobNamespace": "ns", + }, + }, + } + op = DatabricksSubmitRunOperator( + task_id=TASK_ID, + new_cluster=new_cluster, + notebook_task=NOTEBOOK_TASK, + openlineage_inject_parent_job_info=True, + ) + db_mock = db_mock_class.return_value + db_mock.submit_run.return_value = RUN_ID + db_mock.get_run = make_run_with_state_mock("TERMINATED", "SUCCESS") + + op.execute(None) + + submitted = db_mock.submit_run.call_args.args[0] + assert submitted["new_cluster"]["spark_conf"]["spark.executor.memory"] == "8g" + assert submitted["new_cluster"]["spark_conf"]["spark.openlineage.parentJobNamespace"] == "ns" + + class TestDatabricksRunNowOperator: def test_init_with_named_parameters(self): """ diff --git a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py index 8702456ca8629..00eb4f1007d56 100644 --- a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py +++ b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py @@ -34,11 +34,13 @@ from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook from airflow.providers.databricks.utils.openlineage import ( _create_ol_event_pair, + _extract_new_clusters_from_databricks_job, _get_parent_run_facet, _get_queries_details_from_databricks, _process_data_from_api, _run_api_call, emit_openlineage_events_for_databricks_queries, + inject_openlineage_properties_into_databricks_job, ) from airflow.providers.openlineage.conf import namespace from airflow.utils.state import TaskInstanceState @@ -1196,3 +1198,145 @@ def test_emit_openlineage_events_with_old_openlineage_provider(mock_version): ) assert query_ids == original_query_ids # Verify that the input query_ids list is unchanged. fake_adapter.emit.assert_not_called() # No events should be emitted + + +OL_UTILS = "airflow.providers.databricks.utils.openlineage" + + +def test_extract_new_clusters_from_databricks_job(): + top_cluster = {"spark_version": "13.3.x-scala2.12"} + task_cluster = {"spark_version": "14.3.x-scala2.12"} + job_cluster = {"spark_version": "15.4.x-scala2.12"} + job = { + "new_cluster": top_cluster, + "tasks": [ + {"task_key": "a", "new_cluster": task_cluster}, + {"task_key": "b", "existing_cluster_id": "existing"}, + {"task_key": "c", "job_cluster_key": "shared"}, + ], + "job_clusters": [{"job_cluster_key": "shared", "new_cluster": job_cluster}], + } + assert _extract_new_clusters_from_databricks_job(job) == [top_cluster, task_cluster, job_cluster] + + +def test_extract_new_clusters_from_databricks_job_existing_cluster_only(): + assert _extract_new_clusters_from_databricks_job({"existing_cluster_id": "existing"}) == [] + + +@mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=True) +def test_inject_openlineage_properties_disabled_returns_job_unchanged(mock_accessible): + job = {"new_cluster": {"spark_conf": {}}} + result = inject_openlineage_properties_into_databricks_job( + job, context=mock.MagicMock(), inject_parent_job_info=False, inject_transport_info=False + ) + assert result == job + mock_accessible.assert_not_called() + + +@mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=False) +def test_inject_openlineage_properties_provider_inaccessible_returns_job_unchanged(mock_accessible): + job = {"new_cluster": {"spark_conf": {}}} + result = inject_openlineage_properties_into_databricks_job( + job, context=mock.MagicMock(), inject_parent_job_info=True, inject_transport_info=True + ) + assert result == job + + +@mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=True) +def test_inject_openlineage_properties_no_new_cluster_returns_job_unchanged(mock_accessible): + job = {"existing_cluster_id": "existing"} + result = inject_openlineage_properties_into_databricks_job( + job, context=mock.MagicMock(), inject_parent_job_info=True, inject_transport_info=True + ) + assert result == job + + +@mock.patch(f"{OL_UTILS}.inject_transport_information_into_spark_properties") +@mock.patch(f"{OL_UTILS}.inject_parent_job_information_into_spark_properties") +@mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=True) +def test_inject_openlineage_properties_into_top_level_new_cluster( + mock_accessible, mock_parent, mock_transport +): + mock_parent.side_effect = lambda properties, context: { + **properties, + "spark.openlineage.parentJobName": "dag_id.task_id", + } + mock_transport.side_effect = lambda properties, context: { + **properties, + "spark.openlineage.transport.type": "http", + } + job = {"new_cluster": {"spark_conf": {"spark.executor.memory": "8g"}}, "notebook_task": {"x": "y"}} + + result = inject_openlineage_properties_into_databricks_job( + job, context=mock.MagicMock(), inject_parent_job_info=True, inject_transport_info=True + ) + + spark_conf = result["new_cluster"]["spark_conf"] + assert spark_conf["spark.executor.memory"] == "8g" + assert spark_conf["spark.openlineage.parentJobName"] == "dag_id.task_id" + assert spark_conf["spark.openlineage.transport.type"] == "http" + + +@mock.patch(f"{OL_UTILS}.inject_parent_job_information_into_spark_properties") +@mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=True) +def test_inject_openlineage_properties_into_multi_task_new_clusters(mock_accessible, mock_parent): + mock_parent.side_effect = lambda properties, context: { + **properties, + "spark.openlineage.parentJobName": "dag_id.task_id", + } + job = { + "tasks": [ + {"task_key": "a", "new_cluster": {"spark_conf": {}}}, + {"task_key": "b", "new_cluster": {"spark_conf": {}}}, + {"task_key": "c", "existing_cluster_id": "existing"}, + ] + } + + result = inject_openlineage_properties_into_databricks_job( + job, context=mock.MagicMock(), inject_parent_job_info=True, inject_transport_info=False + ) + + assert result["tasks"][0]["new_cluster"]["spark_conf"]["spark.openlineage.parentJobName"] == ( + "dag_id.task_id" + ) + assert result["tasks"][1]["new_cluster"]["spark_conf"]["spark.openlineage.parentJobName"] == ( + "dag_id.task_id" + ) + assert result["tasks"][2] == {"task_key": "c", "existing_cluster_id": "existing"} + + +@mock.patch(f"{OL_UTILS}.inject_parent_job_information_into_spark_properties") +@mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=True) +def test_inject_openlineage_properties_into_shared_job_clusters(mock_accessible, mock_parent): + mock_parent.side_effect = lambda properties, context: { + **properties, + "spark.openlineage.parentJobName": "dag_id.task_id", + } + job = { + "tasks": [{"task_key": "a", "job_cluster_key": "shared"}], + "job_clusters": [{"job_cluster_key": "shared", "new_cluster": {"spark_conf": {}}}], + } + + result = inject_openlineage_properties_into_databricks_job( + job, context=mock.MagicMock(), inject_parent_job_info=True, inject_transport_info=False + ) + + injected = result["job_clusters"][0]["new_cluster"]["spark_conf"]["spark.openlineage.parentJobName"] + assert injected == "dag_id.task_id" + + +@mock.patch(f"{OL_UTILS}.inject_parent_job_information_into_spark_properties") +@mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=True) +def test_inject_openlineage_properties_does_not_mutate_input(mock_accessible, mock_parent): + mock_parent.side_effect = lambda properties, context: { + **properties, + "spark.openlineage.parentJobName": "dag_id.task_id", + } + job = {"new_cluster": {"spark_conf": {"spark.executor.memory": "8g"}}} + original = copy.deepcopy(job) + + inject_openlineage_properties_into_databricks_job( + job, context=mock.MagicMock(), inject_parent_job_info=True, inject_transport_info=False + ) + + assert job == original diff --git a/providers/openlineage/docs/spark.rst b/providers/openlineage/docs/spark.rst index f4301525fee00..0dbdf967c142d 100644 --- a/providers/openlineage/docs/spark.rst +++ b/providers/openlineage/docs/spark.rst @@ -96,6 +96,7 @@ Automatic injection is supported for the following operators: - :class:`~airflow.providers.apache.livy.operators.livy.LivyOperator` - :class:`~airflow.providers.apache.spark.operators.spark_submit.SparkSubmitOperator` +- :class:`~airflow.providers.databricks.operators.databricks.DatabricksSubmitRunOperator` - :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocCreateBatchOperator` - :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocInstantiateInlineWorkflowTemplateOperator` - :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocSubmitJobOperator` From 17c9b8329b67e250afadaa582b08b1b534f0fc35 Mon Sep 17 00:00:00 2001 From: rahul-madaan Date: Wed, 3 Jun 2026 22:40:54 +0530 Subject: [PATCH 2/3] Use debug log level when skipping Databricks OpenLineage injection without a new_cluster Signed-off-by: rahul-madaan --- .../src/airflow/providers/databricks/utils/openlineage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py index 1d61019b84140..463c1ace8126d 100644 --- a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py +++ b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py @@ -430,7 +430,7 @@ def inject_openlineage_properties_into_databricks_job( job = copy.deepcopy(job) new_clusters = _extract_new_clusters_from_databricks_job(job) if not new_clusters: - log.warning( + log.debug( "Could not find a Databricks `new_cluster` definition for automatic OpenLineage " "properties injection. No action will be performed." ) From 38e1d8a50cb2b5fba00c803f58d55bc540c9a264 Mon Sep 17 00:00:00 2001 From: rahul-madaan Date: Wed, 17 Jun 2026 22:57:03 +0530 Subject: [PATCH 3/3] Inject OpenLineage properties into Databricks for_each_task nested clusters Signed-off-by: rahul-madaan --- .../providers/databricks/utils/openlineage.py | 12 +++++-- .../unit/databricks/utils/test_openlineage.py | 36 ++++++++++++++++++- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py index 463c1ace8126d..706afa5df09e2 100644 --- a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py +++ b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py @@ -360,8 +360,9 @@ def _extract_new_clusters_from_databricks_job(job: dict) -> list[dict]: """ Collect every ``new_cluster`` definition that can carry Spark properties in a Databricks job. - A ``runs/submit`` payload can define a ``new_cluster`` in three places: at the top level - (single-task form), inline on each task (``tasks[].new_cluster``), or as a shared job cluster + A ``runs/submit`` payload can define a ``new_cluster`` at the top level (single-task form), inline + on each task (``tasks[].new_cluster``), inside a ``for_each_task``'s nested task + (``tasks[].for_each_task.task.new_cluster``), or as a shared job cluster (``job_clusters[].new_cluster``) referenced by tasks through ``job_cluster_key``. Tasks running on an ``existing_cluster_id`` have no ``new_cluster`` to mutate and are skipped. @@ -379,6 +380,13 @@ def _extract_new_clusters_from_databricks_job(job: dict) -> list[dict]: new_clusters.extend( item["new_cluster"] for item in job[key] if isinstance(item.get("new_cluster"), dict) ) + if isinstance(job.get("tasks"), list): + for task in job["tasks"]: + # ``for_each_task`` wraps a nested task that can carry its own ``new_cluster``. + for_each = task.get("for_each_task") if isinstance(task, dict) else None + nested = for_each.get("task") if isinstance(for_each, dict) else None + if isinstance(nested, dict) and isinstance(nested.get("new_cluster"), dict): + new_clusters.append(nested["new_cluster"]) return new_clusters diff --git a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py index 00eb4f1007d56..cefa881c7ecfe 100644 --- a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py +++ b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py @@ -1207,16 +1207,26 @@ def test_extract_new_clusters_from_databricks_job(): top_cluster = {"spark_version": "13.3.x-scala2.12"} task_cluster = {"spark_version": "14.3.x-scala2.12"} job_cluster = {"spark_version": "15.4.x-scala2.12"} + for_each_cluster = {"spark_version": "16.4.x-scala2.12"} job = { "new_cluster": top_cluster, "tasks": [ {"task_key": "a", "new_cluster": task_cluster}, {"task_key": "b", "existing_cluster_id": "existing"}, {"task_key": "c", "job_cluster_key": "shared"}, + { + "task_key": "d", + "for_each_task": {"task": {"task_key": "d/iter", "new_cluster": for_each_cluster}}, + }, ], "job_clusters": [{"job_cluster_key": "shared", "new_cluster": job_cluster}], } - assert _extract_new_clusters_from_databricks_job(job) == [top_cluster, task_cluster, job_cluster] + assert _extract_new_clusters_from_databricks_job(job) == [ + top_cluster, + task_cluster, + job_cluster, + for_each_cluster, + ] def test_extract_new_clusters_from_databricks_job_existing_cluster_only(): @@ -1325,6 +1335,30 @@ def test_inject_openlineage_properties_into_shared_job_clusters(mock_accessible, assert injected == "dag_id.task_id" +@mock.patch(f"{OL_UTILS}.inject_parent_job_information_into_spark_properties") +@mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=True) +def test_inject_openlineage_properties_into_for_each_task_new_cluster(mock_accessible, mock_parent): + mock_parent.side_effect = lambda properties, context: { + **properties, + "spark.openlineage.parentJobName": "dag_id.task_id", + } + job = { + "tasks": [ + { + "task_key": "a", + "for_each_task": {"task": {"task_key": "a/iter", "new_cluster": {"spark_conf": {}}}}, + } + ] + } + + result = inject_openlineage_properties_into_databricks_job( + job, context=mock.MagicMock(), inject_parent_job_info=True, inject_transport_info=False + ) + + injected = result["tasks"][0]["for_each_task"]["task"]["new_cluster"]["spark_conf"] + assert injected["spark.openlineage.parentJobName"] == "dag_id.task_id" + + @mock.patch(f"{OL_UTILS}.inject_parent_job_information_into_spark_properties") @mock.patch(f"{OL_UTILS}._is_openlineage_provider_accessible", return_value=True) def test_inject_openlineage_properties_does_not_mutate_input(mock_accessible, mock_parent):