Improve task loop information and tests for OpenLineage provider

[kacpermuda]

Task loops are new, and OpenLineage events need to describe them from day one. Inside a loop the iteration of a
task was invisible: taskInstance.map_index is only set for mapped tasks and the runId is the attempt UUID, so
iterations of one task looked identical, and nothing verified that runIds stay consistent across START and
COMPLETE, retries, map indexes and iterations.

• taskInstance facet gets loop_index and loop_max_iterations. Loop groups expose max_iterations,
  has_until, gate_task_id and terminal_task_id in task_group and in AirflowJobFacet.taskGroups.
• The manual state change path passes its computed position to the facet, and TaskInstanceInfo no longer
  raises for a stored task instance it can't place.
• OpenLineageTestOperator checks runId integrity by default: one job and attempt per runId, a complete
  START and COMPLETE/FAIL lifecycle, and one runId per attempt.
• New example Dags for retries, a rescheduled sensor, and loops with a plain task, a retried loop task and a
  retried mapped loop task (Airflow 3.4+).

---

Was generative AI tooling used to co-author this PR?

• Yes (Claude Code)

Generated-by: Claude Code (Sonnet 5.5) following the guidelines

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[ashb]

Just having loop_index will mean this needs to change if/when nested loops are allowed in Airflow.

It also doesn't support mapped tasks inside a loop, which is supported today. This is why Airflow has a "dynamic_region" concept.

[kacpermuda]

Thanks, extended this with region_id. Also looking into exposing a bit more data from core/sdk so that OL can show it.

[github-actions]

uv.lock on main just moved via #74290 ("Add AzureAISearchHook to the Microsoft Azure provider"), commit 935bac0 and this PR currently conflicts.

Quickest fix:

git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

[mobuchowski]

Unrecognized group types now count as "not in a loop" instead of "in a loop"

This function walks up the task's parent groups looking for a loop group. The walk stops as soon as it meets something that isn't a TaskGroup / SerializedTaskGroup. The old _is_task_in_loop then returned group is not None: if the walk stopped on some unexpected non-None object, it assumed the task was in a loop. The new code returns None in that case, so it now assumes the task is not in a loop.

Why the direction of this fallback matters:

• When a task is in a loop, the listener uses the task instance's UUID (ti.id) as the OpenLineage runId. Each iteration gets a new UUID, so iterations stay separate.
• When it's not in a loop, the runId is built from (dag_id, task_id, try_number, logical_date, map_index). Call that the "computed runId": the same inputs always give the same runId.

Every iteration of a looped task has the same values for all five inputs. If loop detection wrongly says "not in a loop", iteration 0 and iteration 1 emit START/COMPLETE under one runId, and consumers merge their lineage into a single run.

I'm not sure which group types can reach this point in practice (maybe a group from a different SDK version, a proxy, or a subclass that doesn't inherit from these two classes). But if one does, the old fallback gave a harmless extra UUID while the new one silently merges runs. Keeping return group is not None, or at least logging when the walk stops on an unexpected type, seems safer.

[kacpermuda]

Good catch, fixed in 7441d2d: _find_enclosing_loop assumes a loop again when the walk stops on a group of an unknown type, and logs a warning, so the worst case is an extra runId rather than merged iterations.

[mobuchowski]

A new database query in the manual state change path can make the whole event disappear

Before this PR, a stored task instance whose task object was attached got its position computed purely in Python, with no query. This line adds a second group, without_iteration: mapped tasks inside a loop. Their iteration number only lives on the DynamicRegion row (the record Airflow keeps for each loop iteration or mapped expansion), so this now queries the database for them.

That code runs from _on_task_instance_manual_state_change, the handler for when the scheduler fails a task (e.g. handle_failure after a heartbeat timeout) or a user marks a task failed or successful in the UI. So this query runs inside the scheduler/API server process.

The scenario I'm worried about:

1. A mapped task inside a loop gets marked failed by the scheduler.
2. get_orm_task_instance_positions now runs the DynamicRegion query for it.
3. The query raises (connection hiccup, the TI's session already detached or closed, etc.).
4. The broad try/except in _on_task_instance_manual_state_change catches the error and returns, so no FAIL event is emitted for that task at all.

All of this is only to fill loop_index, which is descriptive: the runId doesn't depend on it. _add_runtime_loop_iteration states this exact rule in its comment ("failing to read it must not cost the event its run id"). Could the without_iteration lookup get its own try/except, so a failure only leaves loop_index unset?

[kacpermuda]

Fixed in 7441d2d: when the query only fills in mapped tasks' iterations, it runs in a savepoint, and a failure only leaves loop_index unset and logs a warning. The event still goes out and the caller's transaction stays usable. Task instances without their task still raise as before this PR.

[mobuchowski]

Tasks in a mapped task group (with no loop at all) get reported as loop iterations

This else branch handles a stored task instance without its task object attached. That happens for the asset-event source and for the manual state change path. Here the code reads the task's DynamicRegion row and assumes that if the row's node_id isn't the task itself, the row must belong to an enclosing loop.

But DynamicRegion rows are also created for mapped task groups (@task_group + .expand()), not just for loops. Example:

@task_group
def g(x):
    t(x)

g.expand(x=[1, 2, 3])  # no loop anywhere

For task g.t, map index 1:

1. Its region row has node_id = "g", which isn't "g.t", so we land in this else.
2. We set in_loop=True, map_index=-1, loop_index=1, loop_group_id="g".

So the event says "iteration 1 of loop g" for a task that isn't in a loop, drops its real map_index (1), and uses the UUID runId. The same task's START/COMPLETE from the worker uses the computed runId (built from dag/task/try/date/map_index), so the manual-state FAIL and the worker's START end up with different runIds.

Checking the region's kind (loop vs. mapped group) before treating it as a loop would avoid this, if DynamicRegion records that.

[kacpermuda]

I checked this against core: it never creates a region for a mapped task group. Every task that needs expansion, a task in a mapped task group included, gets its own region with node_id == task_id (core's log addressing classifies regions with the same rule), so g.t takes the first branch with its map index. 7441d2d makes the parsing check the shapes it expects (a loop's region is top-level, a mapped task's region sits directly in its loop's) and places anything else as a loop without its iteration, with a warning. test_region_shapes_of_airflow_are_the_ones_openlineage_places builds a real Dag run with these shapes, so a change in core fails the provider's tests.

[mobuchowski]

parent_region_index is the loop iteration only if the loop is the direct parent

For a mapped task in a loop, this takes region.parent_region_index: the index of the mapped task's region inside its parent region. That equals the loop iteration only when the parent region is the loop itself.

If there's another mapped level in between, for example:

loop L
 └── mapped task group G (.expand)
      └── mapped task t (.expand)

then t's region's parent is G's region, so parent_region_index is G's map index, not L's iteration. In iteration 0 of the loop, t could report loop_index=1 just because it sits in G's second expansion.

I haven't checked whether Airflow allows this nesting. If it doesn't, a short comment stating that assumption would help. If it does, this probably needs to walk up the parent regions until it reaches the loop's region.

[kacpermuda]

Airflow doesn't allow that nesting today: the SDK rejects expanding a task inside a mapped task group, and inside a loop each task of a mapped task group gets its region directly under the loop's. 7441d2d states this, uses parent_region_index only when the parent region is the task's own loop and top-level, and pins the SDK restrictions in test_region_shapes_stay_limited_by_the_task_sdk.

[ashb]

Oh, Brent is changing this to be 1-based, not 0-based!

[ashb]

In #74773

[kacpermuda]

Yeah, that's a good call, was confused in the UI. Will adjust the test (here or in the followup PR if there are no other nits in this PR)