feat(openlineage): Add parentRunFacet for DAG events (#57809)

Author: kacpermuda

providers/openlineage/docs/guides/user.rst

@@ -478,6 +478,56 @@ You can enable this automation by setting ``spark_inject_transport_info`` option
   AIRFLOW__OPENLINEAGE__SPARK_INJECT_TRANSPORT_INFO=true
 
 
+Passing parent information to Airflow DAG
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To enable full OpenLineage lineage tracking across dependent DAGs, you can pass parent and root job information
+through the DAG's ``dag_run.conf``. When a DAG run configuration includes an ``openlineage`` section with valid metadata,
+this information is automatically parsed and converted into DAG run's ``parentRunFacet``, from which the root information
+is also propagated to all tasks. If no DAG run ``openlineage`` configuration is provided, the DAG run will not contain
+``parentRunFacet`` and root of all tasks will default to Dag run.
+
+The ``openlineage`` dict in conf should contain the following keys:
+
+
+*(all three values must be included to create a parent reference)*
+
+- **parentRunId** — the unique run ID (uuid) of the direct parent job
+- **parentJobName** — the name of the parent job
+- **parentJobNamespace** — the namespace of the parent job
+
+*(all three values must be included to create a root reference, otherwise parent will be used as root)*
+
+- **rootParentRunId** — the run ID (uuid) of the top-level (root) job
+- **rootParentJobName** — the name of the top-level (root) job
+- **rootParentJobNamespace** — the namespace of the top-level (root) job
+
+.. note::
+
+  We highly recommend providing all six OpenLineage identifiers (parent and root) to ensure complete lineage tracking. If the root information is missing, the parent set will be used as the root; if any of the three parent fields are missing, no parent facet will be created. Partial or mixed configurations are not supported - either all three parent or all three root values must be provided together.
+
+
+Example:
+
+.. code-block:: shell
+
+    curl -X POST "http://<AIRFLOW_HOST>/api/v2/dags/my_dag_name/dagRuns" \
+    -H "Content-Type: application/json" \
+    -d '{
+      "logical_date": "2019-08-24T14:15:22Z",
+      "conf": {
+        "openlineage": {
+          "parentRunId": "3bb703d1-09c1-4a42-8da5-35a0b3216072",
+          "parentJobNamespace": "prod_biz",
+          "parentJobName": "get_files",
+          "rootParentRunId": "9d3b14f7-de91-40b6-aeef-e887e2c7673e",
+          "rootParentJobNamespace": "prod_analytics",
+          "rootParentJobName": "generate_report_sales_e2e"
+        }
+      }
+    }'
+
+
 Troubleshooting
 ===============
 

providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py

@@ -41,7 +41,9 @@
     get_airflow_mapped_task_facet,
     get_airflow_run_facet,
     get_dag_documentation,
+    get_dag_parent_run_facet,
     get_job_name,
+    get_root_information_from_dagrun_conf,
     get_task_documentation,
     get_task_parent_run_facet,
     get_user_provided_run_facets,
@@ -224,7 +226,11 @@ def on_running():
                 task=task_metadata,
                 run_facets={
                     **get_user_provided_run_facets(task_instance, TaskInstanceState.RUNNING),
-                    **get_task_parent_run_facet(parent_run_id=parent_run_id, parent_job_name=dag.dag_id),
+                    **get_task_parent_run_facet(
+                        parent_run_id=parent_run_id,
+                        parent_job_name=dag.dag_id,
+                        **get_root_information_from_dagrun_conf(getattr(dagrun, "conf", {})),
+                    ),
                     **get_airflow_mapped_task_facet(task_instance),
                     **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid),
                     **debug_facet,
@@ -351,7 +357,11 @@ def on_success():
                 nominal_end_time=data_interval_end,
                 run_facets={
                     **get_user_provided_run_facets(task_instance, TaskInstanceState.SUCCESS),
-                    **get_task_parent_run_facet(parent_run_id=parent_run_id, parent_job_name=dag.dag_id),
+                    **get_task_parent_run_facet(
+                        parent_run_id=parent_run_id,
+                        parent_job_name=dag.dag_id,
+                        **get_root_information_from_dagrun_conf(getattr(dagrun, "conf", {})),
+                    ),
                     **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid),
                     **get_airflow_debug_facet(),
                 },
