Improve task loop information and tests for OpenLineage provider

airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests/prepare_dags.py

@@ -50,6 +50,7 @@
 # otherwise raise at parse time and take the whole dag-processor down, not just skip that one DAG.
 MIN_AIRFLOW_VERSION_FOR_DAG: dict[str, tuple[int, int]] = {
     "example_openlineage_hitl_dag.py": (3, 1),
+    "example_openlineage_task_loop_dag.py": (3, 4),
 }
 
 

providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py

@@ -162,8 +162,8 @@ def test_get_parent_run_metadata():
 class TestGenerateOpenLineageEventsFromDbtCloudRun:
     @patch("importlib.metadata.version", return_value="3.0.0")
     @patch("airflow.providers.openlineage.api.emit")
-    @patch("airflow.providers.openlineage.plugins.adapter.OpenLineageAdapter.build_task_instance_run_id")
-    @patch("airflow.providers.openlineage.plugins.adapter.OpenLineageAdapter.build_dag_run_id")
+    @patch("airflow.providers.openlineage.plugins.macros.lineage_run_id")
+    @patch("airflow.providers.openlineage.plugins.macros.lineage_root_run_id")
     @patch.object(DbtCloudHook, "get_job_run")
     @patch.object(DbtCloudHook, "get_project")
     @patch.object(DbtCloudHook, "get_job_run_artifact")

providers/google/tests/unit/google/cloud/operators/test_dataproc.py

@@ -1898,8 +1898,8 @@ def test_dataproc_operator_execute_async_done_before_defer(self, mock_submit_job
         op.execute(context=self.mock_context)
         assert not mock_defer.called
 
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("DataprocHook"))
     def test_execute_openlineage_parent_job_info_injection(
@@ -1969,8 +1969,8 @@ def test_execute_openlineage_parent_job_info_injection(
             metadata=METADATA,
         )
 
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("DataprocHook"))
@@ -2028,8 +2028,8 @@ def test_execute_openlineage_http_transport_info_injection(
             metadata=METADATA,
         )
 
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("DataprocHook"))
@@ -3120,8 +3120,8 @@ def test_wait_for_operation_on_execute(self, mock_hook):
         )
         mock_op.return_value.result.assert_not_called()
 
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("DataprocHook"))
     def test_execute_openlineage_parent_job_info_injection(
@@ -3315,8 +3315,8 @@ def test_execute_openlineage_parent_job_info_injection_skipped_when_ol_not_acces
             metadata=METADATA,
         )
 
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("DataprocHook"))
@@ -3425,8 +3425,8 @@ def test_execute_openlineage_transport_info_injection(
             metadata=METADATA,
         )
 
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("DataprocHook"))
@@ -4124,8 +4124,8 @@ def test_execute_batch_already_exists_cancelled(self, mock_hook, mock_log):
         mock_log.info.assert_any_call("Batch with given id already exists.")
 
     @mock.patch.object(DataprocCreateBatchOperator, "log", new_callable=mock.MagicMock)
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
     @mock.patch(DATAPROC_PATH.format("DataprocHook"))
@@ -4182,8 +4182,8 @@ def test_execute_openlineage_parent_job_info_injection(
         mock_log.info.assert_any_call("Batch job %s completed.\nDriver logs: %s", BATCH_ID, logs_link)
 
     @mock.patch.object(DataprocCreateBatchOperator, "log", new_callable=mock.MagicMock)
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
@@ -4250,8 +4250,8 @@ def test_execute_openlineage_transport_info_injection(
             logs_link,
         )
 
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id")
-    @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id")
+    @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id")
     @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener")
     @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible")
     @mock.patch(DATAPROC_PATH.format("Batch.to_dict"))

providers/openlineage/src/airflow/providers/openlineage/facets/AirflowJobFacet.json

@@ -15,9 +15,30 @@
                 "additionalProperties": true
               },
               "taskGroups": {
-                "description": "Information about all task groups within the DAG.",
+                "description": "Information about all task groups within the DAG, keyed by group id.",
                 "type": "object",
-                "additionalProperties": true
+                "additionalProperties": {
+                  "type": "object",
+                  "properties": {
+                    "max_iterations": {
+                      "description": "Maximum number of iterations of the loop (Airflow 3.4+). Present for a loop group only.",
+                      "type": "integer"
+                    },
+                    "has_until": {
+                      "description": "Whether the loop has an until condition that can stop it before max_iterations (Airflow 3.4+). Present for a loop group only.",
+                      "type": "boolean"
+                    },
+                    "gate_task_id": {
+                      "description": "Id of the loop's gate, the task that ends each iteration and decides whether the loop goes on (Airflow 3.4+). Present for a loop group only.",
+                      "type": "string"
+                    },
+                    "terminal_task_id": {
+                      "description": "Id of the last task of the loop body, whose result the gate and the next iteration read (Airflow 3.4+). Present for a loop group only.",
+                      "type": "string"
+                    }
+                  },
+                  "additionalProperties": true
+                }
               },
               "tasks": {
                 "description": "Details of all individual tasks within the DAG.",

providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json

@@ -142,6 +142,22 @@
             },
             "upstream_task_ids": {
               "type": "string"
+            },
+            "max_iterations": {
+              "description": "Maximum number of iterations of the loop (Airflow 3.4+). Present for a loop group only.",
+              "type": "integer"
+            },
+            "has_until": {
+              "description": "Whether the loop has an until condition that can stop it before max_iterations (Airflow 3.4+). Present for a loop group only.",
+              "type": "boolean"
+            },
+            "gate_task_id": {
+              "description": "Id of the loop's gate, the task that ends each iteration and decides whether the loop goes on (Airflow 3.4+). Present for a loop group only.",
+              "type": "string"
+            },
+            "terminal_task_id": {
+              "description": "Id of the last task of the loop body, whose result the gate and the next iteration read (Airflow 3.4+). Present for a loop group only.",
+              "type": "string"
             }
           },
           "additionalProperties": true,
@@ -193,13 +209,59 @@
         "duration": {
           "type": "number"
         },
+        "id": {
+          "description": "UUID of this attempt (Airflow 3+). Every try, and every task instance a clear creates again, gets a new one.",
+          "type": "string",
+          "format": "uuid"
+        },
         "map_index": {
+          "description": "Position in a mapped task's expansion. Absent for a task that isn't mapped.",
+          "type": "integer"
+        },
+        "loop_index": {
+          "description": "Zero-based iteration of the innermost loop the task runs in (Airflow 3.4+). Absent outside a loop.",
+          "type": "integer"
+        },
+        "loop_max_iterations": {
+          "description": "Maximum number of iterations of the innermost loop the task runs in (Airflow 3.4+). Absent outside a loop.",
+          "type": "integer"
+        },
+        "loop_group_id": {
+          "description": "Group id of the innermost loop the task runs in (Airflow 3.4+), which loop_index counts the iterations of. Absent outside a loop.",
+          "type": "string"
+        },
+        "loop_decision": {
+          "description": "Decision of a loop gate (Airflow 3.4+), on its COMPLETE or FAIL event only. A FAIL with stop false at the last iteration means the loop's until was still false at max_iterations.",
+          "type": "object",
+          "properties": {
+            "stop": {
+              "description": "Whether the gate stopped the loop.",
+              "type": "boolean"
+            },
+            "reason": {
+              "description": "Reason given by until, if any.",
+              "type": "string"
+            }
+          }
+        },
+        "region_id": {
+          "description": "Id of the region (one execution of a loop or of a mapped task's expansion) the task instance belongs to (Airflow 3.4+). Shared by the task instances of that region. Clearing a loop forks its region, so a re-run of an iteration gets a new region_id. Region ids are UUIDv7, so a later trajectory has a greater one, which is an implementation detail. Absent for a task that is neither mapped nor in a loop.",
+          "type": "string",
+          "format": "uuid"
+        },
+        "region_index": {
+          "description": "Position inside the region (Airflow 3.4+): the iteration for a task in a loop that isn't mapped, the map index for a mapped task. Together with dag_id, the Dag run's run_id, task_id, region_id and try_number it identifies one attempt. Absent with region_id.",
           "type": "integer"
         },
         "pool": {
           "type": "string"
         },
         "try_number": {
+          "description": "Try of the task instance, counted per region_id and region_index.",
+          "type": "integer"
+        },
+        "max_tries": {
+          "description": "Last try_number that is retried on failure, so a failed try with try_number > max_tries is final. Equals the task's retries until the task instance is cleared, which sets it to the last try_number plus retries.",
           "type": "integer"
         },
         "queued_dttm": {

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

@@ -53,8 +53,9 @@
     get_dag_run_dag_and_task_from_ti,
     get_job_name,
     get_orm_task_instance_positions,
-    get_regional_task_instance_run_id,
     get_task_documentation,
+    get_task_instance_position,
+    get_task_instance_run_id,
     get_task_parent_run_facet,
     get_user_provided_run_facets,
     is_dag_run_asset_triggered,
@@ -263,15 +264,8 @@ def on_running():
                 clear_number=clear_number,
             )
 
-            task_uuid = get_regional_task_instance_run_id(
-                task_instance
-            ) or self.adapter.build_task_instance_run_id(
-                dag_id=task_instance.dag_id,
-                task_id=task_instance.task_id,
-                try_number=task_instance.try_number,
-                logical_date=date,
-                map_index=task_instance.map_index,
-            )
+            position = get_task_instance_position(task_instance)
+            task_uuid = get_task_instance_run_id(task_instance, logical_date=date, position=position)
             event_type = RunState.RUNNING.value.lower()
             operator_name = task.task_type.lower()
 
@@ -341,6 +335,7 @@ def on_running():
                         task,
                         task_uuid,
                         include_full_task_info=controls.include_full_task_info,
+                        ti_position=position,
                     ),
                     **debug_facet,
                 },
@@ -427,15 +422,8 @@ def on_success():
                 clear_number=dagrun.clear_number,
             )
 
-            task_uuid = get_regional_task_instance_run_id(
-                task_instance
-            ) or self.adapter.build_task_instance_run_id(
-                dag_id=task_instance.dag_id,
-                task_id=task_instance.task_id,
-                try_number=task_instance.try_number,
-                logical_date=date,
-                map_index=task_instance.map_index,
-            )
+            position = get_task_instance_position(task_instance)
+            task_uuid = get_task_instance_run_id(task_instance, logical_date=date, position=position)
             event_type = RunState.COMPLETE.value.lower()
             operator_name = task.task_type.lower()
 
@@ -504,6 +492,7 @@ def on_success():
                         task,
                         task_uuid,
                         include_full_task_info=controls.include_full_task_info,
+                        ti_position=position,
                     ),
                     **get_airflow_debug_facet(),
                 },
@@ -605,15 +594,8 @@ def on_failure():
                 clear_number=dagrun.clear_number,
             )
 
-            task_uuid = get_regional_task_instance_run_id(
-                task_instance
-            ) or self.adapter.build_task_instance_run_id(
-                dag_id=task_instance.dag_id,
-                task_id=task_instance.task_id,
-                try_number=task_instance.try_number,
-                logical_date=date,
-                map_index=task_instance.map_index,
-            )
+            position = get_task_instance_position(task_instance)
+            task_uuid = get_task_instance_run_id(task_instance, logical_date=date, position=position)
             event_type = RunState.FAIL.value.lower()
             operator_name = task.task_type.lower()
 
@@ -682,6 +664,7 @@ def on_failure():
                         task,
                         task_uuid,
                         include_full_task_info=controls.include_full_task_info,
+                        ti_position=position,
                     ),
                     **get_airflow_debug_facet(),
                 },
@@ -759,15 +742,8 @@ def on_skipped():
                 clear_number=dagrun.clear_number,
             )
 
-            task_uuid = get_regional_task_instance_run_id(
-                task_instance
-            ) or self.adapter.build_task_instance_run_id(
-                dag_id=task_instance.dag_id,
-                task_id=task_instance.task_id,
-                try_number=task_instance.try_number,
-                logical_date=date,
-                map_index=task_instance.map_index,
-            )
+            position = get_task_instance_position(task_instance)
+            task_uuid = get_task_instance_run_id(task_instance, logical_date=date, position=position)
             event_type = RunState.COMPLETE.value.lower()
             operator_name = task.task_type.lower()
 
@@ -835,6 +811,7 @@ def on_skipped():
                         task,
                         task_uuid,
                         include_full_task_info=controls.include_full_task_info,
+                        ti_position=position,
                     ),
                     **get_airflow_debug_facet(),
                 },
@@ -913,16 +890,7 @@ def _on_task_instance_manual_state_change(
             # the equivalent note in `on_dag_run_running` (listener.py ~868).
             date = dagrun.logical_date or dagrun.run_after
             position = get_orm_task_instance_positions([ti], session=object_session(ti))[ti.id]
-            if position.in_loop:
-                task_uuid = str(ti.id)
-            else:
-                task_uuid = self.adapter.build_task_instance_run_id(
-                    dag_id=ti.dag_id,
-                    task_id=ti.task_id,
-                    try_number=ti.try_number,
-                    logical_date=date,
-                    map_index=position.map_index,
-                )
+            task_uuid = get_task_instance_run_id(ti, logical_date=date, position=position)
             parent_run_id = self.adapter.build_dag_run_id(
                 dag_id=ti.dag_id,
                 logical_date=date,
@@ -959,6 +927,7 @@ def _on_task_instance_manual_state_change(
                         task,
                         task_uuid,
                         include_full_task_info=include_full_task_info,
+                        ti_position=position,
                     )
 
             adapter_kwargs: dict = {