diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests/prepare_dags.py b/airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests/prepare_dags.py index 7d99e83b09ef4..bff41c6ff6e10 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests/prepare_dags.py +++ b/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), } diff --git a/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py b/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py index 2f38284efbafb..eae4bde3a2f71 100644 --- a/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py +++ b/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") diff --git a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py index c8b61bfb12f8d..29facc7c96788 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py +++ b/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")) diff --git a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowJobFacet.json b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowJobFacet.json index 51a9954de376a..b6095f00677b7 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowJobFacet.json +++ b/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.", diff --git a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json index 621eb6a761453..bfba75b5415fe 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json +++ b/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": { diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py index 6f04c7881f310..dda70e50fecd8 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py +++ b/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 = { diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py index 5ce197f0ae632..5039e235b0b6d 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py @@ -23,9 +23,8 @@ from airflow.providers.openlineage.utils.utils import ( get_job_name, get_parent_information_from_dagrun_conf, - get_regional_task_instance_run_id, get_root_information_from_dagrun_conf, - get_task_instance_map_index, + get_task_instance_run_id, ) from airflow.providers.openlineage.version_compat import AIRFLOW_V_3_0_PLUS @@ -65,13 +64,7 @@ def lineage_run_id(task_instance: TaskInstance): For more information take a look at the guide: :ref:`howto/macros:openlineage` """ - return get_regional_task_instance_run_id(task_instance) or OpenLineageAdapter.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=_get_logical_date(task_instance), - map_index=get_task_instance_map_index(task_instance), - ) + return get_task_instance_run_id(task_instance, logical_date=_get_logical_date(task_instance)) def lineage_parent_id(task_instance: TaskInstance): diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index ca43441870d47..a8a1a76d33217 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1123,9 +1123,67 @@ def team_name(cls, dagrun: DagRun) -> str | None: class TaskInstanceInfo(InfoJsonEncodable): - """Defines encoding TaskInstance object to JSON.""" + """ + Encode a task instance as ``run.facets.airflow.taskInstance``, which tells which execution sent an event. + + The fields that identify the execution, each left out when it doesn't apply: + + * ``id`` (Airflow 3+): the UUID of the attempt. Every try, and every task instance a clear creates again, + gets a new one, so it identifies one execution, and it is the ``runId`` of a task inside a loop. Airflow's + API can look it up, but it says nothing by itself. + * ``try_number``: the try, counted per ``region_id`` and ``region_index`` (see below). + * ``max_tries``: the last ``try_number`` that still gets retried on failure, so a failed try with + ``try_number > max_tries`` is final. It equals the task's ``retries`` (in ``task``) only until a clear: + clearing sets it to the last try number plus ``retries``, giving the task its retries again. + * ``map_index``: the position in a mapped task's expansion that users see. Absent for a task that isn't + mapped. + * ``loop_index``, ``loop_max_iterations`` and ``loop_group_id`` (Airflow 3.4+): the zero-based iteration + of the innermost loop around the task, that loop's limit and its group id, which can differ from + ``task.task_group`` when the task sits in a task group inside the loop. The iterations of loops further + out aren't known. ``(loop_group_id, loop_index)`` names an iteration, but not which trajectory it ran in, + nor in a nested loop which iteration of the outer loop. + * ``region_id`` and ``region_index`` (Airflow 3.4+): the coordinates Airflow stores the task instance under. + A region is one execution of a loop, or of a mapped task's expansion, and is shared by all the task + instances it creates. ``region_index`` is the position inside it: the iteration for a task in a loop that + isn't mapped, the map index for a mapped task. Neither is unique: ``(dag_id, dagRun.run_id, task_id, + region_id, region_index, try_number)`` is, as Airflow enforces it (``task_instance_try_key``). Both are + absent for a task in no region, i.e. one that is neither mapped nor in a loop. + + ``loop_decision`` (Airflow 3.4+): ``stop`` and the optional ``reason`` returned by ``until``, as evaluated by + a loop's gate, the task that ends each iteration (``gate_task_id`` of the loop group, in ``task.task_group`` + and in ``taskGroups`` of ``AirflowJobFacet``). Only the gate's COMPLETE and FAIL events carry it, as the gate + decides while it runs. Read with ``loop_index``, ``loop_max_iterations`` and the group's ``has_until``: + + * COMPLETE with ``stop`` false: the loop goes on to iteration ``loop_index + 1``. + * COMPLETE with ``stop`` true and ``has_until``: ``until`` ended the loop at iteration ``loop_index``, which + can be the last one. + * COMPLETE with ``stop`` true and no ``has_until``: the loop ran all its iterations, as one without ``until`` + always stops at ``loop_index == loop_max_iterations - 1``. + * FAIL with ``stop`` false: ``until`` was still false at the last iteration (``loop_index == + loop_max_iterations - 1``), so the loop ran out of iterations and fails. That's the only failure after the + decision is made. + * FAIL without ``loop_decision``: the gate failed before deciding, e.g. ``until`` raised or returned something + other than a bool or a ``LoopDecision``; the error facet tells which. + + How a cleared loop reads: clearing a loop from iteration ``k`` retries what was cleared in iteration ``k`` in + place, with the same region and the next ``try_number``. The iterations after ``k`` are created again in a + new region forked from the loop's region, and try numbers count per region, so the new execution of + iteration ``k + 1`` can repeat the task, ``loop_index``, ``map_index`` and ``try_number`` of the one it + replaces. Only ``region_id`` (and ``id``) tell the two apart. Each such re-run of a loop is a trajectory. + ``dagRun.clear_number`` counts clears of the whole Dag run only. + + Ordering: region ids are UUIDv7, so of two executions of one task with the same ``loop_index`` and + ``map_index``, the later trajectory has the greater ``region_id``. That is how Airflow generates them, not a + documented guarantee, and it only orders the trajectories seen: an event doesn't say whether it comes from + the second or the third trajectory, nor from which iteration its trajectory was forked. + + What events can't tell: a mapped task's ``region_id`` is the one of its own expansion, which sits inside the + loop's region without naming it. So a mapped task instance in a loop can't be matched to the trajectory of + the loop, or to the task instances of the same pass that aren't mapped (whose ``region_id`` is the loop's), + other than by ``loop_index`` and time. + """ - includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"] + includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "try_number"] casts = { "log_url": lambda ti: getattr(ti, "log_url", None), "note": lambda ti: safe_getattr(ti, "note", None), # From manual state changes only @@ -1135,11 +1193,49 @@ class TaskInstanceInfo(InfoJsonEncodable): "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None, } + def __init__(self, obj, position: TaskInstancePosition | None = None): + # Pass the position when the caller has already placed the task instance, e.g. from its stored region or + # with get_task_instance_position; it is then used as is. + self._position = position + super().__init__(obj) + def _extend_fields(self) -> None: - map_index = get_task_instance_map_index(self.obj) + position = self._position + if position is None: + # A stored task instance whose task isn't attached can't be placed without its region, and a facet + # must not fail the whole event, so then it describes the task instance without its position. + position = _get_task_instance_position(self.obj) + if position is not None: + position = _add_runtime_loop_iteration(self.obj, position) + map_index = position.map_index if position else -1 self.map_index = map_index if map_index != -1 else None self.rendered_map_index = getattr(self.obj, "rendered_map_index", None) if map_index != -1 else None self._fields.extend(("map_index", "rendered_map_index")) + if position and position.loop_index is not None: + self.loop_index = position.loop_index + self._fields.append("loop_index") + if position and position.loop_max_iterations is not None: + self.loop_max_iterations = position.loop_max_iterations + self._fields.append("loop_max_iterations") + if position and position.loop_group_id is not None: + self.loop_group_id = position.loop_group_id + self._fields.append("loop_group_id") + if AIRFLOW_V_3_0_PLUS and (ti_id := getattr(self.obj, "id", None)) is not None: + self.id = str(ti_id) + self._fields.append("id") + # Only a loop gate's runtime task instance has one, once it has evaluated ``until``. Read from a private + # attribute until the task SDK exposes it. + if (decision := getattr(self.obj, "_loop_decision", None)) is not None: + self.loop_decision = {"stop": decision.stop} + if decision.reason is not None: + self.loop_decision["reason"] = decision.reason + self._fields.append("loop_decision") + # A runtime task instance in no region carries the sentinel region id instead of None + region_id = getattr(self.obj, "region_id", None) + if isinstance(region_id, UUID) and region_id.int != 0: + self.region_id = str(region_id) + self.region_index = getattr(self.obj, "region_index", None) + self._fields.extend(("region_id", "region_index")) class AssetInfo(InfoJsonEncodable): @@ -1354,6 +1450,10 @@ class TaskGroupInfo(InfoJsonEncodable): "tooltip", "upstream_group_ids", "upstream_task_ids", + "gate_task_id", # in a loop group + "has_until", # in a loop group + "max_iterations", # in a loop group + "terminal_task_id", # in a loop group ] @@ -1409,13 +1509,14 @@ def get_airflow_run_facet( task: BaseOperator, task_uuid: str, include_full_task_info: bool = False, + ti_position: TaskInstancePosition | None = None, ) -> dict[str, AirflowRunFacet]: runtime_assets = get_runtime_outlet_assets(task_instance) return { "airflow": AirflowRunFacet( dag=DagInfo(dag), dagRun=DagRunInfo(dag_run), - taskInstance=TaskInstanceInfo(task_instance), + taskInstance=TaskInstanceInfo(task_instance, position=ti_position), task=( TaskInfoComplete(task, runtime_assets=runtime_assets) if include_full_task_info @@ -1461,6 +1562,8 @@ def get_task_duration(ti): return { "airflowState": AirflowStateRunFacet( dagRunState=dag_run_state, + # TODO: keyed by task_id only, so the mapped instances and loop iterations of one task collapse + # into a single entry (the last one by region wins). Needs the map index and loop iteration here. tasksState={ti.task_id: ti.state for ti in tis}, tasksDuration={ti.task_id: get_task_duration(ti) for ti in tis}, ) @@ -1489,18 +1592,24 @@ def is_dag_run_asset_triggered( @attrs.frozen class TaskInstancePosition: """ - Where a task instance runs: inside a loop body or not, and the map index users see. + Where a task instance runs: in a loop or not, the map index users see, and its loop iteration. - A plain mapped task keeps the region of its own expansion, so its dag, task, try, logical date and - map index stay unique; inside a loop they repeat on every pass, so only the execution UUID identifies - the run. + The loop fields are ``None`` outside a loop, and also when they can't be known: a stored task whose task + isn't attached has no ``loop_max_iterations``, nor ``loop_group_id`` when it is mapped; a stored mapped task + has no ``loop_index`` unless ``get_orm_task_instance_positions`` could read its region; and a runtime task + instance has no ``loop_index`` nor ``loop_max_iterations`` unless placed with ``get_task_instance_position``, + as only its template context holds them. ``TaskInstanceInfo`` explains what these fields mean in events. """ in_loop: bool map_index: int + loop_index: int | None = None + loop_max_iterations: int | None = None + loop_group_id: str | None = None -def _is_task_in_loop(task: Any) -> bool: +def _find_enclosing_loop(task: Any) -> tuple[bool, Any | None]: + """Tell whether a task runs in a loop, and return that loop's group when it is known.""" from airflow.sdk.definitions._internal.loop import LoopTaskGroup from airflow.sdk.definitions.taskgroup import TaskGroup from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedTaskGroup @@ -1508,9 +1617,19 @@ def _is_task_in_loop(task: Any) -> bool: group = getattr(task, "task_group", None) while isinstance(group, (TaskGroup, SerializedTaskGroup)): if isinstance(group, (LoopTaskGroup, SerializedLoopTaskGroup)): - return True + return True, group group = group.parent_group - return group is not None + if group is None: + return False, None + # A group of an unknown type could be a loop. Assuming one only gives each iteration its own runId, while + # missing one would merge the runs of all its iterations, which share every coordinate but the UUID. + log.warning( + "OpenLineage can't tell whether task `%s` runs in a loop, as its task group has an unknown type `%s`; " + "assuming it does", + getattr(task, "task_id", None), + type(group).__name__, + ) + return True, None def _get_task_instance_position( @@ -1526,9 +1645,25 @@ def _get_task_instance_position( return TaskInstancePosition(in_loop=False, map_index=map_index) if (task := getattr(task_instance, "task", None)) is None: return None + in_loop, enclosing_loop = _find_enclosing_loop(task) + loop_index: int | None = None + loop_max_iterations: int | None = None + loop_group_id = enclosing_loop.group_id if enclosing_loop is not None else None + if enclosing_loop is not None and is_stored: + loop_max_iterations = getattr(enclosing_loop, "max_iterations", None) if is_stored and not task.get_needs_expansion(): + # The stored region_index means a map index only for a mapped task. A task that isn't mapped has none, + # and inside a loop its region_index is the iteration instead. + if in_loop: + loop_index = map_index map_index = -1 - return TaskInstancePosition(in_loop=_is_task_in_loop(task), map_index=map_index) + return TaskInstancePosition( + in_loop=in_loop, + map_index=map_index, + loop_index=loop_index, + loop_max_iterations=loop_max_iterations, + loop_group_id=loop_group_id, + ) def _require_task_instance_position( @@ -1542,56 +1677,146 @@ def _require_task_instance_position( return position +def _add_runtime_loop_iteration( + task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition +) -> TaskInstancePosition: + # Run ids, map indexes and the lineage macros don't need the iteration, so only this reads the template + # context: it re-validates the params on every call, which can raise from inside a user's template. + if not position.in_loop or position.loop_index is not None or isinstance(task_instance, TaskInstance): + return position + try: + loop = task_instance.get_template_context().get("loop") + except Exception: + # The iteration only describes the run, so failing to read it must not cost the event its run id + log.warning("OpenLineage failed to read the loop iteration of task instance %s", task_instance.id) + log.debug("Exception details:", exc_info=True) + return position + if loop is None: + return position + return attrs.evolve(position, loop_index=loop.index, loop_max_iterations=loop.max_iterations) + + +def get_task_instance_position(task_instance: TaskInstance | RuntimeTaskInstance) -> TaskInstancePosition: + """Place a task instance with its loop iteration, once per event, to pass to ``get_airflow_run_facet``.""" + return _add_runtime_loop_iteration(task_instance, _require_task_instance_position(task_instance)) + + def get_orm_task_instance_positions( task_instances: Collection[TaskInstance], *, session: Session | None ) -> dict[UUID, TaskInstancePosition]: """ - Place stored task instances, reading the regions of those without their task in one query. + Place stored task instances, reading in one query the regions of those without their task and of mapped tasks in a loop. - Server side only: workers never hold a session, and an instance that needs its region read without one - raises instead of opening a connection. + Server side only: workers never hold a session. An instance without its task raises when there is no session to + read its region, while a mapped task in a loop is then placed without its iteration. """ positions: dict[UUID, TaskInstancePosition] = {} unplaced: list[TaskInstance] = [] + without_iteration: list[TaskInstance] = [] for task_instance in task_instances: if (position := _get_task_instance_position(task_instance)) is None: unplaced.append(task_instance) else: positions[task_instance.id] = position - if not unplaced: + # A stored mapped task in a loop keeps its iteration on its region only. Without a session it stays unknown. + if position.in_loop and position.loop_index is None and session is not None: + without_iteration.append(task_instance) + if not unplaced and not without_iteration: return positions if session is None: + # Only unplaced instances get here without a session, as no mapped task is added to without_iteration then raise ValueError("Placing a stored task instance without its task needs the session it was loaded in") from sqlalchemy import select + from sqlalchemy.orm import aliased from airflow.models.dynamic_region import DynamicRegion - regions = { - row.id: row - for row in session.execute( - select(DynamicRegion.id, DynamicRegion.node_id, DynamicRegion.parent_region_id).where( - DynamicRegion.id.in_({task_instance.region_id for task_instance in unplaced}) + parent = aliased(DynamicRegion) + query = ( + select( + DynamicRegion.id, + DynamicRegion.node_id, + DynamicRegion.parent_region_id, + DynamicRegion.parent_region_index, + parent.node_id.label("parent_node_id"), + parent.parent_region_id.label("grandparent_region_id"), + ) + .outerjoin(parent, parent.id == DynamicRegion.parent_region_id) + .where( + DynamicRegion.id.in_( + {task_instance.region_id for task_instance in (*unplaced, *without_iteration)} ) ) - } + ) + if unplaced: + regions = {row.id: row for row in session.execute(query)} + else: + # Only mapped tasks' iterations are missing, and they only describe the run, so failing to read them + # must not cost the event. The savepoint keeps the caller's transaction usable after a failure. + try: + with session.begin_nested(): + regions = {row.id: row for row in session.execute(query)} + except Exception: + log.warning( + "OpenLineage failed to read the loop iteration of mapped task instances from their regions" + ) + log.debug("Exception details:", exc_info=True) + return positions + + # Airflow creates a region for each run of a loop and for the expansion of each mapped task, a task in a + # mapped task group included. As loops can't be nested nor sit in a mapped task group, a loop's region is + # top-level and a mapped task's region sits directly in its loop's region, if any. The provider's + # test_region_shapes_* tests check this against Airflow. Any other shape is placed as a loop without its + # iteration, which keeps the runId of each execution unique, instead of guessing. for task_instance in unplaced: region = regions.get(task_instance.region_id) + position = None + map_index = -1 if region is not None and region.node_id == task_instance.task_id: - positions[task_instance.id] = TaskInstancePosition( - in_loop=region.parent_region_id is not None, map_index=task_instance.region_index + map_index = task_instance.region_index + if region.parent_region_id is None: + position = TaskInstancePosition(in_loop=False, map_index=map_index) + elif region.grandparent_region_id is None: + position = TaskInstancePosition( + in_loop=True, + map_index=map_index, + loop_index=region.parent_region_index, + loop_group_id=region.parent_node_id, + ) + elif region is not None and region.parent_region_id is None: + position = TaskInstancePosition( + in_loop=True, + map_index=-1, + loop_index=task_instance.region_index, + loop_group_id=region.node_id, ) + if position is None: + log.warning( + "OpenLineage can't place task instance %s from its region, which has an unexpected shape; " + "assuming it runs in a loop", + task_instance.id, + ) + position = TaskInstancePosition(in_loop=True, map_index=map_index) + positions[task_instance.id] = position + for task_instance in without_iteration: + position = positions[task_instance.id] + region = regions.get(task_instance.region_id) + if ( + region is not None + and region.parent_node_id == position.loop_group_id + and region.grandparent_region_id is None + ): + positions[task_instance.id] = attrs.evolve(position, loop_index=region.parent_region_index) else: - positions[task_instance.id] = TaskInstancePosition(in_loop=True, map_index=-1) + log.warning( + "OpenLineage can't read the loop iteration of task instance %s, as its region isn't directly " + "in its loop's region", + task_instance.id, + ) return positions -def get_regional_task_instance_run_id(task_instance: TaskInstance | RuntimeTaskInstance) -> str | None: - if _require_task_instance_position(task_instance).in_loop: - return str(task_instance.id) - return None - - def build_task_instance_ol_run_id( dag_id: str, task_id: str, @@ -1620,6 +1845,31 @@ def build_task_instance_ol_run_id( ) +def get_task_instance_run_id( + task_instance: TaskInstance | RuntimeTaskInstance, + *, + logical_date: datetime.datetime, + position: TaskInstancePosition | None = None, +) -> str: + """ + Return the OpenLineage ``runId`` of a task instance's attempt, which all its events must share. + + In a loop the iterations share every coordinate but the attempt UUID, so that is the ``runId``. Otherwise + it is built from the coordinates. + """ + if position is None: + position = _require_task_instance_position(task_instance) + if position.in_loop: + return str(task_instance.id) + return build_task_instance_ol_run_id( + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + try_number=task_instance.try_number, + logical_date=logical_date, + map_index=position.map_index, + ) + + def is_valid_uuid(uuid_string: str | None) -> bool: """Validate that a string is a valid UUID format.""" if uuid_string is None: @@ -2049,6 +2299,11 @@ def _get_task_groups_details(dag: DAG | SerializedDAG, edge_map: dict[str, tuple "ui_label": tg.label, "downstream_task_edges": task_edges, "downstream_group_edges": group_edges, + **{ + field: getattr(tg, field) # Only a loop group has these + for field in ("gate_task_id", "has_until", "max_iterations", "terminal_task_id") + if hasattr(tg, field) + }, } return result diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py new file mode 100644 index 0000000000000..67da53b777c4d --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py @@ -0,0 +1,68 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +DAG with a task that fails on its first try and succeeds on the retry. + +It checks: + - the first try emits START and FAIL, the retry emits START and COMPLETE + - taskInstance's `try_number` is 1 for the first try and 2 for the retry, and `max_tries` is 1 for both + - each try has its own `runId` (the `runId` check of OpenLineageTestOperator is enabled by default) +""" + +from __future__ import annotations + +from datetime import datetime, timedelta + +from airflow import DAG + +try: + from airflow.sdk import task +except ImportError: + # Airflow 2 path + from airflow.decorators import task # type: ignore[attr-defined,no-redef] + +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT +from system.openlineage.expected_events import get_expected_event_file_path +from system.openlineage.operator import OpenLineageTestOperator + +DAG_ID = "openlineage_retry_dag" + +with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 1, "retry_delay": timedelta(seconds=1)}, +) as dag: + + @task + def flaky_task(**context): + if context["ti"].try_number == 1: + raise ValueError("Intentional failure of the first try") + + check_events = OpenLineageTestOperator( + task_id="check_events", file_path=get_expected_event_file_path(DAG_ID) + ) + + flaky_task() >> check_events + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py new file mode 100644 index 0000000000000..802974519c687 --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py @@ -0,0 +1,67 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +DAG with a sensor in reschedule mode. + +It checks: + - the sensor emits one START (the START of later pokes is skipped) and one COMPLETE + - taskInstance's `try_number` stays at 1 for both, as rescheduling does not start a new try + - START and COMPLETE share the `runId` (the `runId` check of OpenLineageTestOperator is enabled by default) +""" + +from __future__ import annotations + +import os +from datetime import datetime, timedelta + +from airflow import DAG +from airflow.providers.standard.sensors.time_delta import TimeDeltaSensor + +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT +from system.openlineage.expected_events import get_expected_event_file_path +from system.openlineage.operator import OpenLineageTestOperator + +DAG_ID = "openlineage_sensor_reschedule_dag" + +with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 0}, +) as dag: + # Timedelta is compared to the DAGRun start timestamp, which can occur long before a worker picks up the + # task. We need to ensure the sensor is rescheduled at least once, so setting 180s. + wait = TimeDeltaSensor( + task_id="wait", + delta=timedelta(seconds=int(os.getenv("SYSTEM_TESTS_OL_TIMEDELTA", "180"))), + poke_interval=10, + mode="reschedule", + ) + + check_events = OpenLineageTestOperator( + task_id="check_events", file_path=get_expected_event_file_path(DAG_ID) + ) + + wait >> check_events + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py new file mode 100644 index 0000000000000..e4cf123c12fc7 --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py @@ -0,0 +1,151 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +DAG with a plain task and task loops, some of their tasks retrying and one stopped early by `until` +(requires Airflow 3.4+). + +It checks, for a plain task, tasks, a mapped task and a mapped task group in a loop, and loop gates: + - taskInstance's `try_number`, `map_index`, `loop_index` and `loop_max_iterations` of every event: the + iteration is the zero-based `loop_index`, `map_index` the mapped position, and a plain task has no + `loop_index` + - taskInstance's `loop_group_id` of every event: the loop group of a task in a loop, absent for a plain task + - taskInstance's `max_tries` of the tasks other than the gates, the `retries` of each task + - taskInstance's `id` of every event, and `region_id` and `region_index` of the tasks in a loop: the + `region_index` is the iteration for a task that isn't mapped and the map index for a mapped one, and a + plain task has neither + - the gates' `loop_decision`: absent on START; for the loops without `until`, `stop` true only at the last + iteration and no `reason`; for the loop with `until` (whose gate is named after it), the `reason` it + returns and `stop` true at the iteration it ends the loop, before `max_iterations`; absent for the plain + task + - a mapped task group inside a loop: its task's `map_index` and `region_index` are the mapped position, + its `loop_index` the iteration, and its `loop_group_id` the loop's, unlike its own `task_group` + - the loop groups' `max_iterations`, `has_until`, `gate_task_id` and `terminal_task_id`, in the task's + `task_group` and in the `taskGroups` of the Dag's job facet + - a retried task emits START and FAIL for each failed try, then START and COMPLETE for the last one + - every execution (iteration, map index, try) has its own `runId`, shared by its START and COMPLETE/FAIL + (the `runId` check of OpenLineageTestOperator is enabled by default) +""" + +from __future__ import annotations + +from datetime import datetime, timedelta + +from airflow import DAG +from airflow.sdk import LoopDecision, task, task_group + +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT +from system.openlineage.expected_events import get_expected_event_file_path +from system.openlineage.operator import OpenLineageTestOperator + +DAG_ID = "openlineage_task_loop_dag" + + +def sort_by_execution(event: dict) -> tuple[int, int, int]: + # Dag run events have no airflow facet + task_instance = event["run"]["facets"].get("airflow", {}).get("taskInstance", {}) + return ( + task_instance.get("loop_index", -1), + task_instance.get("map_index") if task_instance.get("map_index") is not None else -1, + task_instance.get("try_number", -1), + ) + + +with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 0}, +) as dag: + + @task + def plain_task(): + return None + + @task_group + def accumulate(): + @task(retries=2, retry_delay=timedelta(seconds=1)) + def increment(*, loop, ti): + # The second iteration fails on its first try + if loop.index == 1 and ti.try_number < 2: + raise ValueError(f"Intentional failure, iteration {loop.index}, try {ti.try_number}") + return (0 if loop.index == 0 else loop.previous) + 1 + + @task + def report(value): + return value + + report(increment()) + + @task_group + def process_batch(): + @task + def choose_items(): + return [1, 2] + + @task( + max_active_tis_per_dagrun=1, # Execute sequentially to not overwrite each other OL events + retries=3, + retry_delay=timedelta(seconds=1), + ) + def process(value, *, loop, ti): + # The mapped instance of the second item fails on its first two tries, in every iteration + if value == 2 and ti.try_number < 3: + raise ValueError(f"Intentional failure, iteration {loop.index}, try {ti.try_number}") + return value + loop.index + + @task_group + def fetch(source): + @task(max_active_tis_per_dagrun=1) # Execute sequentially to not overwrite each other OL events + def download(source, *, loop): + return f"{source}-{loop.index}" + + download(source) + + items = choose_items() + fetch.expand(source=["a", "b"]) >> process.expand(value=items) + + @task_group + def poll(): + @task + def check(*, loop): + return loop.index + + check() + + def ready(*, loop): + # Stops at the second of three iterations, giving a reason both when continuing and when stopping + return LoopDecision(stop=loop.result >= 1, reason=f"check returned {loop.result}") + + accumulate_loop = accumulate.loop(max_iterations=3) + process_batch_loop = process_batch.loop(max_iterations=2) + poll_loop = poll.loop(max_iterations=3, until=ready) + + check_events = OpenLineageTestOperator( + task_id="check_events", + file_path=get_expected_event_file_path(DAG_ID), + event_sort_fn=sort_by_execution, + ) + + [plain_task(), accumulate_loop, process_batch_loop, poll_loop] >> check_events + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_manual_lineage_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_manual_lineage_dag.json index 452f3551746fc..2d3187f82d89c 100644 --- a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_manual_lineage_dag.json +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_manual_lineage_dag.json @@ -1,4 +1,20 @@ [ + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": "{{ result is number }}" + } + } + } + }, + "job": { + "name": "openlineage_manual_lineage_dag.datasets_minimal" + } + }, { "eventType": "RUNNING", "eventTime": "{{ is_datetime(result) }}", @@ -133,6 +149,38 @@ } } }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": "{{ result is number }}" + } + } + } + }, + "job": { + "name": "openlineage_manual_lineage_dag.datasets_minimal" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": "{{ result is number }}" + } + } + } + }, + "job": { + "name": "openlineage_manual_lineage_dag.datasets_maximal" + } + }, { "eventType": "RUNNING", "eventTime": "{{ is_datetime(result) }}", @@ -321,6 +369,22 @@ } } }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": "{{ result is number }}" + } + } + } + }, + "job": { + "name": "openlineage_manual_lineage_dag.datasets_maximal" + } + }, { "eventType": "START", "eventTime": "{{ is_datetime(result) }}", diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_retry_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_retry_dag.json new file mode 100644 index 0000000000000..678816d96241a --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_retry_dag.json @@ -0,0 +1,76 @@ +[ + { + "eventType": "START", + "job": { + "name": "openlineage_retry_dag" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 1 + } + } + } + }, + "job": { + "name": "openlineage_retry_dag.flaky_task" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 1 + } + } + } + }, + "job": { + "name": "openlineage_retry_dag.flaky_task" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "max_tries": 1 + } + } + } + }, + "job": { + "name": "openlineage_retry_dag.flaky_task" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "max_tries": 1 + } + } + } + }, + "job": { + "name": "openlineage_retry_dag.flaky_task" + } + } +] diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_sensor_reschedule_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_sensor_reschedule_dag.json new file mode 100644 index 0000000000000..8ab51fa095221 --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_sensor_reschedule_dag.json @@ -0,0 +1,34 @@ +[ + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1 + } + } + } + }, + "job": { + "name": "openlineage_sensor_reschedule_dag.wait" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1 + } + } + } + }, + "job": { + "name": "openlineage_sensor_reschedule_dag.wait" + } + } +] diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json new file mode 100644 index 0000000000000..6fc30d3d73258 --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json @@ -0,0 +1,2047 @@ +[ + { + "eventType": "START", + "job": { + "name": "openlineage_task_loop_dag", + "facets": { + "airflow": { + "taskGroups": { + "accumulate": { + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + }, + "process_batch": { + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + }, + "poll": { + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + }, + "process_batch.fetch": { + "parent_group": "process_batch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": null, + "loop_max_iterations": null, + "id": "{{ is_uuid(result) }}", + "region_id": null, + "region_index": null, + "max_tries": 0, + "loop_group_id": null, + "loop_decision": null + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.plain_task" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": null, + "loop_max_iterations": null, + "id": "{{ is_uuid(result) }}", + "region_id": null, + "region_index": null, + "max_tries": 0, + "loop_group_id": null, + "loop_decision": null + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.plain_task" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "accumulate", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "accumulate", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "loop_group_id": "accumulate", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "accumulate", + "loop_decision": { + "stop": false + } + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "accumulate", + "loop_decision": { + "stop": false + } + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "loop_group_id": "accumulate", + "loop_decision": { + "stop": true + } + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.choose_items" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.choose_items" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.choose_items" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.choose_items" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 0, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 3, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 0, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 3, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 0, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 3, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 0, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 3, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "process_batch", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "process_batch", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "process_batch", + "loop_decision": { + "stop": false + } + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "process_batch", + "loop_decision": { + "stop": true + } + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "poll", + "max_tries": 0 + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.check" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "poll", + "max_tries": 0 + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.check" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "poll", + "max_tries": 0 + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.check" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "poll", + "max_tries": 0 + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.check" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "poll", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.ready" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "poll", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.ready" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "poll", + "loop_decision": { + "stop": false, + "reason": "check returned 0" + } + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.ready" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "poll", + "loop_decision": { + "stop": true, + "reason": "check returned 1" + } + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.ready" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 0, + "loop_index": 0, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 0, + "loop_index": 1, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 0, + "loop_index": 0, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 0, + "loop_index": 1, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + } +] diff --git a/providers/openlineage/tests/system/openlineage/operator.py b/providers/openlineage/tests/system/openlineage/operator.py index 802ac24320d60..104ee964d404b 100644 --- a/providers/openlineage/tests/system/openlineage/operator.py +++ b/providers/openlineage/tests/system/openlineage/operator.py @@ -23,8 +23,9 @@ import re import time import uuid -from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from collections import defaultdict +from collections.abc import Callable, Collection +from typing import TYPE_CHECKING, Any, NamedTuple from urllib.parse import urlparse from dateutil.parser import parse @@ -251,6 +252,127 @@ def match(expected, result, env: Environment, path: list | None = None) -> bool: return True +class Attempt(NamedTuple): + """ + Which execution of a job an event belongs to. + + ``try_number``, ``map_index``, ``loop_index`` and ``region_id`` are None for a Dag event and when the task has + none. ``clear_number`` plays for a Dag run the role ``try_number`` plays for a task instance, and ``region_id`` + tells apart two executions of one loop iteration after the loop was cleared, which can share every other field. + """ + + dag_run_id: Any + clear_number: Any + try_number: Any + map_index: Any + loop_index: Any + region_id: Any + + +def _get_attempt(event: dict) -> Attempt | None: + """Return the attempt of an event, or None for an event without the airflow facet.""" + facets = (event.get("run") or {}).get("facets") or {} + # Task events carry the airflow facet, Dag run events carry the same content in airflowDagRun. + airflow_facet = facets.get("airflow") or facets.get("airflowDagRun") + if not isinstance(airflow_facet, dict): + return None + dag_run = airflow_facet.get("dagRun") or {} + task_instance = airflow_facet.get("taskInstance") or {} + return Attempt( + dag_run_id=dag_run.get("run_id"), + clear_number=dag_run.get("clear_number"), + try_number=task_instance.get("try_number"), + map_index=task_instance.get("map_index"), + loop_index=task_instance.get("loop_index"), + region_id=task_instance.get("region_id"), + ) + + +def _get_job_type(event: dict) -> str | None: + return (((event.get("job") or {}).get("facets") or {}).get("jobType") or {}).get("jobType") + + +def check_run_id_integrity( + events_by_key: dict[str, list[dict]], repeated_start_keys: Collection[str] = () +) -> list[str]: + """ + Check that ``run.runId`` identifies exactly one attempt of exactly one job, and return the problems found. + + Events are grouped by ``runId``. The attempt of an event is its Dag run id and clear number, plus for a task + event ``(try_number, map_index, loop_index, region_id)``, all from ``run.facets.airflow`` (``airflowDagRun`` for Dag + events). Query events, which are child runs of a task, have no attempt. For every group: + + 1. all events share one ``job.name``; + 2. all events share one attempt, so another Dag run, clear, retry, map index or loop iteration can't reuse + a ``runId``; + 3. the lifecycle is complete, so START and COMPLETE/FAIL of an attempt can't carry different ``runId`` values: + + * task and query: exactly one START, or several when the key of the START events is in + ``repeated_start_keys`` (deferral re-emits START on every resume), and exactly one COMPLETE or FAIL; + * Dag run: one START and at most one COMPLETE or FAIL. The Dag run that contains the checking task is + still running when the check happens, so its COMPLETE/FAIL is missing, while Dag runs triggered by it + can already have finished. + + Across groups: + + 4. one attempt of a job has one ``runId``. + + Only the events passed in are checked, i.e. those of the keys listed in the expected events. + """ + problems: list[str] = [] + groups: dict[str, list[tuple[str, dict]]] = defaultdict(list) + for key, events in events_by_key.items(): + for event in events: + run_id = (event.get("run") or {}).get("runId") + if run_id: + groups[run_id].append((key, event)) + else: + problems.append(f"[{key}] {event.get('eventType')} event has no run.runId") + + run_ids_by_attempt: dict[tuple, set[str]] = defaultdict(set) + for run_id, members in sorted(groups.items()): + job_names = {event["job"]["name"] for _, event in members} + if len(job_names) > 1: + problems.append(f"runId {run_id} is shared by different jobs: {sorted(job_names)}") + continue + job_name = next(iter(job_names)) + + attempts = {attempt for _, event in members if (attempt := _get_attempt(event)) is not None} + if len(attempts) > 1: + problems.append( + f"runId {run_id} of `{job_name}` is shared by different attempts: {sorted(attempts, key=str)}" + ) + for attempt in attempts: + run_ids_by_attempt[(job_name, attempt)].add(run_id) + + job_type = _get_job_type(members[0][1]) + if job_type not in ("TASK", "QUERY", "DAG"): + continue + event_types = [event.get("eventType") for _, event in members] + start_keys = {key for key, event in members if event.get("eventType") == "START"} + terminals = [event_type for event_type in event_types if event_type in ("COMPLETE", "FAIL")] + starts = event_types.count("START") + repeated_starts_allowed = job_type != "DAG" and bool(start_keys & set(repeated_start_keys)) + too_many_starts = starts > 1 and not repeated_starts_allowed + too_many_terminals = len(terminals) > 1 or (job_type != "DAG" and not terminals) + if starts < 1 or too_many_starts or too_many_terminals: + expected = ( + "one START (several only if the Dag expects repeated STARTs, see event_count_assertions) " + "and one COMPLETE/FAIL" + if job_type != "DAG" + else "one START and at most one COMPLETE/FAIL" + ) + problems.append( + f"runId {run_id} of {job_type} `{job_name}` should have {expected}, " + f"got: {sorted(str(event_type) for event_type in event_types)}" + ) + + for (job_name, attempt_fields), run_ids in sorted(run_ids_by_attempt.items(), key=str): + if len(run_ids) > 1: + problems.append(f"`{job_name}` {attempt_fields} has several runIds: {sorted(run_ids)}") + return problems + + class OpenLineageTestOperator(BaseOperator): """ This operator is added for system testing purposes. @@ -289,6 +411,10 @@ class OpenLineageTestOperator(BaseOperator): ascending before comparison. Useful for mapped tasks whose events arrive in non-deterministic order. When the template for a key is a list, sorting is required for pairwise comparison to be meaningful. + :param check_run_ids: if True (default), additionally check the collected events with + ``check_run_id_integrity``: START and COMPLETE/FAIL of one attempt share a ``runId``, while another + try, map index or loop iteration gets its own. Set to False only for a Dag whose expected events + intentionally list a task's lifecycle partially. It runs only once every key matches its template. Special sentinel values in expected dicts: @@ -316,6 +442,7 @@ def __init__( clear_variables: bool = True, fail_fast: bool = False, event_sort_fn: Callable[[dict], Any] | None = None, + check_run_ids: bool = True, **kwargs, ): super().__init__(**kwargs) @@ -325,6 +452,7 @@ def __init__( self.event_count_assertions = event_count_assertions self.clear_variables = clear_variables self.fail_fast = fail_fast + self.check_run_ids = check_run_ids self.event_sort_fn = event_sort_fn if event_templates is None and file_path is None: raise ValueError("Either event_templates or file_path must be provided") @@ -351,15 +479,28 @@ def execute(self, context: Context) -> None: else: self.event_templates[key] = [existing, event] errors: dict[str, str] = {} + events_by_key: dict[str, list[dict]] = {} try: for key, template in self.event_templates.items(): # type: ignore[union-attr] log.info("Checking key: `%s`", key) try: - self._check_key(key, template) + events_by_key[key] = self._check_key(key, template) except ValueError as e: if self.fail_fast: raise errors[key] = str(e) + if self.check_run_ids and errors: + # A key that failed its template is missing from events_by_key, so its runIds would look incomplete + log.info("Skipping the runId integrity check until the event templates match") + elif self.check_run_ids: + if problems := check_run_id_integrity( + events_by_key, + repeated_start_keys={key for key in events_by_key if self._expects_repeated_events(key)}, + ): + message = "\n".join(f" {problem}" for problem in problems) + if self.fail_fast: + raise ValueError(f"runId integrity check failed:\n{message}") + errors[""] = f"\n{message}" finally: if self.clear_variables: for key in self.event_templates: # type: ignore[union-attr] @@ -367,9 +508,10 @@ def execute(self, context: Context) -> None: Variable.delete(key=key) if errors: failures = "\n".join(f" [{key}] {error}" for key, error in errors.items()) + log.error("Event check failures:\n%s", failures) raise ValueError(f"Event check failures:\n{failures}") - def _check_key(self, key: str, template: Any) -> None: + def _check_key(self, key: str, template: Any) -> list[dict]: if isinstance(template, dict) and "$absent" in template: reason = template["$absent"] try: @@ -377,7 +519,7 @@ def _check_key(self, key: str, template: Any) -> None: raise ValueError(f"Expected no events for key {key} ({reason}), got {len(actual)}") except NoVariableError: log.info("Key `%s` absent as expected (%s).", key, reason) - return + return [] try: actual_events = Variable.get(key=key, deserialize_json=True) @@ -403,15 +545,17 @@ def _check_key(self, key: str, template: Any) -> None: default_count = str(len(template)) if isinstance(template, list) else "1" self._assert_event_count(key, len(actual_events), default=default_count) + parsed_events = [json.loads(evt_str) for evt_str in actual_events] if isinstance(template, list): # Multiple expected events: compare each one after sorting. - for i, (tmpl, evt_str) in enumerate(zip(template, actual_events)): - if not match(tmpl, json.loads(evt_str), self.env, [i]): + for i, (tmpl, event) in enumerate(zip(template, parsed_events)): + if not match(tmpl, event, self.env, [i]): raise ValueError(f"Event at index {i} does not match template for key `{key}`") else: # Last event is checked against the template - if not match(template, json.loads(actual_events[-1]), self.env): + if not match(template, parsed_events[-1], self.env): raise ValueError(f"Event received does not match one specified in test for key `{key}`") + return parsed_events def _assert_event_count(self, key: str, actual: int, default: str = "1") -> None: """ @@ -422,13 +566,25 @@ def _assert_event_count(self, key: str, actual: int, default: str = "1") -> None templates, str(len(template)) for list templates — supplied by the caller). Delegates the numeric check to _check_count_expr. """ - expr = default + self._check_count_expr(key, actual, self._get_count_expr(key, default)) + + def _get_count_expr(self, key: str, default: str) -> str: if self.event_count_assertions: for pattern, count_expr in self.event_count_assertions.items(): if re.fullmatch(pattern, key): - expr = str(count_expr) - break - self._check_count_expr(key, actual, expr) + return str(count_expr) + return default + + def _expects_repeated_events(self, key: str) -> bool: + """ + Tell whether the Dag declares that a START key holds several START events of one run. + + An open-ended ``">=N"`` (N >= 2) assertion on a START key does that, e.g. for a deferrable task that + emits START again after every resume. Exact counts don't: they also cover mapped tasks, where each + instance emits its own START under its own ``runId``. + """ + expr = self._get_count_expr(key, default="1").strip() + return key.endswith(".event.start") and expr.startswith(">=") and int(expr[2:]) >= 2 @staticmethod def _check_count_expr(key: str, actual: int, expr: str) -> None: diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index d377e55d74abf..32440f2402697 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -49,6 +49,7 @@ from airflow.providers.openlineage.plugins.macros import lineage_run_id from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy from airflow.providers.openlineage.utils.selective_enable import disable_lineage, enable_lineage +from airflow.providers.openlineage.utils.utils import TaskInstancePosition from airflow.utils import types from airflow.utils.state import DagRunState, State @@ -291,8 +292,27 @@ def test_executor_initializer_with_real_stats(self, mock_configure_orm): mock_configure_orm.assert_called_once() +def _build_fake_task_run_id(dag_id, task_id, try_number, logical_date, map_index): + return f"{logical_date.isoformat()}.{dag_id}.{task_id}.{try_number}.{map_index}" + + +@pytest.fixture +def task_run_id_builder(): + """Patch how a task's runId is built outside a loop, with a readable fake that keeps its inputs.""" + with mock.patch( + "airflow.providers.openlineage.utils.utils.build_task_instance_ol_run_id", + autospec=True, + side_effect=_build_fake_task_run_id, + ) as builder: + yield builder + + @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Airflow 2 tests") class TestOpenLineageListenerAirflow2: + @pytest.fixture(autouse=True) + def _use_task_run_id_builder(self, task_run_id_builder): + self.build_task_run_id = task_run_id_builder + @patch("airflow.models.TaskInstance.xcom_push") @patch("airflow.models.BaseOperator.render_template") def test_listener_does_not_change_task_instance(self, render_mock, xcom_push_mock): @@ -420,9 +440,6 @@ def _create_listener_and_task_instance(self) -> tuple[OpenLineageListener, TaskI def mock_dag_id(dag_id, logical_date, clear_number): return f"{logical_date.isoformat()}.{dag_id}.{clear_number}" - def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): - return f"{logical_date.isoformat()}.{dag_id}.{task_id}.{try_number}.{map_index}" - listener = OpenLineageListener() listener.extractor_manager = mock.Mock() @@ -432,7 +449,6 @@ def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): adapter = mock.Mock() adapter.build_dag_run_id.side_effect = mock_dag_id - adapter.build_task_instance_run_id.side_effect = mock_task_id adapter.start_task = mock.Mock() adapter.fail_task = mock.Mock() adapter.complete_task = mock.Mock() @@ -874,12 +890,12 @@ def test_on_task_instance_running_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is in the running state. This test ensures that when an Airflow task instance transitions to the running state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_running(None, task_instance, None) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -894,7 +910,7 @@ def test_on_task_instance_failed_correctly_calls_openlineage_adapter_run_id_meth """Tests the OpenLineageListener's response when a task instance is in the failed state. This test ensures that when an Airflow task instance transitions to the failed state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() @@ -903,7 +919,7 @@ def test_on_task_instance_failed_correctly_calls_openlineage_adapter_run_id_meth listener.on_task_instance_failed( previous_state=None, task_instance=task_instance, **on_task_failed_kwargs, session=None ) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -918,12 +934,12 @@ def test_on_task_instance_success_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is in the success state. This test ensures that when an Airflow task instance transitions to the success state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_success(None, task_instance, None) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -1010,7 +1026,7 @@ def test_listener_on_task_instance_running_do_not_call_adapter_when_disabled_ope listener.on_task_instance_running(None, task_instance, None) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.start_task.assert_not_called() @@ -1036,7 +1052,7 @@ def test_listener_on_task_instance_failed_do_not_call_adapter_when_disabled_oper ) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.fail_task.assert_not_called() @@ -1058,7 +1074,7 @@ def test_listener_on_task_instance_success_do_not_call_adapter_when_disabled_ope listener.on_task_instance_success(None, task_instance, None) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.complete_task.assert_not_called() @@ -1140,6 +1156,10 @@ def set_result(*args, **kwargs): @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Airflow 3 tests") class TestOpenLineageListenerAirflow3: + @pytest.fixture(autouse=True) + def _use_task_run_id_builder(self, task_run_id_builder): + self.build_task_run_id = task_run_id_builder + @pytest.mark.parametrize("state", ["success", "failed", "skipped"]) @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") def test_manual_regional_event_uses_same_execution_uuid(self, mocker, state): @@ -1149,9 +1169,6 @@ def test_manual_regional_event_uses_same_execution_uuid(self, mocker, state): listener, ti = self._create_listener_and_task_instance(runtime_ti=False) listener._executor = mocker.Mock(spec=ProcessPoolExecutor) listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id - listener.adapter.build_task_instance_run_id.side_effect = ( - OpenLineageAdapter.build_task_instance_run_id - ) with DAG("loop_dag") as loop_dag: @task_group @@ -1172,7 +1189,7 @@ def body(): listener._on_task_instance_manual_state_change(ti, ti.dag_run, state) assert submit.call_args.kwargs["run_id"] == str(ti.id) - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) @pytest.mark.parametrize("map_index", [-1, 2]) @@ -1194,9 +1211,6 @@ def body(): create_loop(body, max_iterations=2) ti.task = loop_dag.get_task("body.work") listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id - listener.adapter.build_task_instance_run_id.side_effect = ( - OpenLineageAdapter.build_task_instance_run_id - ) ti.region_id = uuid.uuid4() ti.region_index = 2 ti.map_index = map_index @@ -1207,7 +1221,113 @@ def body(): method = {"running": "start_task", "failed": "fail_task"}.get(state, "complete_task") assert getattr(listener.adapter, method).call_args.kwargs["run_id"] == str(ti.id) assert lineage_run_id(ti) == str(ti.id) - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() + + @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_worker_event_of_task_in_loop_has_its_iteration(self, state): + from airflow.sdk import DAG, task_group + from airflow.sdk.api.datamodels._generated import LoopContext + from airflow.sdk.definitions._internal.loop import create_loop + + listener, ti = self._create_listener_and_task_instance() + listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id + with DAG("loop_dag") as loop_dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti.task = loop_dag.get_task("body.work") + ti.region_id = uuid.uuid4() + ti.region_index = 1 + ti._ti_context_from_server.loop = LoopContext( + node_id="body", index=1, max_iterations=3, terminal_task_id="body.work", terminal_is_mapped=False + ) + kwargs = {"error": ValueError("test")} if state == "failed" else {} + + getattr(listener, f"on_task_instance_{state}")(None, ti, **kwargs) + + method = {"running": "start_task", "failed": "fail_task"}.get(state, "complete_task") + run_facets = getattr(listener.adapter, method).call_args.kwargs["run_facets"] + task_instance_facet = run_facets["airflow"].taskInstance + assert ( + task_instance_facet["loop_index"], + task_instance_facet["loop_max_iterations"], + task_instance_facet["loop_group_id"], + ) == (1, 3, "body") + + @pytest.mark.parametrize("state", ["running", "success", "failed"]) + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_worker_event_of_loop_gate_has_its_decision_once_evaluated(self, state): + from airflow.sdk.api.datamodels._generated import LoopDecisionPayload + + listener, ti = self._create_listener_and_task_instance() + listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id + if state != "running": + ti._loop_decision = LoopDecisionPayload(stop=state == "success", reason="checked") + kwargs = {"error": ValueError("test")} if state == "failed" else {} + + getattr(listener, f"on_task_instance_{state}")(None, ti, **kwargs) + + method = {"running": "start_task", "failed": "fail_task"}.get(state, "complete_task") + task_instance_facet = ( + getattr(listener.adapter, method).call_args.kwargs["run_facets"]["airflow"].taskInstance + ) + if state == "running": + assert "loop_decision" not in task_instance_facet + else: + assert task_instance_facet["loop_decision"] == {"stop": state == "success", "reason": "checked"} + + @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_worker_event_run_id_of_mapped_task_has_its_map_index(self, state): + listener, ti = self._create_listener_and_task_instance() + ti.map_index = 2 + kwargs = {"error": ValueError("test")} if state == "failed" else {} + + getattr(listener, f"on_task_instance_{state}")(None, ti, **kwargs) + + assert self.build_task_run_id.call_args.kwargs["map_index"] == 2 + + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") + def test_manual_event_passes_its_stored_position_to_the_airflow_run_facet(self, mocker): + from airflow.sdk import DAG, task_group + from airflow.sdk.definitions._internal.loop import create_loop + + listener, ti = self._create_listener_and_task_instance(runtime_ti=False) + listener._executor = mocker.Mock(spec=ProcessPoolExecutor) + with DAG("loop_dag") as loop_dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti.task = loop_dag.get_task("body.work") + ti.region_id = uuid.uuid4() + ti.region_index = 2 + get_airflow_run_facet = mocker.patch( + "airflow.providers.openlineage.plugins.listener.get_airflow_run_facet", + autospec=True, + return_value={}, + ) + mocker.patch.object(listener, "submit_callable", autospec=True) + + listener._on_task_instance_manual_state_change(ti, ti.dag_run, "success") + + assert get_airflow_run_facet.call_args.kwargs["ti_position"] == TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=2, loop_max_iterations=3, loop_group_id="body" + ) @pytest.mark.skip("Rendering fields is not migrated yet in Airflow 3") @patch("airflow.models.BaseOperator.render_template") @@ -1465,9 +1585,6 @@ def _create_listener_and_task_instance( def mock_dag_id(dag_id, logical_date, clear_number): return f"{logical_date.isoformat()}.{dag_id}.{clear_number}" - def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): - return f"{logical_date.isoformat()}.{dag_id}.{task_id}.{try_number}.{map_index}" - listener = OpenLineageListener() listener.extractor_manager = mock.Mock() @@ -1477,7 +1594,6 @@ def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): adapter = mock.Mock() adapter.build_dag_run_id.side_effect = mock_dag_id - adapter.build_task_instance_run_id.side_effect = mock_task_id adapter.start_task = mock.Mock() adapter.fail_task = mock.Mock() adapter.complete_task = mock.Mock() @@ -2372,12 +2488,12 @@ def test_on_task_instance_running_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is in the running state. This test ensures that when an Airflow task instance transitions to the running state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_running(None, task_instance) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -2450,7 +2566,7 @@ def test_on_task_instance_failed_correctly_calls_openlineage_adapter_run_id_meth """Tests the OpenLineageListener's response when a task instance is in the failed state. This test ensures that when an Airflow task instance transitions to the failed state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() @@ -2459,7 +2575,7 @@ def test_on_task_instance_failed_correctly_calls_openlineage_adapter_run_id_meth listener.on_task_instance_failed( previous_state=None, task_instance=task_instance, **on_task_failed_kwargs ) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -2474,12 +2590,12 @@ def test_on_task_instance_success_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is in the success state. This test ensures that when an Airflow task instance transitions to the success state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_success(None, task_instance) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -2507,7 +2623,7 @@ def test_listener_on_task_instance_running_do_not_call_adapter_when_disabled_ope listener.on_task_instance_running(None, task_instance) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.start_task.assert_not_called() @@ -2533,7 +2649,7 @@ def test_listener_on_task_instance_failed_do_not_call_adapter_when_disabled_oper ) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.fail_task.assert_not_called() @@ -2555,7 +2671,7 @@ def test_listener_on_task_instance_success_do_not_call_adapter_when_disabled_ope listener.on_task_instance_success(None, task_instance) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.complete_task.assert_not_called() @@ -2566,12 +2682,12 @@ def test_on_task_instance_skipped_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is skipped. This test ensures that when an Airflow task instance is skipped via AirflowSkipException, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_skipped(previous_state=None, task_instance=task_instance) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -2597,7 +2713,7 @@ def test_listener_on_task_instance_skipped_do_not_call_adapter_when_disabled_ope listener.on_task_instance_skipped(previous_state=None, task_instance=task_instance) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.complete_task.assert_not_called() diff --git a/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py new file mode 100644 index 0000000000000..fb4ccc0f30225 --- /dev/null +++ b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py @@ -0,0 +1,267 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tests for ``check_run_id_integrity``, the runId check of the OpenLineage system tests' checking operator.""" + +from __future__ import annotations + +from unittest import mock + +import pytest + +from system.openlineage.operator import OpenLineageTestOperator, check_run_id_integrity + + +def _event( + run_id: str | None, + event_type: str = "START", + *, + job: str = "dag.task", + job_type: str = "TASK", + dag_run_id: str = "run_1", + clear_number: int = 0, + try_number: int = 1, + map_index: int | None = None, + loop_index: int | None = None, + region_id: str | None = None, +) -> dict: + run: dict = {"facets": {}} + if run_id is not None: + run["runId"] = run_id + dag_run = {"run_id": dag_run_id, "clear_number": clear_number} + if job_type == "TASK": + task_instance: dict[str, int | str] = {"try_number": try_number} + if map_index is not None: + task_instance["map_index"] = map_index + if loop_index is not None: + task_instance["loop_index"] = loop_index + if region_id is not None: + task_instance["region_id"] = region_id + run["facets"]["airflow"] = {"dagRun": dag_run, "taskInstance": task_instance} + elif job_type == "DAG": + run["facets"]["airflowDagRun"] = {"dagRun": dag_run} + return { + "eventType": event_type, + "job": {"name": job, "facets": {"jobType": {"jobType": job_type}}}, + "run": run, + } + + +def _dag_event(run_id: str, event_type: str = "START", **kwargs) -> dict: + return _event(run_id, event_type, job="dag", job_type="DAG", **kwargs) + + +def _query_event(run_id: str, event_type: str = "START", **kwargs) -> dict: + return _event(run_id, event_type, job="dag.task.query.1", job_type="QUERY", **kwargs) + + +VALID_SCENARIOS = { + "start_and_complete": [_event("a"), _event("a", "COMPLETE")], + "start_and_fail": [_event("a"), _event("a", "FAIL")], + "retry": [ + _event("a"), + _event("a", "FAIL"), + _event("b", try_number=2), + _event("b", "COMPLETE", try_number=2), + ], + "mapped": [ + _event("a", map_index=0), + _event("a", "COMPLETE", map_index=0), + _event("b", map_index=1), + _event("b", "COMPLETE", map_index=1), + ], + "loop_iterations": [ + event + for index, run_id in enumerate("abc") + for event in (_event(run_id, loop_index=index), _event(run_id, "COMPLETE", loop_index=index)) + ], + "mapped_in_loop": [ + event + for run_id, map_index, loop_index in (("a", 0, 0), ("b", 1, 0), ("c", 0, 1), ("d", 1, 1)) + for event in ( + _event(run_id, map_index=map_index, loop_index=loop_index), + _event(run_id, "COMPLETE", map_index=map_index, loop_index=loop_index), + ) + ], + "loop_iteration_run_again_after_a_clear": [ + _event("a", loop_index=1, region_id="loop_region"), + _event("a", "COMPLETE", loop_index=1, region_id="loop_region"), + _event("b", loop_index=1, region_id="forked_loop_region"), + _event("b", "COMPLETE", loop_index=1, region_id="forked_loop_region"), + ], + "same_task_in_two_dag_runs": [ + _event("a", dag_run_id="run_1"), + _event("a", "COMPLETE", dag_run_id="run_1"), + _event("b", dag_run_id="run_2"), + _event("b", "COMPLETE", dag_run_id="run_2"), + ], + "query_start_and_complete": [_query_event("a"), _query_event("a", "COMPLETE")], + "query_start_and_fail": [_query_event("a"), _query_event("a", "FAIL")], + "dag_run_start_only": [_dag_event("a")], + "dag_run_start_and_complete": [_dag_event("a"), _dag_event("a", "COMPLETE")], + "dag_run_cleared": [ + _dag_event("a"), + _dag_event("a", "FAIL"), + _dag_event("b", clear_number=1), + _dag_event("b", "COMPLETE", clear_number=1), + ], +} + +INVALID_SCENARIOS = { + "start_and_complete_with_different_run_ids": ( + [_event("a"), _event("b", "COMPLETE")], + "has several runIds", + ), + "complete_without_start": ([_event("a", "COMPLETE")], "one START"), + "start_without_terminal": ([_event("a")], "one COMPLETE/FAIL"), + "two_terminal_events": ([_event("a"), _event("a", "FAIL"), _event("a", "COMPLETE")], "one COMPLETE/FAIL"), + "repeated_start_not_declared": ( + [_event("a"), _event("a"), _event("a", "COMPLETE")], + "one START (several only if", + ), + "retry_reusing_run_id": ( + [_event("a"), _event("a", "FAIL"), _event("a", try_number=2), _event("a", "COMPLETE", try_number=2)], + "shared by different attempts", + ), + "map_indexes_sharing_run_id": ( + [ + _event("a", map_index=0), + _event("a", map_index=1), + _event("a", "COMPLETE", map_index=0), + ], + "shared by different attempts", + ), + "loop_iterations_sharing_run_id": ( + [_event("a", loop_index=0), _event("a", "COMPLETE", loop_index=0), _event("a", loop_index=1)], + "shared by different attempts", + ), + "loop_iteration_runs_of_two_trajectories_sharing_run_id": ( + [ + _event("a", loop_index=1, region_id="loop_region"), + _event("a", "COMPLETE", loop_index=1, region_id="loop_region"), + _event("a", loop_index=1, region_id="forked_loop_region"), + ], + "shared by different attempts", + ), + "dag_runs_sharing_run_id": ( + [_event("a", dag_run_id="run_1"), _event("a", "COMPLETE", dag_run_id="run_2")], + "shared by different attempts", + ), + "clear_numbers_sharing_run_id": ( + [_dag_event("a"), _dag_event("a", "COMPLETE", clear_number=1)], + "shared by different attempts", + ), + "dag_run_with_different_run_ids_per_event": ( + [_dag_event("a"), _dag_event("b", "COMPLETE")], + "has several runIds", + ), + "query_start_without_terminal": ([_query_event("a")], "one COMPLETE/FAIL"), + "query_complete_without_start": ([_query_event("a", "COMPLETE")], "one START"), + "query_two_terminal_events": ( + [_query_event("a"), _query_event("a", "COMPLETE"), _query_event("a", "FAIL")], + "one COMPLETE/FAIL", + ), + "dag_run_complete_without_start": ([_dag_event("a", "COMPLETE")], "one START and at most one"), + "dag_run_two_terminal_events": ( + [_dag_event("a"), _dag_event("a", "COMPLETE"), _dag_event("a", "FAIL")], + "one START and at most one", + ), + "jobs_sharing_run_id": ( + [_event("a", job="dag.one"), _event("a", "COMPLETE", job="dag.two")], + "shared by different jobs", + ), + "missing_run_id": ([_event(None)], "has no run.runId"), +} + + +@pytest.mark.parametrize("events", VALID_SCENARIOS.values(), ids=VALID_SCENARIOS.keys()) +def test_valid_events_have_no_problems(events): + assert check_run_id_integrity({"dag.task.event.start": events}) == [] + + +@pytest.mark.parametrize( + ("events", "expected_problem"), INVALID_SCENARIOS.values(), ids=INVALID_SCENARIOS.keys() +) +def test_invalid_events_are_reported(events, expected_problem): + problems = check_run_id_integrity({"dag.task.event.start": events}) + + assert any(expected_problem in problem for problem in problems), problems + + +def test_events_of_one_attempt_are_joined_across_keys(): + events_by_key = { + "dag.task.event.start": [_event("a")], + "dag.task.event.complete": [_event("a", "COMPLETE")], + } + + assert check_run_id_integrity(events_by_key) == [] + + +@pytest.mark.parametrize( + ("repeated_start_keys", "has_problems"), + [((), True), (("dag.task.event.start",), False), (("other.event.start",), True)], +) +def test_repeated_starts_are_allowed_only_for_declared_keys(repeated_start_keys, has_problems): + events_by_key = { + "dag.task.event.start": [_event("a"), _event("a"), _event("a")], + "dag.task.event.complete": [_event("a", "COMPLETE")], + } + + problems = check_run_id_integrity(events_by_key, repeated_start_keys=repeated_start_keys) + + assert bool(problems) is has_problems + + +def test_repeated_starts_are_never_allowed_for_a_dag_run(): + events_by_key = {"dag.event.start": [_dag_event("a"), _dag_event("a")]} + + assert check_run_id_integrity(events_by_key, repeated_start_keys=("dag.event.start",)) + + +@pytest.mark.parametrize( + ("event_count_assertions", "expected"), + [ + ({"dag.task.event.start": ">=2"}, True), + ({"dag.task.event.start": ">=3"}, True), + ({"dag.task.event.start": ">=1"}, False), + ({"dag.task.event.start": "==2"}, False), + ({"dag.task.event.start": 2}, False), + ({".*\\.event\\.start": ">=2"}, True), + ({"dag.task.event.complete": ">=2"}, False), + (None, False), + ], +) +def test_operator_derives_repeated_starts_from_event_count_assertions(event_count_assertions, expected): + operator = OpenLineageTestOperator( + task_id="check_events", event_templates={}, event_count_assertions=event_count_assertions + ) + + assert operator._expects_repeated_events("dag.task.event.start") is expected + + +@mock.patch("system.openlineage.operator.time.sleep") +@mock.patch("system.openlineage.operator.check_run_id_integrity", autospec=True) +def test_operator_skips_the_run_id_check_when_a_template_fails(check_integrity, _sleep): + operator = OpenLineageTestOperator( + task_id="check_events", event_templates={"dag.task.event.start": {}}, clear_variables=False + ) + + with mock.patch.object(operator, "_check_key", side_effect=ValueError("template mismatch")): + with pytest.raises(ValueError, match="template mismatch") as err: + operator.execute(context={}) + + check_integrity.assert_not_called() + assert "runId integrity" not in str(err.value) diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index 1b66b80ab5ad0..03f817745f999 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -29,6 +29,7 @@ import pendulum import pytest from openlineage.client.facet_v2 import parent_run +from sqlalchemy.orm import Session from sqlalchemy.orm.exc import DetachedInstanceError from uuid6 import uuid7 @@ -77,10 +78,12 @@ get_operator_provider_version, get_orm_task_instance_positions, get_parent_information_from_dagrun_conf, - get_regional_task_instance_run_id, get_root_information_from_dagrun_conf, get_runtime_outlet_assets, get_task_documentation, + get_task_instance_map_index, + get_task_instance_position, + get_task_instance_run_id, get_task_parent_run_facet, get_user_provided_run_facets, is_dag_run_asset_triggered, @@ -2011,6 +2014,33 @@ def test_get_task_groups_details(): assert result == expected +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +def test_get_task_groups_details_has_loop_attributes_only_for_a_loop(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("loop_dag", schedule=None, start_date=datetime.datetime(2024, 6, 1)) as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + with TaskGroup("plain"): + EmptyOperator(task_id="work") + + result = _get_task_groups_details(dag, edge_map=_build_labeled_edge_map(dag)) + + loop_attributes = {"max_iterations", "has_until", "terminal_task_id", "gate_task_id"} + assert {k: result["body"][k] for k in loop_attributes} == { + "max_iterations": 3, + "has_until": False, + "terminal_task_id": "body.work", + "gate_task_id": "body.__loop_gate", + } + assert not loop_attributes & set(result["plain"]) + + def test_get_task_groups_details_nested(): with DAG("test_dag", schedule=None, start_date=datetime.datetime(2024, 6, 1)) as dag: with TaskGroup("tg1", prefix_group_id=True) as tg: @@ -3266,10 +3296,13 @@ def test_taskinstance_info_af3(): bundle_instance = MagicMock(version="bundle_version") bundle_instance.name = "bundle_name" runtime_ti.bundle_instance = bundle_instance + runtime_ti.max_tries = 3 assert dict(TaskInstanceInfo(runtime_ti)) == { + "id": str(runtime_ti.id), "log_url": runtime_ti.log_url, "map_index": 2, + "max_tries": 3, "note": None, "rendered_map_index": None, "try_number": 1, @@ -3305,6 +3338,7 @@ def test_taskinstance_info_af2(): "rendered_map_index": None, "pool": "default_pool", "try_number": 0, + "max_tries": 0, "queued_dttm": "2024-06-01T00:00:00+00:00", "log_url": "some_log_url", "dag_bundle_name": None, @@ -3910,7 +3944,7 @@ def test_orm_lineage_facets_project_public_index(dag_maker, session, mapped): @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") @pytest.mark.parametrize("in_loop", [False, True]) -def test_regional_run_id_is_the_execution_uuid_only_inside_a_loop(in_loop): +def test_task_instance_run_id_is_the_attempt_uuid_only_inside_a_loop(in_loop): from airflow.sdk import task_group from airflow.sdk.definitions._internal.loop import create_loop @@ -3924,15 +3958,600 @@ def body(): create_loop(body, max_iterations=2) else: body() - ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), task=dag.get_task("body.work")) + ti = SimpleNamespace( + id=uuid4(), + dag_id="run_id_dag", + task_id="body.work", + try_number=1, + region_id=uuid4(), + map_index=-1, + task=dag.get_task("body.work"), + get_template_context=MagicMock(), + ) + logical_date = timezone.datetime(2024, 1, 1) - assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop else None) + run_id = get_task_instance_run_id(ti, logical_date=logical_date) + + assert run_id == ( + str(ti.id) + if in_loop + else build_task_instance_ol_run_id( + dag_id="run_id_dag", task_id="body.work", try_number=1, logical_date=logical_date, map_index=-1 + ) + ) + assert get_task_instance_map_index(ti) == -1 + ti.get_template_context.assert_not_called() + + +def test_task_instance_run_id_uses_the_map_index_of_the_given_position(): + ti = SimpleNamespace(id=uuid4(), dag_id="dag", task_id="task", try_number=2) + logical_date = timezone.datetime(2024, 1, 1) + + run_id = get_task_instance_run_id( + ti, logical_date=logical_date, position=TaskInstancePosition(in_loop=False, map_index=3) + ) + + assert run_id == build_task_instance_ol_run_id( + dag_id="dag", task_id="task", try_number=2, logical_date=logical_date, map_index=3 + ) + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +@pytest.mark.parametrize("has_loop_context", [True, False]) +def test_runtime_task_instance_position_has_the_iteration_of_its_loop_context(has_loop_context): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("position_dag") as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + loop = SimpleNamespace(index=2, max_iterations=3) if has_loop_context else None + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=-1, + task=dag.get_task("body.work"), + get_template_context=MagicMock(return_value={"loop": loop}), + ) + + assert get_task_instance_position(ti) == TaskInstancePosition( + in_loop=True, + map_index=-1, + loop_index=2 if has_loop_context else None, + loop_max_iterations=3 if has_loop_context else None, + loop_group_id="body", + ) + ti.get_template_context.assert_called_once_with() + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +def test_runtime_task_instance_position_without_its_iteration_when_the_context_fails(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("position_dag") as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=-1, + task=dag.get_task("body.work"), + get_template_context=MagicMock(side_effect=ValueError("invalid params")), + ) + + assert get_task_instance_position(ti) == TaskInstancePosition( + in_loop=True, map_index=-1, loop_group_id="body" + ) + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +def test_stored_mapped_task_in_loop_is_placed_without_iteration_when_its_region_query_fails(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("stored_mapped_loop_dag") as dag: + + @task_group + def body(): + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + session = MagicMock(spec=Session) + session.execute.side_effect = OSError("connection lost") + + positions = get_orm_task_instance_positions([ti], session=session) + + assert positions == { + ti.id: TaskInstancePosition(in_loop=True, map_index=1, loop_max_iterations=3, loop_group_id="body") + } + session.begin_nested.assert_called_once_with() -@pytest.mark.db_test @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +def test_stored_task_instance_without_task_still_raises_when_its_region_query_fails(): + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + ti.task = None + session = MagicMock(spec=Session) + session.execute.side_effect = OSError("connection lost") + + with pytest.raises(OSError, match="connection lost"): + get_orm_task_instance_positions([ti], session=session) + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +def test_task_in_a_task_group_of_unknown_type_is_assumed_in_a_loop(): + task = SimpleNamespace(task_id="work", task_group=SimpleNamespace(group_id="unknown")) + ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), map_index=-1, task=task) + + assert get_task_instance_position(ti) == TaskInstancePosition(in_loop=True, map_index=-1) + + +def test_task_instance_info_uses_a_given_position_as_is(): + ti = SimpleNamespace(id=uuid4(), map_index=-1, get_template_context=MagicMock()) + + info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=True, map_index=-1))) + + assert "loop_index" not in info + ti.get_template_context.assert_not_called() + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +def test_stored_task_instance_without_task_or_region_has_no_loop_index(): + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + ti.dag_id = "dag_id" + ti.task_id = "task_id" + ti.region_id = uuid4() + ti.region_index = 1 + ti.task = None + session = MagicMock(spec=Session) + session.execute.return_value = [] + + assert get_orm_task_instance_positions([ti], session=session) == { + ti.id: TaskInstancePosition(in_loop=True, map_index=-1, loop_index=None) + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") @pytest.mark.parametrize("mapped", [False, True]) -def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag_maker, session, mapped): +@pytest.mark.parametrize("in_loop", [False, True]) +def test_runtime_task_instance_info_has_loop_index_only_inside_a_loop(in_loop, mapped): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("loop_index_dag") as dag: + + @task_group + def body(): + if mapped: + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + else: + EmptyOperator(task_id="work") + + if in_loop: + create_loop(body, max_iterations=3) + else: + body() + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=1 if mapped else -1, + task=dag.get_task("body.work"), + ) + if in_loop: + ti.get_template_context = lambda: {"loop": SimpleNamespace(index=2, max_iterations=3)} + + info = dict(TaskInstanceInfo(ti)) + + assert info["map_index"] == (1 if mapped else None) + if in_loop: + assert (info["loop_index"], info["loop_max_iterations"], info["loop_group_id"]) == (2, 3, "body") + else: + assert "loop_index" not in info + assert "loop_max_iterations" not in info + assert "loop_group_id" not in info + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +@pytest.mark.parametrize( + ("until", "loop_index", "expected", "runs_out"), + [ + pytest.param(None, 1, {"stop": False}, False, id="no_until_continues"), + pytest.param(None, 2, {"stop": True}, False, id="no_until_stops_at_last"), + pytest.param("done", 0, {"stop": True, "reason": "done"}, False, id="until_stops_with_reason"), + pytest.param(False, 2, {"stop": False}, True, id="until_false_at_last"), + ], +) +def test_task_instance_info_has_the_decision_of_a_loop_gate(until, loop_index, expected, runs_out): + from airflow.sdk import LoopDecision + from airflow.sdk.definitions._internal.loop import LoopGateOperator + from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded, execute_loop_gate + + if until is False: + until_callable = lambda: False + elif until is not None: + until_callable = lambda: LoopDecision(stop=True, reason=until) + else: + until_callable = None + gate = LoopGateOperator(task_id="__loop_gate", until=until_callable) + ti = SimpleNamespace(id=uuid7(), map_index=-1) + context = {"loop": SimpleNamespace(index=loop_index, max_iterations=3), "ti": ti} + + if runs_out: + with pytest.raises(LoopMaxIterationsExceeded): + execute_loop_gate(gate, context) + else: + execute_loop_gate(gate, context) + info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1))) + + assert info["loop_decision"] == expected + + +def test_task_instance_info_has_no_loop_decision_for_other_tasks(): + ti = SimpleNamespace(id=uuid7(), map_index=-1) + + info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1))) + + assert "loop_decision" not in info + + +@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Attempt UUIDs require Airflow 3") +@pytest.mark.parametrize("stored", [False, True]) +@pytest.mark.parametrize("in_region", [False, True]) +def test_task_instance_info_has_attempt_id_and_region_only_in_a_region(stored, in_region): + region_id = uuid4() if in_region else UUID(int=0) + if stored: + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + else: + ti = SimpleNamespace(id=uuid7(), map_index=-1) + ti.region_id = region_id + ti.region_index = 2 if in_region else -1 + + info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1))) + + assert info["id"] == str(ti.id) + if in_region: + assert (info["region_id"], info["region_index"]) == (str(region_id), 2) + else: + assert "region_id" not in info + assert "region_index" not in info + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +def test_stored_task_instance_in_loop_has_loop_index_and_max_iterations(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("stored_loop_dag") as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 2 + + assert get_orm_task_instance_positions([ti], session=None) == { + ti.id: TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=2, loop_max_iterations=3, loop_group_id="body" + ) + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +@pytest.mark.parametrize("with_session", [True, False]) +def test_stored_mapped_task_in_loop_gets_its_iteration_from_its_region(with_session): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("stored_mapped_loop_dag") as dag: + + @task_group + def body(): + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + session = MagicMock(spec=Session) + session.execute.return_value = [_region_row(ti.region_id, "body.work", parent=(uuid4(), 2, "body"))] + + positions = get_orm_task_instance_positions([ti], session=session if with_session else None) + + assert positions == { + ti.id: TaskInstancePosition( + in_loop=True, + map_index=1, + loop_index=2 if with_session else None, + loop_max_iterations=3, + loop_group_id="body", + ) + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +def test_task_instance_info_of_stored_instance_without_task_does_not_raise(): + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + ti.dag_id = "dag_id" + ti.task_id = "task_id" + ti.region_id = uuid4() + ti.region_index = 2 + ti.task = None + + info = dict(TaskInstanceInfo(ti)) + + assert info["map_index"] is None + assert "loop_index" not in info + + +def test_task_instance_info_uses_the_position_it_is_given(): + ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), map_index=-1) + position = TaskInstancePosition( + in_loop=True, map_index=1, loop_index=2, loop_max_iterations=3, loop_group_id="body" + ) + + info = dict(TaskInstanceInfo(ti, position=position)) + + assert (info["map_index"], info["loop_index"], info["loop_max_iterations"], info["loop_group_id"]) == ( + 1, + 2, + 3, + "body", + ) + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +def test_task_group_info_has_loop_attributes_only_for_a_loop(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("loop_group_dag"): + + @task_group + def body(): + EmptyOperator(task_id="work") + + loop = create_loop(body, max_iterations=3) + + @task_group + def plain(): + EmptyOperator(task_id="work") + + plain_group = plain() + + loop_info = dict(TaskGroupInfo(loop)) + assert loop_info["max_iterations"] == 3 + assert loop_info["has_until"] is False + assert loop_info["terminal_task_id"] == "body.work" + assert loop_info["gate_task_id"] == "body.__loop_gate" + assert not {"max_iterations", "has_until", "terminal_task_id", "gate_task_id"} & set( + TaskGroupInfo(plain_group) + ) + + +def _region_row(region_id, node_id, *, parent=None, grandparent_region_id=None): + """Build a row of the region query: ``parent`` is ``(parent_region_id, parent_region_index, parent_node_id)``.""" + parent_region_id, parent_region_index, parent_node_id = parent or (None, None, None) + return SimpleNamespace( + id=region_id, + node_id=node_id, + parent_region_id=parent_region_id, + parent_region_index=parent_region_index, + parent_node_id=parent_node_id, + grandparent_region_id=grandparent_region_id, + ) + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +@pytest.mark.parametrize( + ("own_region", "parent", "grandparent_region_id", "map_index"), + [ + pytest.param(True, (uuid4(), 2, "outer"), uuid4(), 1, id="mapped_task_two_levels_deep"), + pytest.param(False, (uuid4(), 2, "outer"), None, -1, id="enclosing_region_not_top_level"), + pytest.param(False, None, None, -1, id="region_missing"), + ], +) +def test_stored_task_instance_in_an_unexpected_region_shape_is_placed_without_iteration( + own_region, parent, grandparent_region_id, map_index +): + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + ti.task_id = "work" + ti.region_id = uuid4() + ti.region_index = 1 + ti.task = None + session = MagicMock(spec=Session) + rows = [] + if parent is not None or own_region: + rows = [ + _region_row( + ti.region_id, + "work" if own_region else "group", + parent=parent, + grandparent_region_id=grandparent_region_id, + ) + ] + session.execute.return_value = rows + + assert get_orm_task_instance_positions([ti], session=session) == { + ti.id: TaskInstancePosition(in_loop=True, map_index=map_index) + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +@pytest.mark.parametrize( + ("parent", "grandparent_region_id"), + [ + pytest.param((uuid4(), 2, "another_loop"), None, id="parent_is_not_its_loop"), + pytest.param((uuid4(), 2, "body"), uuid4(), id="loop_region_not_top_level"), + ], +) +def test_stored_mapped_task_in_loop_keeps_no_iteration_from_an_unexpected_region_shape( + parent, grandparent_region_id +): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("stored_mapped_loop_dag") as dag: + + @task_group + def body(): + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + session = MagicMock(spec=Session) + session.execute.return_value = [ + _region_row(ti.region_id, "body.work", parent=parent, grandparent_region_id=grandparent_region_id) + ] + + assert get_orm_task_instance_positions([ti], session=session) == { + ti.id: TaskInstancePosition(in_loop=True, map_index=1, loop_max_iterations=3, loop_group_id="body") + } + + +@pytest.mark.db_test +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +@pytest.mark.parametrize("task_attached", [False, True]) +def test_region_shapes_of_airflow_are_the_ones_openlineage_places(dag_maker, session, task_attached): + """ + Pin the regions Airflow creates, which ``get_orm_task_instance_positions`` relies on to place task instances. + + If this fails, Airflow changed which regions it creates or how they nest: update the region parsing in + ``get_orm_task_instance_positions`` and its comment before changing these expectations. + """ + from airflow.sdk import task, task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with dag_maker(serialized=True) as dag: + + @task_group + def body(): + plain = EmptyOperator(task_id="plain") + mapped = PythonOperator.partial(task_id="mapped", python_callable=list).expand(op_kwargs=[{}, {}]) + + @task_group + def mapped_group(x): + @task + def member(x): + return x + + member(x) + + last = EmptyOperator(task_id="last") + plain >> mapped >> mapped_group.expand(x=[1, 2]) >> last + + create_loop(body, max_iterations=2) + + @task_group + def top_group(x): + @task + def top_member(x): + return x + + top_member(x) + + top_group.expand(x=[1, 2]) + run = dag_maker.create_dagrun() + tis = [ti for ti in run.task_instances if ti.task_id != "body.__loop_gate"] + for ti in tis: + ti.task = dag.get_task(ti.task_id) if task_attached else None + in_loop_extra = {"loop_max_iterations": 2} if task_attached else {} + + positions = get_orm_task_instance_positions(tis, session=session) + + assert {(ti.task_id, ti.region_index): positions[ti.id] for ti in tis} == { + ("body.plain", 0): TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=0, loop_group_id="body", **in_loop_extra + ), + ("body.last", 0): TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=0, loop_group_id="body", **in_loop_extra + ), + **{ + (task_id, map_index): TaskInstancePosition( + in_loop=True, map_index=map_index, loop_index=0, loop_group_id="body", **in_loop_extra + ) + for task_id in ("body.mapped", "body.mapped_group.member") + for map_index in (0, 1) + }, + **{ + ("top_group.top_member", map_index): TaskInstancePosition(in_loop=False, map_index=map_index) + for map_index in (0, 1) + }, + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +@pytest.mark.parametrize("shape", ["nested_loop", "loop_in_mapped_group", "expand_in_mapped_group"]) +def test_region_shapes_stay_limited_by_the_task_sdk(shape): + """ + Pin the Dag shapes the task SDK rejects, which ``get_orm_task_instance_positions`` relies on. + + If one of them stops raising, a task can sit two region levels deep, which OpenLineage places as a loop + without its iteration: support the new shape in ``get_orm_task_instance_positions``, then update this test. + """ + from airflow.sdk import task, task_group + from airflow.sdk.definitions._internal.loop import create_loop + + def define_dag(): + with DAG(f"{shape}_dag"): + + @task_group + def inner(): + EmptyOperator(task_id="work") + + if shape == "nested_loop": + + @task_group + def outer(): + create_loop(inner, max_iterations=2) + + create_loop(outer, max_iterations=2) + elif shape == "loop_in_mapped_group": + + @task_group + def mapped_outer(x): + create_loop(inner, max_iterations=2) + + mapped_outer.expand(x=[1, 2]) + else: + + @task_group + def mapped_outer(x): + @task + def member(y): + return y + + member.expand(y=[x]) + + mapped_outer.expand(x=[1, 2]) + + with pytest.raises(NotImplementedError): + define_dag() + + +@pytest.mark.db_test +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +@pytest.mark.parametrize("kind", ["plain_loop", "mapped", "mapped_in_loop"]) +def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag_maker, session, kind): + mapped = kind != "plain_loop" with dag_maker(serialized=True): if mapped: PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) @@ -3940,22 +4559,34 @@ def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag EmptyOperator(task_id="work") run = dag_maker.create_dagrun() ti = max(run.task_instances, key=lambda ti: ti.region_index) - if not mapped: + if kind == "plain_loop": region = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id, node_id="loop") session.add(region) session.flush() ti.region_id = region.id ti.region_index = 2 + elif kind == "mapped_in_loop": + loop_region = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id, node_id="loop") + session.add(loop_region) + session.flush() + mapped_region = session.get(DynamicRegion, ti.region_id) + mapped_region.parent_region_id = loop_region.id + mapped_region.parent_region_index = 2 session.flush() ti.task = None with pytest.raises(ValueError, match="has no task attached"): - get_regional_task_instance_run_id(ti) + get_task_instance_run_id(ti, logical_date=timezone.datetime(2024, 1, 1)) with assert_queries_count(1): positions = get_orm_task_instance_positions([ti], session=session) assert positions == { - ti.id: TaskInstancePosition(in_loop=not mapped, map_index=ti.region_index if mapped else -1) + ti.id: TaskInstancePosition( + in_loop=kind != "mapped", + map_index=ti.region_index if mapped else -1, + loop_index=2 if kind != "mapped" else None, + loop_group_id="loop" if kind != "mapped" else None, + ) }