@@ -489,7 +499,11 @@ def on_failure():
                 job_description_type=doc_type,
                 run_facets={
                     **get_user_provided_run_facets(task_instance, TaskInstanceState.FAILED),
-                    **get_task_parent_run_facet(parent_run_id=parent_run_id, parent_job_name=dag.dag_id),
+                    **get_task_parent_run_facet(
+                        parent_run_id=parent_run_id,
+                        parent_job_name=dag.dag_id,
+                        **get_root_information_from_dagrun_conf(getattr(dagrun, "conf", {})),
+                    ),
                     **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid),
                     **get_airflow_debug_facet(),
                 },
@@ -540,7 +554,11 @@ def on_state_change():
                 "job_description": None,
                 "job_description_type": None,
                 "run_facets": {
-                    **get_task_parent_run_facet(parent_run_id=parent_run_id, parent_job_name=ti.dag_id),
+                    **get_task_parent_run_facet(
+                        parent_run_id=parent_run_id,
+                        parent_job_name=ti.dag_id,
+                        **get_root_information_from_dagrun_conf(getattr(dagrun, "conf", {})),
+                    ),
                     **get_airflow_debug_facet(),
                 },
             }
@@ -645,8 +663,6 @@ def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None:
             )
             data_interval_end = dag_run.data_interval_end.isoformat() if dag_run.data_interval_end else None
 
-            run_facets = {**get_airflow_dag_run_facet(dag_run)}
-
             date = dag_run.logical_date
             if AIRFLOW_V_3_0_PLUS and date is None:
                 date = dag_run.run_after
@@ -660,7 +676,6 @@ def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None:
                 start_date=dag_run.start_date,
                 nominal_start_time=data_interval_start,
                 nominal_end_time=data_interval_end,
-                run_facets=run_facets,
                 clear_number=dag_run.clear_number,
                 owners=[x.strip() for x in dag_run.dag.owner.split(",")] if dag_run.dag else None,
                 job_description=doc,
@@ -669,6 +684,10 @@ def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None:
                 # AirflowJobFacet should be created outside ProcessPoolExecutor that pickles objects,
                 # as it causes lack of some TaskGroup attributes and crashes event emission.
                 job_facets=get_airflow_job_facet(dag_run=dag_run),
+                run_facets={
+                    **get_airflow_dag_run_facet(dag_run),
+                    **get_dag_parent_run_facet(getattr(dag_run, "conf", {})),
+                },
             )
         except BaseException as e:
             self.log.warning("OpenLineage received exception in method on_dag_run_running", exc_info=e)
@@ -716,7 +735,10 @@ def on_dag_run_success(self, dag_run: DagRun, msg: str) -> None:
                 job_description_type=doc_type,
                 task_ids=task_ids,
                 dag_run_state=dag_run.get_state(),
-                run_facets={**get_airflow_dag_run_facet(dag_run)},
+                run_facets={
+                    **get_airflow_dag_run_facet(dag_run),
+                    **get_dag_parent_run_facet(getattr(dag_run, "conf", {})),
+                },
             )
         except BaseException as e:
             self.log.warning("OpenLineage received exception in method on_dag_run_success", exc_info=e)
@@ -765,7 +787,10 @@ def on_dag_run_failed(self, dag_run: DagRun, msg: str) -> None:
                 dag_run_state=dag_run.get_state(),
                 task_ids=task_ids,
                 msg=msg,
-                run_facets={**get_airflow_dag_run_facet(dag_run)},
+                run_facets={
+                    **get_airflow_dag_run_facet(dag_run),
+                    **get_dag_parent_run_facet(getattr(dag_run, "conf", {})),
+                },
             )
         except BaseException as e:
             self.log.warning("OpenLineage received exception in method on_dag_run_failed", exc_info=e)

providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py

@@ -20,7 +20,7 @@
 
 from airflow.providers.openlineage import conf
 from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter
-from airflow.providers.openlineage.utils.utils import get_job_name
+from airflow.providers.openlineage.utils.utils import get_job_name, get_root_information_from_dagrun_conf
 from airflow.providers.openlineage.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
@@ -102,51 +102,69 @@ def lineage_root_parent_id(task_instance: TaskInstance):
     """
     return "/".join(
         (
-            lineage_job_namespace(),
+            lineage_root_job_namespace(task_instance),
             lineage_root_job_name(task_instance),
             lineage_root_run_id(task_instance),
         )
     )
 
 
 def lineage_root_job_name(task_instance: TaskInstance):
+    root_parent_job_name = _get_ol_root_id("root_parent_job_name", task_instance)
+    if root_parent_job_name:
+        return root_parent_job_name
     return task_instance.dag_id
 
 
 def lineage_root_run_id(task_instance: TaskInstance):
+    root_parent_run_id = _get_ol_root_id("root_parent_run_id", task_instance)
+    if root_parent_run_id:
+        return root_parent_run_id
     return OpenLineageAdapter.build_dag_run_id(
         dag_id=task_instance.dag_id,
         logical_date=_get_logical_date(task_instance),
         clear_number=_get_dag_run_clear_number(task_instance),
     )
 
 
+def lineage_root_job_namespace(task_instance: TaskInstance):
+    root_parent_job_namespace = _get_ol_root_id("root_parent_job_namespace", task_instance)
+    if root_parent_job_namespace:
+        return root_parent_job_namespace
+    return conf.namespace()
+
+
+def _get_ol_root_id(id_key: str, task_instance: TaskInstance) -> str | None:
+    dr_conf = _get_dag_run_conf(task_instance=task_instance)
+    ol_root_info = get_root_information_from_dagrun_conf(dr_conf=dr_conf)
+    if ol_root_info and ol_root_info.get(id_key):
+        return ol_root_info[id_key]
+    return None
+
+
+def _get_dagrun_from_ti(task_instance: TaskInstance):
+    context = task_instance.get_template_context()
+    if getattr(task_instance, "dag_run", None):
+        return task_instance.dag_run
+    return context["dag_run"]
+
+
+def _get_dag_run_conf(task_instance: TaskInstance) -> dict:
+    dr = _get_dagrun_from_ti(task_instance=task_instance)
+    return dr.conf or {}
+
+
 def _get_dag_run_clear_number(task_instance: TaskInstance):
-    # todo: remove when min airflow version >= 3.0
-    if AIRFLOW_V_3_0_PLUS:
-        context = task_instance.get_template_context()
-        if hasattr(task_instance, "dag_run"):
-            dag_run = task_instance.dag_run
-        else:
-            dag_run = context["dag_run"]
-        return dag_run.clear_number
-    return task_instance.dag_run.clear_number
+    dr = _get_dagrun_from_ti(task_instance=task_instance)
+    return dr.clear_number
 
 
 def _get_logical_date(task_instance):
-    # todo: remove when min airflow version >= 3.0
     if AIRFLOW_V_3_0_PLUS:
-        context = task_instance.get_template_context()
-        if hasattr(task_instance, "dag_run"):
-            dag_run = task_instance.dag_run
-        else:
-            dag_run = context["dag_run"]
-        if hasattr(dag_run, "logical_date") and dag_run.logical_date:
-            date = dag_run.logical_date
-        else:
-            date = dag_run.run_after
-    elif hasattr(task_instance, "logical_date"):
-        date = task_instance.logical_date
-    else:
-        date = task_instance.execution_date
-    return date
+        dr = _get_dagrun_from_ti(task_instance=task_instance)
+        if getattr(dr, "logical_date", None):
+            return dr.logical_date
+        return dr.run_after
+    if getattr(task_instance, "logical_date", None):
+        return task_instance.logical_date
+    return task_instance.execution_date

providers/openlineage/src/airflow/providers/openlineage/plugins/openlineage.py

@@ -28,6 +28,7 @@
         lineage_job_namespace,
         lineage_parent_id,
         lineage_root_job_name,
+        lineage_root_job_namespace,
         lineage_root_parent_id,
         lineage_root_run_id,
         lineage_run_id,
@@ -51,6 +52,7 @@ class OpenLineageProviderPlugin(AirflowPlugin):
             lineage_parent_id,
             lineage_root_run_id,
             lineage_root_job_name,
+            lineage_root_job_namespace,
             lineage_root_parent_id,
         ]
         listeners = [get_openlineage_listener()]