[DOP-22750] Add owner parsing for airflow run events

Author: TiGrib

data_rentgen/consumer/extractors/run.py

@@ -182,6 +182,15 @@ def enrich_run_user(run: RunDTO, event: OpenLineageRunEvent) -> RunDTO:
     spark_application_details = event.run.facets.spark_applicationDetails
     if spark_application_details:
         run.user = UserDTO(name=spark_application_details.userName)
+
     # Airflow DAG and task have 'owner' field, but if can be either user or group name,
     # and also it does not mean that this exact user started this run.
+    # Airflow using different facets for version above provider-opelineage/1.11.0.
+    airflow_application_details = event.run.facets.airflow
+    if airflow_application_details and airflow_application_details.dag.owner != "airflow":
+        run.user = UserDTO(name=airflow_application_details.dag.owner)
+    airflow_application_dag_details = event.run.facets.airflowDagRun
+    if airflow_application_dag_details and airflow_application_dag_details.dag.owner != "airflow":
+        run.user = UserDTO(name=airflow_application_dag_details.dag.owner)
+
     return run

data_rentgen/consumer/openlineage/run_facets/airflow.py

@@ -14,6 +14,7 @@ class OpenLineageAirflowDagInfo(OpenLineageBase):
     """
 
     dag_id: str
+    owner: str
 
 
 class OpenLineageAirflowDagRunType(Enum):

tests/test_consumer/test_extractors/test_extractors_batch_airflow.py

@@ -66,7 +66,7 @@ def airflow_dag_run_event_start() -> OpenLineageRunEvent:
                     openlineageAdapterVersion=Version("1.10.0"),
                 ),
                 airflowDagRun=OpenLineageAirflowDagRunFacet(
-                    dag=OpenLineageAirflowDagInfo(dag_id="mydag"),
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="manual__2024-07-05T09:04:13:979349+00:00",
                         run_type=OpenLineageAirflowDagRunType.MANUAL,
@@ -137,7 +137,7 @@ def airflow_task_run_event_start() -> OpenLineageRunEvent:
                     openlineageAdapterVersion=Version("1.10.0"),
                 ),
                 airflow=OpenLineageAirflowTaskRunFacet(
-                    dag=OpenLineageAirflowDagInfo(dag_id="mydag"),
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="manual__2024-07-05T09:04:13:979349+00:00",
                         run_type=OpenLineageAirflowDagRunType.MANUAL,

tests/test_consumer/test_extractors/test_extractors_run.py

@@ -174,7 +174,7 @@ def test_extractors_extract_run_airflow_dag_2_3_plus():
                     openlineageAdapterVersion=Version("1.10.0"),
                 ),
                 airflowDagRun=OpenLineageAirflowDagRunFacet(
-                    dag=OpenLineageAirflowDagInfo(dag_id="mydag"),
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="manual__2024-07-05T09:04:13:979349+00:00",
                         run_type=OpenLineageAirflowDagRunType.MANUAL,
@@ -237,7 +237,7 @@ def test_extractors_extract_run_airflow_dag_2_x():
                     openlineageAdapterVersion=Version("1.10.0"),
                 ),
                 airflowDagRun=OpenLineageAirflowDagRunFacet(
-                    dag=OpenLineageAirflowDagInfo(dag_id="mydag"),
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="manual__2024-07-05T09:04:13:979349+00:00",
                         run_type=OpenLineageAirflowDagRunType.MANUAL,
@@ -300,7 +300,7 @@ def test_extractors_extract_run_airflow_task_with_ti_persistent_log_url():
                     openlineageAdapterVersion=Version("1.10.0"),
                 ),
                 airflow=OpenLineageAirflowTaskRunFacet(
-                    dag=OpenLineageAirflowDagInfo(dag_id="mydag"),
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="manual__2024-07-05T09:04:13:979349+00:00",
                         run_type=OpenLineageAirflowDagRunType.MANUAL,
@@ -372,7 +372,7 @@ def test_extractors_extract_run_airflow_task_2_9_plus():
                     openlineageAdapterVersion=Version("1.9.0"),
                 ),
                 airflow=OpenLineageAirflowTaskRunFacet(
-                    dag=OpenLineageAirflowDagInfo(dag_id="mydag"),
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="backfill__2024-07-05T09:04:13:979349+00:00",
                         run_type=OpenLineageAirflowDagRunType.BACKFILL_JOB,
@@ -436,7 +436,7 @@ def test_extractors_extract_run_airflow_task_2_x():
             runId=run_id,
             facets=OpenLineageRunFacets(
                 airflow=OpenLineageAirflowTaskRunFacet(
-                    dag=OpenLineageAirflowDagInfo(dag_id="mydag"),
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="scheduled__2024-07-05T09:04:13:979349+00:00",
                         run_type=OpenLineageAirflowDagRunType.SCHEDULED,

tests/test_consumer/test_handlers/test_runs_handler_airflow.py

@@ -17,6 +17,7 @@
     Run,
     RunStartReason,
     RunStatus,
+    User,
 )
 
 RESOURCES_PATH = Path(__file__).parent.parent.joinpath("resources").resolve()
@@ -97,7 +98,10 @@ async def test_runs_handler_airflow(
     assert task_run.status == RunStatus.SUCCEEDED
     assert task_run.started_at == datetime(2024, 7, 5, 9, 4, 20, 783845, tzinfo=timezone.utc)
     assert task_run.ended_at == datetime(2024, 7, 5, 9, 7, 37, 858423, tzinfo=timezone.utc)
-    assert task_run.started_by_user_id is None
+    user_query = select(User).where(User.name == "myuser")
+    user_scalars = await async_session.scalars(user_query)
+    user = user_scalars.one_or_none()
+    assert task_run.started_by_user_id == user.id
     assert task_run.start_reason == RunStartReason.MANUAL
     assert task_run.end_reason is None
     assert task_run.external_id == "manual__2024-07-05T09:04:12.162809+00:00"

tests/test_consumer/test_openlineage/test_run_event_airflow.py

@@ -150,6 +150,7 @@ def test_run_event_airflow_dag_start():
                 airflowDagRun=OpenLineageAirflowDagRunFacet(
                     dag=OpenLineageAirflowDagInfo(
                         dag_id="mydag",
+                        owner="myuser",
                     ),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="manual__2024-07-05T09:04:12.162809+00:00",
@@ -400,6 +401,7 @@ def test_run_event_airflow_task_start():
                 airflow=OpenLineageAirflowTaskRunFacet(
                     dag=OpenLineageAirflowDagInfo(
                         dag_id="mydag",
+                        owner="myuser",
                     ),
                     dagRun=OpenLineageAirflowDagRunInfo(
                         run_id="manual__2024-07-05T09:04:12.162809+00:00",