fixup! Resolve live producers by region and read their XCom by attempt

Remove the loaded_as_persistent listener once the lookup test is done.

The session fixture yields the thread-local Session that the scoped session
hands out to every caller, so event.listen(session, ...) registers on an
object that outlives the test and the listener was never removed. Its
closure kept every instance loaded by later tests alive in a list. The
identity map only holds weak references, and Variable.update writes with a
Core upsert that does not expire the Variable it just loaded, while
XComModel.set only expires XComModelV2 instances and not XComModel ones. A
row loaded earlier in the same test was therefore returned stale ('value1'
instead of 'value2') by the next read. test_variable_update and
test_xcom_set_again_replace passed alone and failed whenever this test ran
first on SQLite.

Register a named listener and remove it in a finally block, as the other
loaded_as_persistent tests in the stack already do.

Author: ashb

airflow-core/src/airflow/models/dagrun.py

@@ -1738,8 +1738,8 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
 
         # Check dependencies.
         expansion_happened = False
-        # Set of task ids for which was already done _revise_map_indexes_if_mapped
-        revised_map_index_task_ids: set[tuple[str, UUID]] = set()
+        # (task id, region id) pairs for which _revise_map_indexes_if_mapped was already done
+        revised_expansion_keys: set[tuple[str, UUID]] = set()
         for schedulable in itertools.chain(schedulable_tis, additional_tis):
             if TYPE_CHECKING:
                 assert isinstance(schedulable.task, Operator)
@@ -1764,18 +1764,18 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
                     # later in this same pass must see the post-expansion count).
                     dep_context.invalidate_upstream_task_id_counts()
             if new_tis is None and schedulable.state in SCHEDULEABLE_STATES:
-                # It's enough to revise map index once per task id,
+                # It's enough to revise map index once per task and region,
                 # checking the map index for each mapped task significantly slows down scheduling
                 expansion_key = (schedulable.task.task_id, schedulable.region_id)
-                if expansion_key not in revised_map_index_task_ids:
+                if expansion_key not in revised_expansion_keys:
                     revised_tis = self._revise_map_indexes_if_mapped(
                         schedulable.task,
                         dag_version_id=schedulable.dag_version_id,
                         region_id=schedulable.region_id,
                         session=session,
                     )
                     ready_tis.extend(revised_tis)
-                    revised_map_index_task_ids.add(expansion_key)
+                    revised_expansion_keys.add(expansion_key)
                     if revised_tis:
                         # Revising a mapped task can add new instances, growing its instance count
                         # the same way expansion does. Drop the upstream-count memo so a downstream

airflow-core/src/airflow/models/xcom.py

@@ -357,6 +357,7 @@ def get_many(
             if (
                 any(value is not None for value in (task_ids, dag_ids, map_indexes))
                 or include_prior_dates
+                or try_number is not None
                 or region_id != SENTINEL_REGION_ID
             ):
                 raise ValueError("producer_ids cannot be combined with coordinate filters")

airflow-core/tests/unit/models/test_dynamic_region.py

@@ -251,19 +251,71 @@ def test_loop_iteration_lookup_loads_only_the_selected_iterations_rows(producer_
     expected_ids = {ti.id for ti in expected}
     session.expunge_all()
     loaded = []
-    event.listen(session, "loaded_as_persistent", lambda _, instance: loaded.append(instance))
-    selected = resolve_current_producers(
-        dag_id=mapped.dag_id,
-        run_id=mapped.run_id,
-        task_id="mapped",
-        is_mapped=True,
-        context=ProducerContext(consumer.region_id, consumer.region_index, "loop"),
-        session=session,
-    )
+
+    def record_loaded(_, instance):
+        loaded.append(instance)
+
+    event.listen(session, "loaded_as_persistent", record_loaded)
+    try:
+        selected = resolve_current_producers(
+            dag_id=mapped.dag_id,
+            run_id=mapped.run_id,
+            task_id="mapped",
+            is_mapped=True,
+            context=ProducerContext(consumer.region_id, consumer.region_index, "loop"),
+            session=session,
+        )
+    finally:
+        event.remove(session, "loaded_as_persistent", record_loaded)
     assert {ti.id for ti in selected} == expected_ids
     assert {instance.id for instance in loaded if isinstance(instance, TaskInstance)} == expected_ids
 
 
+def test_loop_iteration_lookup_ignores_another_loops_nested_regions(producer_tis, session):
+    tis, regions, _ = producer_tis
+    consumer, mapped = tis["consumer"], tis["mapped"]
+    other_loop = DynamicRegion(dag_id=mapped.dag_id, run_id=mapped.run_id, node_id="other_loop")
+    session.add(other_loop)
+    session.flush()
+    placed = {}
+    for owner, parent in (("own", regions[0]), ("other", other_loop)):
+        nested = DynamicRegion(
+            dag_id=mapped.dag_id,
+            run_id=mapped.run_id,
+            node_id="mapped",
+            parent_region_id=parent.id,
+            parent_region_index=consumer.region_index,
+        )
+        session.add(nested)
+        session.flush()
+        placed[owner] = TaskInstance(
+            mapped.task, mapped.dag_version_id, run_id=mapped.run_id, map_index=0, region_id=nested.id
+        )
+        session.add(placed[owner])
+    session.flush()
+    own_id, other_id = placed["own"].id, placed["other"].id
+    session.expunge_all()
+    loaded = []
+
+    def record_loaded(_, instance):
+        loaded.append(instance)
+
+    event.listen(session, "loaded_as_persistent", record_loaded)
+    try:
+        selected = resolve_current_producers(
+            dag_id=mapped.dag_id,
+            run_id=mapped.run_id,
+            task_id="mapped",
+            is_mapped=True,
+            context=ProducerContext(consumer.region_id, consumer.region_index, "loop"),
+            session=session,
+        )
+    finally:
+        event.remove(session, "loaded_as_persistent", record_loaded)
+    assert [ti.id for ti in selected] == [own_id]
+    assert other_id not in {instance.id for instance in loaded if isinstance(instance, TaskInstance)}
+
+
 def test_previous_iteration_zero_is_missing(producer_tis, session):
     tis, _, _ = producer_tis
     consumer = tis["consumer"]
@@ -404,6 +456,7 @@ def test_get_many_rejects_region_filter_with_prior_dates(regional_tis):
         {"region_id": None},
         {"region_id": uuid4()},
         {"include_prior_dates": True},
+        {"try_number": 2},
     ],
 )
 def test_get_many_rejects_producer_ids_with_coordinate_filters(coordinate_filter):

airflow-core/tests/unit/models/test_xcom_arg.py

@@ -25,6 +25,7 @@
 from airflow.models.xcom_arg import XComArg
 from airflow.providers.standard.operators.bash import BashOperator
 from airflow.providers.standard.operators.python import PythonOperator
+from airflow.sdk import task_group
 from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
 from airflow.serialization.definitions.notset import NOTSET
 from airflow.serialization.definitions.xcom_arg import (
@@ -420,7 +421,6 @@ def source(value):
 
 
 def test_member_of_mapped_task_group_has_no_map_length(dag_maker, session):
-    from airflow.sdk import task_group
 
     with dag_maker(session=session, serialized=True):