This is an automated email from the ASF dual-hosted git repository.
mobuchowski pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d8d78647229 Propagate OpenLineage context to Databricks job runs
(#72643)
d8d78647229 is described below
commit d8d7864722975bb3ad979016b7d9f45146a0ae27
Author: Kacper Muda <[email protected]>
AuthorDate: Wed Sep 9 15:30:37 2026 +0200
Propagate OpenLineage context to Databricks job runs (#72643)
---
providers/databricks/docs/operators/run_now.rst | 29 +++
providers/databricks/docs/operators/submit_run.rst | 16 ++
.../providers/databricks/operators/databricks.py | 93 ++++++++-
.../providers/databricks/utils/openlineage.py | 101 ++++++++++
.../unit/databricks/operators/test_databricks.py | 211 +++++++++++++++++++++
.../unit/databricks/utils/test_openlineage.py | 140 ++++++++++++++
providers/openlineage/docs/spark.rst | 6 +
7 files changed, 586 insertions(+), 10 deletions(-)
diff --git a/providers/databricks/docs/operators/run_now.rst
b/providers/databricks/docs/operators/run_now.rst
index 8681934e2b6..5e706e15dba 100644
--- a/providers/databricks/docs/operators/run_now.rst
+++ b/providers/databricks/docs/operators/run_now.rst
@@ -81,6 +81,35 @@ disable this parameter forwarding behavior.
# job_parameters={"env": "staging", "batch_size": "42"}
# i.e. the same dict, passed straight through to the run-now request body.
+OpenLineage parent job information
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Set ``openlineage_inject_parent_job_info=True`` to add the standardized
OpenLineage context to the
+run's ``job_parameters`` under the ``OPENLINEAGE_CONTEXT`` key. The JSON value
contains the Airflow
+task as the parent job, with ``jobType`` set to ``BATCH/AIRFLOW/TASK``. When
the current Airflow Dag
+is the root job, its ``jobType`` is set to ``BATCH/AIRFLOW/DAG``. For a root
inherited through
+``DagRun.conf``, the job type is copied from
``conf["openlineage"]["rootParentJobType"]`` when that
+mapping is present; otherwise, the root job has no ``jobType`` facet. This
follows the context format
+introduced in `OpenLineage #4682
<https://github.com/OpenLineage/OpenLineage/pull/4682>`_.
+
+``OPENLINEAGE_CONTEXT`` is a Databricks job parameter, not an operating-system
environment variable.
+The Databricks task must expose the parameter to the downstream OpenLineage
integration through its
+own configuration channel (for example, ``spark.openlineage.context`` for
OpenLineage Spark or the
+``OPENLINEAGE_CONTEXT`` environment variable for dbt).
+
+The option defaults to the ``openlineage.spark_inject_parent_job_info``
configuration value. Existing
+OpenLineage context job parameters are preserved. Injection is skipped when a
legacy parameter slot
+such as ``notebook_params`` or ``spark_submit_params`` is used because
Databricks does not allow those
+slots to be combined with ``job_parameters``.
+
+.. code-block:: python
+
+ run_now = DatabricksRunNowOperator(
+ task_id="run_now",
+ job_id=123,
+ openlineage_inject_parent_job_info=True,
+ )
+
Durable execution
^^^^^^^^^^^^^^^^^
diff --git a/providers/databricks/docs/operators/submit_run.rst
b/providers/databricks/docs/operators/submit_run.rst
index 0e85b7dd8f6..f5a88056c43 100644
--- a/providers/databricks/docs/operators/submit_run.rst
+++ b/providers/databricks/docs/operators/submit_run.rst
@@ -138,6 +138,22 @@ or ``tasks`` argument.
# i.e. the same dict, copied into the task's dict-shaped parameter slot.
+OpenLineage parent context
+--------------------------
+
+Set ``openlineage_inject_parent_job_info=True`` to inject the standardized
OpenLineage context into
+the dict-shaped parameter slot of each supported task. The context is passed
under the
+``OPENLINEAGE_CONTEXT`` key and includes the Airflow task as a
``BATCH/AIRFLOW/TASK`` parent. The root
+job is marked as ``BATCH/AIRFLOW/DAG`` when it is the current Airflow Dag. For
a root inherited
+through ``DagRun.conf``, its job type is copied from
+``conf["openlineage"]["rootParentJobType"]`` when available and otherwise
omitted.
+
+For tasks with a ``new_cluster``, the operator also retains the existing
injection into
+``new_cluster.spark_conf``. Tasks that only accept positional list parameters
+(``spark_jar_task``, ``spark_python_task``, or ``spark_submit_task``) receive
the context only when
+their Spark configuration can be injected.
+
+
Examples
--------
diff --git
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
index 009c803899a..4ce3d890d2e 100644
---
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
+++
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
@@ -337,6 +337,20 @@ def _inject_airflow_params_into_task(task: dict, params:
dict) -> None:
task_def[field] = dict(params)
+def _inject_openlineage_context_into_task_parameters(task: dict, context:
Context) -> None:
+ """Inject OpenLineage context into each dict-shaped parameter field
supported by a task."""
+ from airflow.providers.databricks.utils.openlineage import (
+ inject_openlineage_context_into_databricks_job_parameters,
+ )
+
+ for task_key, field in _DICT_PARAM_FIELD_BY_TASK.items():
+ task_def = task.get(task_key)
+ if isinstance(task_def, dict):
+ task_def[field] =
inject_openlineage_context_into_databricks_job_parameters(
+ job_parameters=task_def.get(field) or {}, context=context
+ )
+
+
def _coerce_json_to_dict(json: Any) -> dict[str, Any]:
if json is None:
return {}
@@ -731,8 +745,8 @@ class DatabricksSubmitRunOperator(ResumableJobMixin,
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
+ into dict-shaped task parameters and the ``new_cluster``
``spark_conf`` so the Databricks
+ job can emit 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
@@ -947,12 +961,30 @@ class DatabricksSubmitRunOperator(ResumableJobMixin,
BaseOperator):
_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.")
+ self.log.info("Automatic injection of OpenLineage information is
enabled.")
json =
self._inject_openlineage_properties_into_databricks_job(json, context)
return cast("dict[str, Any]", normalise_json_content(json))
def _inject_openlineage_properties_into_databricks_job(self, json: dict,
context: Context) -> dict:
+ if self.openlineage_inject_parent_job_info:
+ try:
+ context_json = copy.deepcopy(json)
+ tasks = context_json.get("tasks")
+ if isinstance(tasks, list):
+ for task in tasks:
+ if isinstance(task, dict):
+
_inject_openlineage_context_into_task_parameters(task, context)
+ else:
+
_inject_openlineage_context_into_task_parameters(context_json, context)
+ json = context_json
+ except Exception as e:
+ self.log.warning(
+ "An error occurred while trying to inject OpenLineage
context. "
+ "Databricks task parameters have not been modified by
OpenLineage.",
+ exc_info=e,
+ )
+
try:
from airflow.providers.databricks.utils.openlineage import (
inject_openlineage_properties_into_databricks_job,
@@ -1227,6 +1259,10 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
:param do_xcom_push: Whether we should push run_id and run_page_url to
xcom.
:param wait_for_termination: if we should wait for termination of the job
run. ``True`` by default.
:param deferrable: Run operator in the deferrable mode.
+ :param openlineage_inject_parent_job_info: If True, injects the
standardized OpenLineage parent-run
+ context into the ``OPENLINEAGE_CONTEXT`` job parameter so Databricks
tasks can link their
+ OpenLineage events back to the Airflow task. Defaults to the
+ ``openlineage.spark_inject_parent_job_info`` config value.
:param repair_run: Repair the databricks run in case of failure.
:param databricks_repair_reason_new_settings: A dict of reason and
new_settings JSON object for which
to repair the run. `None` by default. `None` means to repair at
all cases with existing job
@@ -1300,6 +1336,9 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
do_xcom_push: bool = True,
wait_for_termination: bool = True,
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
+ ),
repair_run: bool = False,
databricks_repair_reason_new_settings: dict[str, Any] | None = None,
cancel_previous_runs: bool = False,
@@ -1332,6 +1371,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
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.repair_run = repair_run
self.databricks_repair_reason_new_settings =
databricks_repair_reason_new_settings or {}
self.cancel_previous_runs = cancel_previous_runs
@@ -1379,13 +1419,13 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
def execute(self, context: Context):
if self.deferrable:
- json = self._prepare_run_now_json()
+ json = self._prepare_run_now_json(context)
self.run_id = self._hook.run_now(json)
_handle_deferrable_databricks_operator_execution(self, self._hook,
self.log, context)
else:
return self.execute_resumable(context)
- def _build_run_now_payload(self) -> dict[str, Any]:
+ def _build_run_now_payload(self, context: Context) -> dict[str, Any]:
# Utility to build the run payload: merge, validate, resolve job_name
-> job_id, inject params.
# Kept separate from cancel_previous_runs so the reconnect path can
rebuild the payload
# (for repair_run) without re-cancelling the run it is reconnecting to.
@@ -1404,6 +1444,9 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
json["job_id"] = job_id
del json["job_name"]
+ return self._inject_run_now_job_parameters(json, context)
+
+ def _inject_run_now_job_parameters(self, json: dict[str, Any], context:
Context) -> dict[str, Any]:
if (
self.forward_dag_params
and not json.get("job_parameters")
@@ -1412,10 +1455,40 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
):
json["job_parameters"] = dict(self.params)
+ if self.openlineage_inject_parent_job_info:
+ if any(k in json for k in
_RUN_NOW_PARAM_SLOTS_CONFLICTING_WITH_JOB_PARAMETERS):
+ self.log.info(
+ "Skipping OpenLineage parent job information injection
because the Databricks "
+ "run uses a legacy parameter slot that cannot be combined
with job_parameters."
+ )
+ else:
+ json =
self._inject_openlineage_properties_into_databricks_job(json, context)
+
return json
- def _prepare_run_now_json(self) -> dict[str, Any]:
- json = self._build_run_now_payload()
+ def _inject_openlineage_properties_into_databricks_job(
+ self, json: dict[str, Any], context: Context
+ ) -> dict[str, Any]:
+ try:
+ from airflow.providers.databricks.utils.openlineage import (
+ inject_openlineage_context_into_databricks_job_parameters,
+ )
+
+ json = dict(json)
+ json["job_parameters"] =
inject_openlineage_context_into_databricks_job_parameters(
+ job_parameters=json.get("job_parameters", {}), context=context
+ )
+ return json
+ except Exception as e:
+ self.log.warning(
+ "An error occurred while trying to inject OpenLineage context.
"
+ "Databricks job parameters have not been modified by
OpenLineage.",
+ exc_info=e,
+ )
+ return json
+
+ def _prepare_run_now_json(self, context: Context) -> dict[str, Any]:
+ json = self._build_run_now_payload(context)
if self.cancel_previous_runs:
if (job_id := json.get("job_id")) is None:
raise ValueError(
@@ -1427,7 +1500,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
return json
def submit_job(self, context: Context) -> int:
- json = self._prepare_run_now_json()
+ json = self._prepare_run_now_json(context)
# Set run_id the instant the run exists so on_kill can cancel it even
if the worker dies
# before polling begins.
self.run_id = self._hook.run_now(json)
@@ -1470,7 +1543,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
# in the poll helper). _build_run_now_payload resolves job_name ->
job_id but omits
# cancel_previous_runs, which would otherwise cancel the run we are
reconnecting to.
if not getattr(self, "_merged_json", None):
- self._merged_json = self._build_run_now_payload()
+ self._merged_json = self._build_run_now_payload(context)
# The run already exists here (fresh submit logged in submit_job, or
reconnect logged by the
# mixin), so the poll helper must not announce a submission.
_handle_databricks_operator_execution(self, self._hook, self.log,
context, announce_submission=False)
@@ -1509,7 +1582,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
# of reading a mutated self.json: on a deferral resume this is
a fresh process, so any
# value written to self.json in execute() is gone.
_get_merged_json() also recovers a
# job_parameters supplied via the named ``job_parameters=``
argument, not only inside json=.
- merged = self._get_merged_json()
+ merged =
self._inject_run_now_job_parameters(self._get_merged_json(), context)
if "job_parameters" in merged:
repair_json["job_parameters"] = merged["job_parameters"]
self._hook.repair_run(repair_json)
diff --git
a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py
b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py
index 706afa5df09..d241db12c94 100644
--- a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py
+++ b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py
@@ -21,6 +21,7 @@ import datetime
import json
import logging
from typing import TYPE_CHECKING, Any
+from uuid import UUID
import requests
@@ -41,6 +42,12 @@ if TYPE_CHECKING:
log = logging.getLogger(__name__)
+_AIRFLOW_DAG_JOB_TYPE = {
+ "processingType": "BATCH",
+ "integration": "AIRFLOW",
+ "jobType": "DAG",
+}
+
def _get_parent_run_facet(task_instance):
"""
@@ -356,6 +363,100 @@ def _is_openlineage_provider_accessible() -> bool:
return True
+def _get_dag_run_conf(task_instance) -> dict[str, Any]:
+ dag_run = getattr(task_instance, "dag_run", None)
+ if dag_run is None:
+ dag_run = task_instance.get_template_context()["dag_run"]
+ conf = getattr(dag_run, "conf", None)
+ return conf if isinstance(conf, dict) else {}
+
+
+def _has_valid_job_identifiers(openlineage_conf: dict[str, Any], keys:
tuple[str, str, str]) -> bool:
+ run_id, namespace, name = (openlineage_conf.get(key) for key in keys)
+ if not all((run_id, namespace, name)):
+ return False
+ try:
+ UUID(str(run_id))
+ except ValueError:
+ return False
+ return True
+
+
+def _get_root_job_type(task_instance) -> dict[str, Any] | None:
+ openlineage_conf = _get_dag_run_conf(task_instance).get("openlineage")
+ if not isinstance(openlineage_conf, dict):
+ return dict(_AIRFLOW_DAG_JOB_TYPE)
+
+ has_inherited_root = _has_valid_job_identifiers(
+ openlineage_conf,
+ ("rootParentRunId", "rootParentJobNamespace", "rootParentJobName"),
+ ) or _has_valid_job_identifiers(
+ openlineage_conf,
+ ("parentRunId", "parentJobNamespace", "parentJobName"),
+ )
+ if not has_inherited_root:
+ return dict(_AIRFLOW_DAG_JOB_TYPE)
+
+ root_job_type = openlineage_conf.get("rootParentJobType")
+ return copy.deepcopy(root_job_type) if isinstance(root_job_type, dict)
else None
+
+
+def _build_openlineage_parent_run_context(task_instance) -> dict[str, Any]:
+ """Build the standardized OpenLineage parent-run context for a task
instance."""
+ parent_run_facet = _get_parent_run_facet(task_instance)
+ root_job: dict[str, Any] = {
+ "namespace": parent_run_facet.root.job.namespace,
+ "name": parent_run_facet.root.job.name,
+ }
+ if root_job_type := _get_root_job_type(task_instance):
+ root_job["facets"] = {"jobType": root_job_type}
+
+ return {
+ "parent": {
+ "run": {"runId": parent_run_facet.run.runId},
+ "job": {
+ "namespace": parent_run_facet.job.namespace,
+ "name": parent_run_facet.job.name,
+ "facets": {
+ "jobType": {
+ "processingType": "BATCH",
+ "integration": "AIRFLOW",
+ "jobType": "TASK",
+ }
+ },
+ },
+ "root": {
+ "run": {"runId": parent_run_facet.root.run.runId},
+ "job": root_job,
+ },
+ }
+ }
+
+
+def inject_openlineage_context_into_databricks_job_parameters(job_parameters:
dict, context: Context) -> dict:
+ """Inject the standardized OpenLineage context into Databricks job
parameters."""
+ if any(key in job_parameters for key in ("OPENLINEAGE_CONTEXT",
"spark.openlineage.context")):
+ log.info(
+ "OpenLineage context is already present in Databricks job
parameters. Skipping the injection."
+ )
+ return job_parameters
+
+ if not _is_openlineage_provider_accessible():
+ log.warning(
+ "Could not access OpenLineage provider for automatic OpenLineage
context injection into "
+ "Databricks job parameters. No action will be performed."
+ )
+ return job_parameters
+
+ log.debug("Injecting OpenLineage context into Databricks job parameters.")
+ return {
+ **job_parameters,
+ "OPENLINEAGE_CONTEXT": json.dumps(
+ _build_openlineage_parent_run_context(context["ti"]),
separators=(",", ":")
+ ),
+ }
+
+
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.
diff --git
a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
index 990e15f1098..d05be38716d 100644
--- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
@@ -1712,6 +1712,217 @@ class
TestDatabricksSubmitRunOperatorOpenLineageInjection:
assert submitted["new_cluster"]["spark_conf"]["spark.executor.memory"]
== "8g"
assert
submitted["new_cluster"]["spark_conf"]["spark.openlineage.parentJobNamespace"]
== "ns"
+ @mock.patch(
+
"airflow.providers.databricks.utils.openlineage.inject_openlineage_context_into_databricks_job_parameters",
+ autospec=True,
+ )
+ @mock.patch(
+
"airflow.providers.databricks.utils.openlineage.inject_openlineage_properties_into_databricks_job",
+ autospec=True,
+ )
+ def test_injects_context_into_task_parameters_in_addition_to_spark_conf(
+ self, mock_spark_inject, mock_context_inject
+ ):
+ mock_context_inject.side_effect = lambda job_parameters, context: {
+ **job_parameters,
+ "OPENLINEAGE_CONTEXT": "context",
+ }
+ mock_spark_inject.side_effect = lambda job, context,
inject_parent_job_info, inject_transport_info: {
+ **job,
+ "new_cluster": {
+ **job["new_cluster"],
+ "spark_conf": {"spark.openlineage.context": "spark-context"},
+ },
+ }
+ context = {"ti": object()}
+ op = DatabricksSubmitRunOperator(
+ task_id=TASK_ID,
+ new_cluster=NEW_CLUSTER,
+ notebook_task={"notebook_path": "/Users/me/notebook"},
+ openlineage_inject_parent_job_info=True,
+ )
+
+ result = op._prepare_submit_json(context)
+
+ assert result["notebook_task"]["base_parameters"] ==
{"OPENLINEAGE_CONTEXT": "context"}
+ assert result["new_cluster"]["spark_conf"] ==
{"spark.openlineage.context": "spark-context"}
+ mock_context_inject.assert_called_once_with(job_parameters={},
context=context)
+ mock_spark_inject.assert_called_once()
+ spark_call = mock_spark_inject.call_args.kwargs
+ assert spark_call["context"] == context
+ assert spark_call["inject_parent_job_info"] is True
+ assert spark_call["inject_transport_info"] is False
+ assert spark_call["job"]["notebook_task"]["base_parameters"] ==
{"OPENLINEAGE_CONTEXT": "context"}
+
+ @mock.patch(
+
"airflow.providers.databricks.utils.openlineage.inject_openlineage_context_into_databricks_job_parameters",
+ autospec=True,
+ )
+ @mock.patch(
+
"airflow.providers.databricks.utils.openlineage.inject_openlineage_properties_into_databricks_job",
+ autospec=True,
+ )
+ def test_injects_context_into_supported_tasks_without_new_cluster(
+ self, mock_spark_inject, mock_context_inject
+ ):
+ mock_context_inject.side_effect = lambda job_parameters, context: {
+ **job_parameters,
+ "OPENLINEAGE_CONTEXT": "context",
+ }
+ mock_spark_inject.side_effect = lambda job, context,
inject_parent_job_info, inject_transport_info: (
+ job
+ )
+ context = {"ti": object()}
+ op = DatabricksSubmitRunOperator(
+ task_id=TASK_ID,
+ tasks=[
+ {"task_key": "notebook", "notebook_task": {"notebook_path":
"/notebook"}},
+ {"task_key": "wheel", "python_wheel_task": {"package_name":
"package"}},
+ {"task_key": "jar", "spark_jar_task": {"main_class_name":
"Main"}},
+ ],
+ openlineage_inject_parent_job_info=True,
+ )
+
+ result = op._prepare_submit_json(context)
+
+ assert result["tasks"][0]["notebook_task"]["base_parameters"] ==
{"OPENLINEAGE_CONTEXT": "context"}
+ assert result["tasks"][1]["python_wheel_task"]["named_parameters"] == {
+ "OPENLINEAGE_CONTEXT": "context"
+ }
+ assert "OPENLINEAGE_CONTEXT" not in
result["tasks"][2]["spark_jar_task"]
+ assert mock_context_inject.call_count == 2
+ mock_spark_inject.assert_called_once()
+
+ @mock.patch(
+
"airflow.providers.databricks.utils.openlineage.inject_openlineage_context_into_databricks_job_parameters",
+ autospec=True,
+ side_effect=RuntimeError("context generation failed"),
+ )
+ @mock.patch(
+
"airflow.providers.databricks.utils.openlineage.inject_openlineage_properties_into_databricks_job",
+ autospec=True,
+ )
+ def test_context_injection_failure_does_not_block_spark_conf_injection(
+ self, mock_spark_inject, mock_context_inject
+ ):
+ mock_spark_inject.side_effect = lambda job, context,
inject_parent_job_info, inject_transport_info: (
+ job
+ )
+ op = DatabricksSubmitRunOperator(
+ task_id=TASK_ID,
+ notebook_task={"notebook_path": "/Users/me/notebook"},
+ new_cluster=NEW_CLUSTER,
+ openlineage_inject_parent_job_info=True,
+ )
+
+ result = op._prepare_submit_json({"ti": object()})
+
+ assert "base_parameters" not in result["notebook_task"]
+ mock_context_inject.assert_called_once()
+ mock_spark_inject.assert_called_once()
+
+
+class TestDatabricksRunNowOperatorOpenLineageInjection:
+ @mock.patch(
+ "airflow.providers.databricks.utils.openlineage."
+ "inject_openlineage_context_into_databricks_job_parameters",
+ autospec=True,
+ )
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook",
autospec=True)
+ def test_inject_openlineage_context_into_job_parameters(self,
db_mock_class, mock_inject):
+ mock_inject.side_effect = lambda job_parameters, context: {
+ **job_parameters,
+ "OPENLINEAGE_CONTEXT":
'{"parent":{"job":{"facets":{"jobType":{"jobType":"TASK"}}}}}',
+ }
+ op = DatabricksRunNowOperator(
+ durable=False,
+ task_id=TASK_ID,
+ job_id=JOB_ID,
+ wait_for_termination=False,
+ openlineage_inject_parent_job_info=True,
+ )
+ db_mock = db_mock_class.return_value
+ db_mock.run_now.return_value = RUN_ID
+ context = {"ti": MagicMock(spec=["stats_tags", "xcom_push"],
stats_tags={})}
+
+ op.execute(context)
+
+ mock_inject.assert_called_once_with(job_parameters={}, context=context)
+ submitted = db_mock.run_now.call_args.args[0]
+ assert submitted["job_parameters"] == {
+ "OPENLINEAGE_CONTEXT":
'{"parent":{"job":{"facets":{"jobType":{"jobType":"TASK"}}}}}'
+ }
+
+ @mock.patch(
+ "airflow.providers.databricks.utils.openlineage."
+ "inject_openlineage_context_into_databricks_job_parameters",
+ autospec=True,
+ side_effect=RuntimeError("context generation failed"),
+ )
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook",
autospec=True)
+ def test_injection_failure_leaves_job_parameters_unchanged(self,
db_mock_class, mock_inject):
+ op = DatabricksRunNowOperator(
+ durable=False,
+ task_id=TASK_ID,
+ job_id=JOB_ID,
+ wait_for_termination=False,
+ openlineage_inject_parent_job_info=True,
+ )
+ db_mock = db_mock_class.return_value
+ db_mock.run_now.return_value = RUN_ID
+
+ op.execute({"ti": MagicMock(spec=["stats_tags", "xcom_push"],
stats_tags={})})
+
+ mock_inject.assert_called_once()
+ submitted = db_mock.run_now.call_args.args[0]
+ assert "job_parameters" not in submitted
+
+ @mock.patch(
+ "airflow.providers.databricks.utils.openlineage."
+ "inject_openlineage_context_into_databricks_job_parameters",
+ autospec=True,
+ )
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook",
autospec=True)
+ def test_does_not_inject_parent_job_info_when_disabled(self,
db_mock_class, mock_inject):
+ op = DatabricksRunNowOperator(
+ durable=False,
+ task_id=TASK_ID,
+ job_id=JOB_ID,
+ wait_for_termination=False,
+ openlineage_inject_parent_job_info=False,
+ )
+ db_mock = db_mock_class.return_value
+ db_mock.run_now.return_value = RUN_ID
+
+ op.execute({"ti": MagicMock(spec=["stats_tags", "xcom_push"],
stats_tags={})})
+
+ mock_inject.assert_not_called()
+
+ @mock.patch(
+ "airflow.providers.databricks.utils.openlineage."
+ "inject_openlineage_context_into_databricks_job_parameters",
+ autospec=True,
+ )
+
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook",
autospec=True)
+ def test_skips_injection_with_legacy_parameter_slots(self, db_mock_class,
mock_inject):
+ op = DatabricksRunNowOperator(
+ durable=False,
+ task_id=TASK_ID,
+ job_id=JOB_ID,
+ notebook_params={"input": "value"},
+ wait_for_termination=False,
+ openlineage_inject_parent_job_info=True,
+ )
+ db_mock = db_mock_class.return_value
+ db_mock.run_now.return_value = RUN_ID
+
+ op.execute({"ti": MagicMock(spec=["stats_tags", "xcom_push"],
stats_tags={})})
+
+ mock_inject.assert_not_called()
+ submitted = db_mock.run_now.call_args.args[0]
+ assert submitted["notebook_params"] == {"input": "value"}
+ assert "job_parameters" not in submitted
+
@pytest.mark.skipif(
not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
diff --git
a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
index cefa881c7ec..7517a160370 100644
--- a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
+++ b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
@@ -18,6 +18,8 @@ from __future__ import annotations
import copy
import datetime
+import json
+from types import SimpleNamespace
from unittest import mock
import pytest
@@ -40,6 +42,7 @@ from airflow.providers.databricks.utils.openlineage import (
_process_data_from_api,
_run_api_call,
emit_openlineage_events_for_databricks_queries,
+ inject_openlineage_context_into_databricks_job_parameters,
inject_openlineage_properties_into_databricks_job,
)
from airflow.providers.openlineage.conf import namespace
@@ -1203,6 +1206,143 @@ def
test_emit_openlineage_events_with_old_openlineage_provider(mock_version):
OL_UTILS = "airflow.providers.databricks.utils.openlineage"
[email protected](
+ ("dag_run_conf", "expected_root_job_type"),
+ [
+ pytest.param(
+ {},
+ {"processingType": "BATCH", "integration": "AIRFLOW", "jobType":
"DAG"},
+ id="current-dag-is-root",
+ ),
+ pytest.param(
+ {
+ "openlineage": {
+ "rootParentRunId": "11111111-1111-1111-1111-111111111111",
+ "rootParentJobNamespace": "external_namespace",
+ "rootParentJobName": "external_job",
+ }
+ },
+ None,
+ id="inherited-root-without-job-type",
+ ),
+ pytest.param(
+ {
+ "openlineage": {
+ "rootParentRunId": "11111111-1111-1111-1111-111111111111",
+ "rootParentJobNamespace": "external_namespace",
+ "rootParentJobName": "external_job",
+ "rootParentJobType": {
+ "processingType": "STREAMING",
+ "integration": "CUSTOM",
+ "jobType": "PIPELINE",
+ },
+ }
+ },
+ {"processingType": "STREAMING", "integration": "CUSTOM",
"jobType": "PIPELINE"},
+ id="inherited-root-with-job-type",
+ ),
+ pytest.param(
+ {
+ "openlineage": {
+ "parentRunId": "22222222-2222-2222-2222-222222222222",
+ "parentJobNamespace": "external_namespace",
+ "parentJobName": "external_job",
+ "rootParentJobType": {
+ "processingType": "BATCH",
+ "integration": "DBT",
+ "jobType": "JOB",
+ },
+ }
+ },
+ {"processingType": "BATCH", "integration": "DBT", "jobType":
"JOB"},
+ id="parent-is-inherited-root-with-job-type",
+ ),
+ pytest.param(
+ {
+ "openlineage": {
+ "rootParentRunId": "invalid-run-id",
+ "rootParentJobNamespace": "external_namespace",
+ "rootParentJobName": "external_job",
+ }
+ },
+ {"processingType": "BATCH", "integration": "AIRFLOW", "jobType":
"DAG"},
+ id="invalid-inherited-root-falls-back-to-current-dag",
+ ),
+ ],
+)
[email protected](f"{OL_UTILS}._get_parent_run_facet", autospec=True)
[email protected](f"{OL_UTILS}._is_openlineage_provider_accessible", autospec=True,
return_value=True)
+def test_inject_openlineage_context_into_job_parameters(
+ mock_accessible, mock_parent, dag_run_conf, expected_root_job_type
+):
+ context = {"ti":
SimpleNamespace(dag_run=SimpleNamespace(conf=dag_run_conf))}
+ parent_run_facet = SimpleNamespace(
+ run=SimpleNamespace(runId="run_id"),
+ job=SimpleNamespace(namespace="namespace", name="dag_id.task_id"),
+ root=SimpleNamespace(
+ run=SimpleNamespace(runId="root_run_id"),
+ job=SimpleNamespace(namespace="namespace", name="dag_id"),
+ ),
+ )
+ mock_parent.return_value = parent_run_facet
+ job_parameters = {"input": "value"}
+
+ result =
inject_openlineage_context_into_databricks_job_parameters(job_parameters,
context)
+
+ assert result["input"] == "value"
+ context_value = json.loads(result["OPENLINEAGE_CONTEXT"])
+ expected_root_job = {
+ "namespace": "namespace",
+ "name": "dag_id",
+ }
+ if expected_root_job_type is not None:
+ expected_root_job["facets"] = {"jobType": expected_root_job_type}
+
+ assert context_value == {
+ "parent": {
+ "run": {"runId": "run_id"},
+ "job": {
+ "namespace": "namespace",
+ "name": "dag_id.task_id",
+ "facets": {
+ "jobType": {"processingType": "BATCH", "integration":
"AIRFLOW", "jobType": "TASK"}
+ },
+ },
+ "root": {
+ "run": {"runId": "root_run_id"},
+ "job": expected_root_job,
+ },
+ }
+ }
+ assert job_parameters == {"input": "value"}
+ mock_parent.assert_called_once_with(context["ti"])
+ mock_accessible.assert_called_once_with()
+
+
[email protected](f"{OL_UTILS}._get_parent_run_facet", autospec=True)
[email protected](f"{OL_UTILS}._is_openlineage_provider_accessible", autospec=True,
return_value=True)
+def
test_inject_openlineage_context_preserves_existing_parameters(mock_accessible,
mock_parent):
+ job_parameters = {"OPENLINEAGE_CONTEXT": "manual_context", "input":
"value"}
+
+ result =
inject_openlineage_context_into_databricks_job_parameters(job_parameters,
{"ti": object()})
+
+ assert result == job_parameters
+ mock_accessible.assert_not_called()
+ mock_parent.assert_not_called()
+
+
[email protected](f"{OL_UTILS}._get_parent_run_facet", autospec=True)
[email protected](f"{OL_UTILS}._is_openlineage_provider_accessible", autospec=True,
return_value=False)
+def test_inject_openlineage_context_provider_inaccessible(mock_accessible,
mock_parent):
+ job_parameters = {"input": "value"}
+
+ result =
inject_openlineage_context_into_databricks_job_parameters(job_parameters,
{"ti": object()})
+
+ assert result == job_parameters
+ mock_accessible.assert_called_once_with()
+ mock_parent.assert_not_called()
+
+
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"}
diff --git a/providers/openlineage/docs/spark.rst
b/providers/openlineage/docs/spark.rst
index a324508b55c..d7c9b12b845 100644
--- a/providers/openlineage/docs/spark.rst
+++ b/providers/openlineage/docs/spark.rst
@@ -100,6 +100,7 @@ Automatic injection is supported for the following
operators:
- :class:`~airflow.providers.amazon.aws.operators.glue.GlueJobOperator`
- :class:`~airflow.providers.apache.livy.operators.livy.LivyOperator`
-
:class:`~airflow.providers.apache.spark.operators.spark_submit.SparkSubmitOperator`
+-
:class:`~airflow.providers.databricks.operators.databricks.DatabricksRunNowOperator`
-
:class:`~airflow.providers.databricks.operators.databricks.DatabricksSubmitRunOperator`
-
:class:`~airflow.providers.google.cloud.operators.dataproc.DataprocCreateBatchOperator`
-
:class:`~airflow.providers.google.cloud.operators.dataproc.DataprocInstantiateInlineWorkflowTemplateOperator`
@@ -125,6 +126,11 @@ Automatic injection is supported for the following
operators:
:class:`~airflow.providers.amazon.aws.operators.emr.EmrServerlessStartJobOperator`
inject parent and transport
properties through their ``spark-defaults`` configuration.
+
:class:`~airflow.providers.databricks.operators.databricks.DatabricksRunNowOperator`
injects the standardized
+ ``OPENLINEAGE_CONTEXT`` through Databricks job parameters. The
+
:class:`~airflow.providers.databricks.operators.databricks.DatabricksSubmitRunOperator`
supports the same
+ context in dict-shaped task parameters and retains its Spark
``spark_conf`` injection for Spark jobs.
+
.. _options:spark_inject_parent_job_info:
