Add AirflowRun on COMPLETE/FAIL events (#40996)

Signed-off-by: Maciej Obuchowski <obuchowski.maciej@gmail.com>
apache · Jul 28, 2024 · 4ab0183 · 4ab0183
1 parent a482d0f
commit 4ab0183
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 11 deletions.
diff --git a/airflow/providers/openlineage/plugins/adapter.py b/airflow/providers/openlineage/plugins/adapter.py
@@ -178,7 +178,7 @@ def start_task(
         nominal_end_time: str | None,
         owners: list[str],
         task: OperatorLineage | None,
-        run_facets: dict[str, RunFacet] | None = None,  # Custom run facets
+        run_facets: dict[str, RunFacet] | None = None,
     ) -> RunEvent:
         """
         Emit openlineage event of type START.
@@ -243,7 +243,7 @@ def complete_task(
         parent_run_id: str | None,
         end_time: str,
         task: OperatorLineage,
-        run_facets: dict[str, RunFacet] | None = None,  # Custom run facets
+        run_facets: dict[str, RunFacet] | None = None,
     ) -> RunEvent:
         """
         Emit openlineage event of type COMPLETE.
@@ -255,7 +255,7 @@ def complete_task(
         :param parent_run_id: identifier of job spawning this task
         :param end_time: time of task completion
         :param task: metadata container with information extracted from operator
-        :param run_facets: custom run facets
+        :param run_facets: additional run facets
         """
         run_facets = run_facets or {}
         if task:
@@ -285,8 +285,8 @@ def fail_task(
         parent_run_id: str | None,
         end_time: str,
         task: OperatorLineage,
-        run_facets: dict[str, RunFacet] | None = None,  # Custom run facets
         error: str | BaseException | None = None,
+        run_facets: dict[str, RunFacet] | None = None,
     ) -> RunEvent:
         """
         Emit openlineage event of type FAIL.
@@ -300,6 +300,7 @@ def fail_task(
         :param task: metadata container with information extracted from operator
         :param run_facets: custom run facets
         :param error: error
+        :param run_facets: additional run facets
         """
         run_facets = run_facets or {}
         if task:

diff --git a/airflow/providers/openlineage/plugins/listener.py b/airflow/providers/openlineage/plugins/listener.py
@@ -235,7 +235,10 @@ def on_success():
                 parent_run_id=parent_run_id,
                 end_time=end_date.isoformat(),
                 task=task_metadata,
-                run_facets=get_user_provided_run_facets(task_instance, TaskInstanceState.SUCCESS),
+                run_facets={
+                    **get_user_provided_run_facets(task_instance, TaskInstanceState.SUCCESS),
+                    **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid),
+                },
             )
             Stats.gauge(
                 f"ol.event.size.{event_type}.{operator_name}",
@@ -330,8 +333,11 @@ def on_failure():
                 parent_run_id=parent_run_id,
                 end_time=end_date.isoformat(),
                 task=task_metadata,
-                run_facets=get_user_provided_run_facets(task_instance, TaskInstanceState.FAILED),
                 error=error,
+                run_facets={
+                    **get_user_provided_run_facets(task_instance, TaskInstanceState.FAILED),
+                    **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid),
+                },
             )
             Stats.gauge(
                 f"ol.event.size.{event_type}.{operator_name}",

diff --git a/tests/providers/openlineage/plugins/test_adapter.py b/tests/providers/openlineage/plugins/test_adapter.py
@@ -45,7 +45,10 @@
 from airflow.providers.openlineage.conf import namespace
 from airflow.providers.openlineage.extractors import OperatorLineage
 from airflow.providers.openlineage.plugins.adapter import _PRODUCER, OpenLineageAdapter
-from airflow.providers.openlineage.plugins.facets import AirflowDagRunFacet, AirflowStateRunFacet
+from airflow.providers.openlineage.plugins.facets import (
+    AirflowDagRunFacet,
+    AirflowStateRunFacet,
+)
 from airflow.providers.openlineage.utils.utils import get_airflow_job_facet
 from airflow.utils.task_group import TaskGroup
 from tests.test_utils.config import conf_vars

diff --git a/tests/providers/openlineage/plugins/test_listener.py b/tests/providers/openlineage/plugins/test_listener.py
@@ -264,11 +264,16 @@ def test_adapter_start_task_is_called_with_proper_arguments(
 
 @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled")
 @mock.patch("airflow.providers.openlineage.plugins.listener.OpenLineageAdapter")
+@mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet")
 @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets")
 @mock.patch("airflow.providers.openlineage.plugins.listener.get_job_name")
 @mock.patch("airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call)
 def test_adapter_fail_task_is_called_with_proper_arguments(
-    mock_get_job_name, mock_get_user_provided_run_facets, mocked_adapter, mock_disabled
+    mock_get_job_name,
+    mock_get_user_provided_run_facets,
+    mock_get_airflow_run_facet,
+    mocked_adapter,
+    mock_disabled,
 ):
     """Tests that the 'fail_task' method of the OpenLineageAdapter is invoked with the correct arguments.
 
@@ -289,6 +294,7 @@ def mock_task_id(dag_id, task_id, try_number, execution_date):
     mocked_adapter.build_dag_run_id.side_effect = mock_dag_id
     mocked_adapter.build_task_instance_run_id.side_effect = mock_task_id
     mock_get_user_provided_run_facets.return_value = {"custom_user_facet": 2}
+    mock_get_airflow_run_facet.return_value = {"airflow": {"task": "..."}}
     mock_disabled.return_value = False
 
     err = ValueError("test")
@@ -305,18 +311,23 @@ def mock_task_id(dag_id, task_id, try_number, execution_date):
         parent_run_id="execution_date.dag_id",
         run_id="execution_date.dag_id.task_id.1",
         task=listener.extractor_manager.extract_metadata(),
-        run_facets={"custom_user_facet": 2},
+        run_facets={"custom_user_facet": 2, "airflow": {"task": "..."}},
         **expected_err_kwargs,
     )
 
 
 @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled")
 @mock.patch("airflow.providers.openlineage.plugins.listener.OpenLineageAdapter")
+@mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet")
 @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets")
 @mock.patch("airflow.providers.openlineage.plugins.listener.get_job_name")
 @mock.patch("airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call)
 def test_adapter_complete_task_is_called_with_proper_arguments(
-    mock_get_job_name, mock_get_user_provided_run_facets, mocked_adapter, mock_disabled
+    mock_get_job_name,
+    mock_get_user_provided_run_facets,
+    mock_get_airflow_run_facet,
+    mocked_adapter,
+    mock_disabled,
 ):
     """Tests that the 'complete_task' method of the OpenLineageAdapter is called with the correct arguments.
 
@@ -338,6 +349,7 @@ def mock_task_id(dag_id, task_id, try_number, execution_date):
     mocked_adapter.build_dag_run_id.side_effect = mock_dag_id
     mocked_adapter.build_task_instance_run_id.side_effect = mock_task_id
     mock_get_user_provided_run_facets.return_value = {"custom_user_facet": 2}
+    mock_get_airflow_run_facet.return_value = {"airflow": {"task": "..."}}
     mock_disabled.return_value = False
 
     listener.on_task_instance_success(None, task_instance, None)
@@ -352,7 +364,7 @@ def mock_task_id(dag_id, task_id, try_number, execution_date):
         parent_run_id="execution_date.dag_id",
         run_id=f"execution_date.dag_id.task_id.{EXPECTED_TRY_NUMBER_1}",
         task=listener.extractor_manager.extract_metadata(),
-        run_facets={"custom_user_facet": 2},
+        run_facets={"custom_user_facet": 2, "airflow": {"task": "..."}},
     )