kaxil commented on code in PR #74892:
URL: https://github.com/apache/airflow/pull/74892#discussion_r4299785512
##########
airflow-core/src/airflow/models/dynamic_region.py:
##########
@@ -404,6 +446,70 @@ def _set_slot_key(mapper: Mapper, connection: Connection,
region: DynamicRegion)
)
+def _get_coordinates_cache(region: DynamicRegion) -> dict[UUID, tuple[UUID,
bytes]]:
+ """
+ Remember each region's coordinates until the transaction ends, so siblings
cost one read of their parent.
Review Comment:
`after_transaction_end` also fires for the subtransaction every flush opens
and for each `begin_nested()`, so this cache is cleared at the end of every
flush, not the transaction. Behaviour is fine (a `get_or_create_many` batch
shares one flush, and clearing on savepoint end keeps a rolled-back region
out), so maybe just say "until the flush or savepoint ends" here.
`test_coordinates_of_a_rolled_back_region_are_forgotten` and
`test_coordinates_are_forgotten_when_the_transaction_commits` have the same
mismatch: `add_region` has already flushed, so the cache is gone before the
commit or rollback they're named for.
##########
airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py:
##########
@@ -157,6 +157,10 @@ def upgrade():
sa.Column("parent_region_id", CompactUUID(), nullable=True),
sa.Column("parent_region_index", sa.Integer(), nullable=True),
sa.Column("forked_from_region_id", CompactUUID(), nullable=True),
+ sa.Column("family_root_id", CompactUUID(), nullable=False),
Review Comment:
A 3.4.0b1 database is already stamped past this revision (b1 shipped 0143 to
0146), so `airflow db migrate` won't add these columns and nothing tells the
user. The first sign is a missing `family_root_id` error from
`DynamicRegion.get_or_create_many`, which `_create_tasks` calls for every
mapped task, so a b1 user who never wrote a loop hits it too. Could we add a
check to `_check_migration_errors` in `utils/db.py` that stops with a "migrated
by 3.4.0b1, downgrade or reset" message when `dynamic_region` has no
`family_root_id`? `upgradedb` already exits on those before running any
migration.
##########
airflow-core/src/airflow/models/dynamic_region.py:
##########
@@ -457,59 +563,37 @@ def loop_position(
return None
[email protected](frozen=True)
-class _LoopPassRegions:
- """The regions that hold one iteration of a loop: its fork family and
everything nested under it."""
-
- family_ids: tuple[UUID, ...]
- nested_ids: tuple[UUID, ...]
- iteration: int
-
-
-def _load_loop_pass_regions(position: tuple[UUID, int], *, session: Session)
-> _LoopPassRegions:
+def build_loop_family_legs(
+ query: Select[Any], family_member: DynamicRegion, *, iteration: int | None
= None
+) -> tuple[Select[Any], Select[Any]]:
"""
- Resolve, in one statement, the regions that hold the loop iteration
``position`` names.
-
- Both recursive CTEs run against ``dynamic_region`` only. The
``task_instance`` query then takes the
- ids as literals, which every backend turns into an index lookup on
``task_instance_current_key``;
- a subquery in that position leaves MySQL and Postgres scanning every row
of the task.
+ Split ``query``, a selection from ``task_instance``, into the rows that
belong to one loop family.
+
+ ``family_member`` is any region of the loop's family; the family's regions
share ``family_root_id`` and
+ ``path``. The first leg holds the rows that sit directly in a region of
the family, where
+ ``region_index`` is the iteration. The second holds the rows nested under
the family, however deep,
+ whose ``path`` starts with the family and the iteration. Pass
``iteration`` to keep one iteration, or
+ leave it out for all of them.
+
+ Each leg is joined to :class:`DynamicRegion`, so a caller can add columns
to both (the iteration of
+ a nested row is ``DynamicRegion.parent_region_index``) before combining
them with ``UNION ALL``. The
Review Comment:
For a row two or more levels under the family, the joined region's
`parent_region_index` is its index inside the intermediate region, not the loop
iteration. The `SUBSTRING(path, LENGTH(m.path) + 17, 4)` expression in the
class docstring is the right one there. Worth saying "for a row directly under
the family" or pointing at that instead?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
