From cdda9db06f09c96a83820ac0f66ca5c208804f9b Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Fri, 2 Oct 2026 18:34:32 +0100 Subject: [PATCH] Introduce the (internal/db level) concept of "dynamic regions" for TIs As mentioned in the first docs PR of this stack, and the precursor PR #74222, we need a way of distinguishing between "interation 2" and "iteration 2 after we already cleared and re-ran loop 2 but-these-are-totally-different-TIs, not retries.", or more simply in the digram as `2'` (said as "two prime"). Task loops need the same task_id to exist once per pass, and a clear needs to put replacement task beside the work it has "retired" or superceded. A task instance is identified by dag_id, task_id, run_id and map_index today, so neither is possible: the second row would collide with the first as they both have the same task_id and try_number. A "dynamic region is one execution of a construct that creates task instances at run time: a loop, or the expansion of a mapped task or task group. Each region is a row recording which Dag run it belongs to, which node it executes (node_id, the id of the task group or mapped task), which enclosing region and position contains it and (so that we can have mapped tasks inside loops!), when a clear replaced an earlier execution, which region it was forked from and where it resumes (i.e. "clear iteration 2 onwards"). A region is a row rather than a position (i.e. it's own table, not just a colum on TI) because a position can be reused/"overwritten" after a clear, and would then stop saying which precise TI created the data. Regions are immutable for the same reason; everything that changes stays on the task instances. This change only adds the storage to hopefully make PR creation easier. Nothing creates a region yet, and existing rows take the all-zero sentinel, meaning outside any region, because rewriting existing rows on upgrade was measured at production scale and rejected (4+ hrs! But I never let it run to completion. Even 30mins is too long). XCom and rendered fields are keyed by attempt UUID since the precursor ownership PR and so don't need a region column. `region_index` is an ORM synonym of `map_index` for now. Later changes will build upon this to make mapped expansions and loops create regions, resolve and write task data by (region_id, region_index), add the task loop gate that decides whether a loop runs another pass/iteration, and let a clear replace only the selected work by forking a region instead of rewriting history. Downgrade refuses when rows differ only by region, because dropping the column would otherwise merge distinct executions. The check scans task_instance, so expect it to be slow on very large installations. --- .../task-and-asset-state-store.rst | 4 +- .../task-and-asset-state-store.rst | 2 +- .../docs/core-concepts/task-state-store.rst | 4 +- airflow-core/docs/migrations-ref.rst | 4 +- airflow-core/src/airflow/migrations/utils.py | 72 ++++++ ...0142_3_4_0_unify_task_attempt_ownership.py | 26 +- .../0143_3_4_0_add_dynamic_region_storage.py | 211 ++++++++++++++++ airflow-core/src/airflow/models/__init__.py | 1 + airflow-core/src/airflow/models/dagrun.py | 4 + .../src/airflow/models/dynamic_region.py | 80 ++++++ .../src/airflow/models/task_state_store.py | 19 +- .../src/airflow/models/taskinstance.py | 92 ++++++- airflow-core/src/airflow/state/metastore.py | 12 +- airflow-core/src/airflow/utils/db.py | 5 +- airflow-core/src/airflow/utils/db_cleanup.py | 8 +- airflow-core/src/airflow/utils/sqlalchemy.py | 57 ++++- .../test_0143_add_dynamic_region_storage.py | 231 ++++++++++++++++++ .../tests/unit/models/test_taskinstance.py | 1 + .../tests/unit/state/test_metastore.py | 61 +++++ airflow-core/tests/unit/utils/test_db.py | 20 ++ .../tests/unit/utils/test_db_cleanup.py | 41 ++++ .../tests/unit/utils/test_sqlalchemy.py | 62 ++++- .../common/io/docs/state_store_backend.rst | 4 +- providers/common/io/provider.yaml | 4 +- .../providers/common/io/get_provider_info.py | 2 +- .../common/io/state_store/backend.py | 32 +-- .../providers/common/io/version_compat.py | 1 + .../common/io/state_store/test_backend.py | 49 +++- .../src/airflow_shared/state/__init__.py | 17 +- shared/state/tests/state/test_state.py | 15 ++ 30 files changed, 1060 insertions(+), 81 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py create mode 100644 airflow-core/src/airflow/models/dynamic_region.py create mode 100644 airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py diff --git a/airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst b/airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst index d4a6859129417..9b13da2abc114 100644 --- a/airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst +++ b/airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst @@ -203,7 +203,9 @@ Example skeleton: class S3StateBackend(BaseStoreBackend): def _task_ref(self, scope: TaskScope, key: str) -> str: - return f"airflow/task-store/{scope.dag_id}/{scope.run_id}/{scope.task_id}/{scope.map_index}/{key}" + # Two regions can hold the same task_id and map_index, so region_id must be part of the ref. + region = f"region={scope.region_id}/" if scope.region_id.int else "" + return f"airflow/task-store/{scope.dag_id}/{scope.run_id}/{region}{scope.task_id}/{scope.map_index}/{key}" def _asset_ref(self, scope: AssetScope, key: str) -> str: import hashlib diff --git a/airflow-core/docs/core-concepts/task-and-asset-state-store.rst b/airflow-core/docs/core-concepts/task-and-asset-state-store.rst index 098cffe7da468..264531932dfae 100644 --- a/airflow-core/docs/core-concepts/task-and-asset-state-store.rst +++ b/airflow-core/docs/core-concepts/task-and-asset-state-store.rst @@ -38,7 +38,7 @@ Task and Asset state store provide two key/value stores to persist data like a j - Default lifetime - Primary use case * - **Task state store** - - A single task Instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``) + - A single task Instance (``dag_id`` + ``run_id`` + ``task_id`` + ``region_id`` + ``map_index``) - Configurable retention; cleared on task success when ``clear_on_success = True`` - Survive retries, track in-flight jobs, checkpoint progress within a run, resume progress from checkpoint set by a past run * - **Asset state store** diff --git a/airflow-core/docs/core-concepts/task-state-store.rst b/airflow-core/docs/core-concepts/task-state-store.rst index c972fd7861ca2..7b181439a455d 100644 --- a/airflow-core/docs/core-concepts/task-state-store.rst +++ b/airflow-core/docs/core-concepts/task-state-store.rst @@ -30,7 +30,9 @@ Task State Store .. versionadded:: 3.3 -Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. +Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``region_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. + +``region_id`` is the all-zero UUID for task instances outside any dynamic region. Task instances created by dynamically expanded work, such as a loop, belong to a region, and two of them can share the same ``dag_id``, ``run_id``, ``task_id`` and ``map_index``. A custom backend must therefore include ``region_id`` in its uniqueness key. Because it outlives a worker crash and stays readable by the next attempt, the task state store is the mechanism behind :ref:`durable execution `, where a task continues from where it stopped instead of repeating work or submitting a duplicate external job. Provider operators that advertise a ``durable`` parameter are built on the API described here. diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 513d0089dcbac..aa991bfc2b7fa 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``e7c2a91bd540`` (head) | ``90e4d18ccadf`` | ``3.4.0`` | Unify task attempt ownership without rewriting legacy XCom | +| ``54a27b6f9d01`` (head) | ``e7c2a91bd540`` | ``3.4.0`` | Add dynamic region storage. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``e7c2a91bd540`` | ``90e4d18ccadf`` | ``3.4.0`` | Unify task attempt ownership without rewriting legacy XCom | | | | | data. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``90e4d18ccadf`` | ``e5a91c7f42b3`` | ``3.4.0`` | Add timetable_asset_gated to DagModel. | diff --git a/airflow-core/src/airflow/migrations/utils.py b/airflow-core/src/airflow/migrations/utils.py index 4eeaf373c6a87..f61cd4e6c5814 100644 --- a/airflow-core/src/airflow/migrations/utils.py +++ b/airflow-core/src/airflow/migrations/utils.py @@ -18,6 +18,7 @@ import contextlib from contextlib import contextmanager +from textwrap import dedent @contextmanager @@ -30,6 +31,25 @@ def disable_sqlite_fkeys(op): yield op +# Unlike disable_sqlite_fkeys, a failed SQLite upgrade rolls back atomically and foreign_keys is always restored. +@contextmanager +def sqlite_rebuilds(op): + if op.get_bind().dialect.name != "sqlite": + yield + return + if op.get_context().as_sql: + raise RuntimeError("SQLite offline SQL cannot render this migration's table rebuilds") + enabled = op.get_bind().exec_driver_sql("PRAGMA foreign_keys").scalar() + with op.get_context().autocommit_block(): + op.execute("PRAGMA foreign_keys=OFF") + try: + with op.get_bind().begin_nested(): + yield + finally: + with op.get_context().autocommit_block(): + op.execute(f"PRAGMA foreign_keys={int(enabled)}") + + def mysql_drop_foreignkey_if_exists(constraint_name, table_name, op): """Older Mysql versions do not support DROP FOREIGN KEY IF EXISTS.""" op.execute(f""" @@ -55,6 +75,58 @@ def mysql_drop_foreignkey_if_exists(constraint_name, table_name, op): """) +def raise_if_rows_exist(select_sql: str, message: str, op) -> None: + """ + Abort the migration when ``select_sql`` returns a row, using SQL so it also runs offline. + + Call it before any DDL: MySQL DDL is not transactional, so a late failure leaves a half-applied + migration. ``message`` must not contain single quotes or percent signs and, for MySQL, must be at + most 128 characters. PostgreSQL and MySQL only. + """ + dialect = op.get_context().dialect.name + if dialect == "postgresql": + op.execute( + dedent(""" + DO $$ + BEGIN + IF EXISTS ({select_sql}) THEN + RAISE EXCEPTION '{message}'; + END IF; + END $$ + """).format(select_sql=select_sql, message=message) + ) + elif dialect == "mysql": + create = ( + dedent(""" + CREATE PROCEDURE AirflowMigrationGuard() + BEGIN + IF EXISTS ({select_sql}) THEN + SIGNAL SQLSTATE '45000' SET MESSAGE_TEXT = '{message}'; + END IF; + END + """) + .strip() + .format(select_sql=select_sql, message=message) + ) + context = op.get_context() + if context.as_sql: + # The mysql client splits on ';' unless the procedure body is fenced by DELIMITER. + context.output_buffer.write( + "DROP PROCEDURE IF EXISTS AirflowMigrationGuard;\n" + f"DELIMITER //\n{create}//\nDELIMITER ;\n" + "CALL AirflowMigrationGuard();\nDROP PROCEDURE AirflowMigrationGuard;\n\n" + ) + else: + op.execute("DROP PROCEDURE IF EXISTS AirflowMigrationGuard") + op.execute(create) + try: + op.execute("CALL AirflowMigrationGuard()") + finally: + op.execute("DROP PROCEDURE AirflowMigrationGuard") + else: + raise NotImplementedError(f"No SQL guard for dialect {dialect}") + + def ignore_sqlite_value_error(): from alembic import op diff --git a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py index 1ded51930f844..c4a6e9a04013f 100644 --- a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py +++ b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py @@ -26,13 +26,12 @@ from __future__ import annotations -from contextlib import contextmanager - import sqlalchemy as sa from alembic import op from sqlalchemy.dialects import postgresql, sqlite from airflow.migrations.db_types import TIMESTAMP, StringID +from airflow.migrations.utils import sqlite_rebuilds from airflow.utils.sqlalchemy import ExecutorConfigType, ExtendedJSON, UtcDateTime revision = "e7c2a91bd540" @@ -99,25 +98,6 @@ ) -# Unlike disable_sqlite_fkeys, a failed SQLite upgrade rolls back atomically and foreign_keys is always restored. -@contextmanager -def _sqlite_rebuilds(): - if op.get_bind().dialect.name != "sqlite": - yield - return - if op.get_context().as_sql: - raise RuntimeError("SQLite offline SQL cannot render this migration's table rebuilds") - enabled = op.get_bind().exec_driver_sql("PRAGMA foreign_keys").scalar() - with op.get_context().autocommit_block(): - op.execute("PRAGMA foreign_keys=OFF") - try: - with op.get_bind().begin_nested(): - yield - finally: - with op.get_context().autocommit_block(): - op.execute(f"PRAGMA foreign_keys={int(enabled)}") - - def _check_source(): bind = op.get_bind() if op.get_context().as_sql: @@ -190,7 +170,7 @@ def _redirect_legacy(table, constraint, target, *, onupdate=None, not_valid=Fals def upgrade(): """Retain attempts and give legacy and new task data immutable UUID owners.""" _check_source() - with _sqlite_rebuilds(): + with sqlite_rebuilds(op): owner = op.create_table( "legacy_task_data_owner", sa.Column("dag_id", StringID(), nullable=False), @@ -387,7 +367,7 @@ def downgrade(): ) ): raise RuntimeError("Cannot downgrade attempt ownership after legacy owners changed coordinates") - with _sqlite_rebuilds(): + with sqlite_rebuilds(op): op.drop_table("rtif_v2") op.drop_table("xcom_v2") op.drop_index("ti_current_state", table_name="task_instance") diff --git a/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py b/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py new file mode 100644 index 0000000000000..80d6c5e1aadf2 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py @@ -0,0 +1,211 @@ +# 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. + +""" +Add dynamic region storage. + +Revision ID: 54a27b6f9d01 +Revises: e7c2a91bd540 +Create Date: 2026-09-27 20:00:00.000000 +""" + +from __future__ import annotations + +from textwrap import dedent +from uuid import UUID + +import sqlalchemy as sa +from alembic import op + +from airflow.migrations.utils import raise_if_rows_exist, sqlite_rebuilds +from airflow.models.base import StringID +from airflow.utils.sqlalchemy import CompactUUID, UtcDateTime, compact_uuid_default + +revision = "54a27b6f9d01" +down_revision = "e7c2a91bd540" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + +_SENTINEL = UUID(int=0) +_KEYS = ( + ( + "task_instance", + "task_instance_current_key", + ("dag_id", "task_id", "run_id", "map_index", "working_set"), + "working_set IS NOT NULL", + ), + ( + "task_instance", + "task_instance_try_key", + ("dag_id", "task_id", "run_id", "map_index", "try_number"), + None, + ), + ("task_state_store", "task_state_store_uq", ("dag_run_id", "task_id", "map_index", "key"), None), +) + + +def _replace_unique(table_name, constraint_name, columns): + context = op.get_context() + dialect = context.dialect.name + if dialect in ("postgresql", "mysql"): + quote = context.dialect.identifier_preparer.quote + column_list = ", ".join(quote(name) for name in columns) + if dialect == "postgresql": + sql = f""" + ALTER TABLE {table_name} + DROP CONSTRAINT {constraint_name}, + ADD CONSTRAINT {constraint_name} UNIQUE ({column_list}) + """ + else: + sql = f""" + ALTER TABLE {table_name} + DROP INDEX {constraint_name}, + ADD UNIQUE KEY {constraint_name} ({column_list}) + """ + op.execute(dedent(sql)) + else: + with op.batch_alter_table(table_name) as batch_op: + batch_op.drop_constraint(constraint_name, type_="unique") + batch_op.create_unique_constraint(constraint_name, columns) + + +def _build_collision_query(table_name, columns, where): + """Select old-key groups holding more than one row, scanning only DagRuns that have regions.""" + table = sa.table(table_name, *(sa.column(name) for name in columns)) + region = sa.table("dynamic_region", sa.column("dag_id"), sa.column("run_id")) + if table_name == "task_state_store": + dag_run = sa.table("dag_run", sa.column("id"), sa.column("dag_id"), sa.column("run_id")) + has_region = ( + sa.select(1) + .select_from( + dag_run.join( + region, sa.and_(region.c.dag_id == dag_run.c.dag_id, region.c.run_id == dag_run.c.run_id) + ) + ) + .where(dag_run.c.id == table.c.dag_run_id) + .exists() + ) + else: + has_region = ( + sa.select(1) + .select_from(region) + .where(region.c.dag_id == table.c.dag_id, region.c.run_id == table.c.run_id) + .exists() + ) + query = sa.select(1).select_from(table).where(has_region) + if where: + query = query.where(sa.text(where)) + return query.group_by(*(table.c[name] for name in columns)).having(sa.func.count() > 1).limit(1) + + +def _assert_downgrade_is_lossless(): + """ + Abort the downgrade if dropping ``region_id`` would merge rows, before any DDL runs. + + MySQL DDL cannot be rolled back, so every check has to happen first. The SQL guard also works + in offline mode but does not support SQLite, which is checked by running the query directly. + The query scans the whole ``task_instance`` table, so expect it to be slow on very large installations. + """ + dialect = op.get_context().dialect + for table_name, _, columns, where in _KEYS: + query = _build_collision_query(table_name, columns, where) + message = ( + f"Cannot downgrade: {table_name} has region rows colliding on the old key. " + "Delete the affected DagRuns first." + ) + if dialect.name == "sqlite": + if op.get_bind().execute(query).scalar_one_or_none() is not None: + raise RuntimeError(message) + else: + sql = str(query.compile(dialect=dialect, compile_kwargs={"literal_binds": True})) + raise_if_rows_exist(sql, message, op) + + +def _configure_index_builds(): + if op.get_bind().dialect.name == "postgresql": + op.execute("SET LOCAL statement_timeout = 0") + op.execute( + "SELECT set_config('maintenance_work_mem', '256MB', true) " + "WHERE (SELECT setting::bigint FROM pg_settings WHERE name = 'maintenance_work_mem') < 262144" + ) + + +def upgrade(): + _configure_index_builds() + with sqlite_rebuilds(op): + op.create_table( + "dynamic_region", + sa.Column("id", CompactUUID(), nullable=False), + sa.Column("dag_id", StringID(), nullable=False), + sa.Column("run_id", StringID(), nullable=False), + sa.Column("node_id", StringID(), nullable=False), + 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("resumes_from_index", sa.Integer(), nullable=False, server_default="0"), + sa.Column("created_at", UtcDateTime(), nullable=False), + sa.PrimaryKeyConstraint("id", name="dynamic_region_pkey"), + sa.ForeignKeyConstraint( + ["dag_id", "run_id"], + ["dag_run.dag_id", "dag_run.run_id"], + name="dynamic_region_dag_run_fkey", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["parent_region_id"], + ["dynamic_region.id"], + name="dynamic_region_parent_region_id_fkey", + ondelete="CASCADE", + ), + # Fork lineage is unbounded; a cascading self-FK would exceed MySQL's cascade depth limit. + sa.UniqueConstraint("forked_from_region_id", name="dynamic_region_forked_from_region_id_uq"), + sa.CheckConstraint( + "(parent_region_id IS NULL AND parent_region_index IS NULL) OR " + "(parent_region_id IS NOT NULL AND parent_region_index IS NOT NULL)", + name="parent_coordinates_paired", + ), + sa.CheckConstraint("resumes_from_index >= 0", name="resumes_from_index_nonnegative"), + ) + op.create_index( + "idx_dynamic_region_slot", + "dynamic_region", + ["dag_id", "run_id", "node_id", "parent_region_id", "parent_region_index"], + ) + op.create_index("idx_dynamic_region_parent_region_id", "dynamic_region", ["parent_region_id"]) + for table_name in ("task_instance", "task_state_store"): + op.add_column( + table_name, + sa.Column( + "region_id", CompactUUID(), nullable=False, server_default=compact_uuid_default(_SENTINEL) + ), + ) + for table_name, constraint_name, columns, _ in _KEYS: + new_columns = list(columns) + new_columns.insert(new_columns.index("map_index"), "region_id") + _replace_unique(table_name, constraint_name, new_columns) + + +def downgrade(): + _configure_index_builds() + _assert_downgrade_is_lossless() + with sqlite_rebuilds(op): + for table_name, constraint_name, columns, _ in _KEYS: + _replace_unique(table_name, constraint_name, list(columns)) + for table_name in ("task_instance", "task_state_store"): + op.drop_column(table_name, "region_id") + op.drop_table("dynamic_region") diff --git a/airflow-core/src/airflow/models/__init__.py b/airflow-core/src/airflow/models/__init__.py index c88209f37d3d8..ac16af219134c 100644 --- a/airflow-core/src/airflow/models/__init__.py +++ b/airflow-core/src/airflow/models/__init__.py @@ -69,6 +69,7 @@ def import_all_models(): import airflow.models.dagbundle import airflow.models.dagwarning import airflow.models.deadline_alert + import airflow.models.dynamic_region import airflow.models.errors import airflow.models.revoked_token import airflow.models.serialized_dag diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index db332645fe449..23467c2506730 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -88,6 +88,7 @@ from airflow.models.backfill import Backfill from airflow.models.base import Base, StringID from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel +from airflow.models.dynamic_region import SENTINEL_REGION_ID from airflow.models.taskinstance import TaskInstance as TI, _add_and_prime_mapped_ti, clear_task_instances from airflow.models.tasklog import LogTemplate from airflow.serialization.definitions.deadline import SerializedReferenceModels @@ -2180,6 +2181,7 @@ def _revise_map_indexes_if_mapped( TI.dag_id == self.dag_id, TI.task_id == task.task_id, TI.run_id == self.run_id, + TI.region_id == SENTINEL_REGION_ID, ) ) existing_indexes = set(query) @@ -2192,6 +2194,7 @@ def _revise_map_indexes_if_mapped( TI.dag_id == self.dag_id, TI.task_id == task.task_id, TI.run_id == self.run_id, + TI.region_id == SENTINEL_REGION_ID, TI.map_index.in_(removed_indexes), ) .values(state=TaskInstanceState.REMOVED) @@ -2206,6 +2209,7 @@ def _revise_map_indexes_if_mapped( task_id=task.task_id, run_id=self.run_id, map_indexes=missing_indexes, + region_id=SENTINEL_REGION_ID, session=session, ) diff --git a/airflow-core/src/airflow/models/dynamic_region.py b/airflow-core/src/airflow/models/dynamic_region.py new file mode 100644 index 0000000000000..19363c9e6ce1b --- /dev/null +++ b/airflow-core/src/airflow/models/dynamic_region.py @@ -0,0 +1,80 @@ +# 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. +from __future__ import annotations + +from datetime import datetime +from uuid import UUID + +import uuid6 +from sqlalchemy import CheckConstraint, ForeignKeyConstraint, Index, Integer, UniqueConstraint +from sqlalchemy.orm import Mapped, mapped_column + +from airflow._shared.timezones import timezone +from airflow.models.base import Base, StringID +from airflow.utils.sqlalchemy import CompactUUID, UtcDateTime + +SENTINEL_REGION_ID = UUID(int=0) + + +class DynamicRegion(Base): + """ + One execution of a construct that creates task instances at run time, within a Dag run. + + That is a loop, or the expansion of a mapped task or task group. Task instances created by the + execution carry its ``id`` as ``region_id``, so rows that share dag, task, run and index but come + from different executions can coexist. A region can sit inside another one (``parent_region_*``). + When a clear replaces an execution, the successor records the region it was forked from and where + it resumes. Rows never change after they are created; everything that does change lives on the + task instances. + """ + + __tablename__ = "dynamic_region" + + id: Mapped[UUID] = mapped_column(CompactUUID(), primary_key=True, default=uuid6.uuid7) + dag_id: Mapped[str] = mapped_column(StringID(), nullable=False) + run_id: Mapped[str] = mapped_column(StringID(), nullable=False) + # The id of the task group (a loop) or mapped task that this region executes. + node_id: Mapped[str] = mapped_column(StringID(), nullable=False) + parent_region_id: Mapped[UUID | None] = mapped_column(CompactUUID(), nullable=True) + parent_region_index: Mapped[int | None] = mapped_column(Integer, nullable=True) + forked_from_region_id: Mapped[UUID | None] = mapped_column(CompactUUID(), nullable=True) + resumes_from_index: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0") + created_at: Mapped[datetime] = mapped_column(UtcDateTime, nullable=False, default=timezone.utcnow) + + __table_args__ = ( + ForeignKeyConstraint( + [dag_id, run_id], + ["dag_run.dag_id", "dag_run.run_id"], + name="dynamic_region_dag_run_fkey", + ondelete="CASCADE", + ), + ForeignKeyConstraint( + [parent_region_id], + ["dynamic_region.id"], + name="dynamic_region_parent_region_id_fkey", + ondelete="CASCADE", + ), + UniqueConstraint("forked_from_region_id", name="dynamic_region_forked_from_region_id_uq"), + CheckConstraint( + "(parent_region_id IS NULL AND parent_region_index IS NULL) OR " + "(parent_region_id IS NOT NULL AND parent_region_index IS NOT NULL)", + name="parent_coordinates_paired", + ), + CheckConstraint("resumes_from_index >= 0", name="resumes_from_index_nonnegative"), + Index("idx_dynamic_region_slot", dag_id, run_id, node_id, parent_region_id, parent_region_index), + Index("idx_dynamic_region_parent_region_id", parent_region_id), + ) diff --git a/airflow-core/src/airflow/models/task_state_store.py b/airflow-core/src/airflow/models/task_state_store.py index 7ca497defe55b..0ea7a9b79b0a6 100644 --- a/airflow-core/src/airflow/models/task_state_store.py +++ b/airflow-core/src/airflow/models/task_state_store.py @@ -17,21 +17,23 @@ from __future__ import annotations from datetime import datetime +from uuid import UUID from sqlalchemy import ForeignKeyConstraint, Index, Integer, String, Text, UniqueConstraint from sqlalchemy.dialects.mysql import MEDIUMTEXT -from sqlalchemy.orm import Mapped, mapped_column +from sqlalchemy.orm import Mapped, mapped_column, synonym from airflow._shared.timezones import timezone from airflow.models.base import COLLATION_ARGS, Base, StringID -from airflow.utils.sqlalchemy import UtcDateTime +from airflow.models.dynamic_region import SENTINEL_REGION_ID +from airflow.utils.sqlalchemy import CompactUUID, UtcDateTime, compact_uuid_default class TaskStateStoreModel(Base): """ Persists key/value state for a task within a single DAG run. - Scoped to (dag_run_id, task_id, map_index). Retries of the same task share + Scoped to (dag_run_id, task_id, region_id, map_index). Retries of the same task share the same rows — that is the point. Different DAG runs have different dag_run_id values so they get independent namespaces automatically. """ @@ -43,6 +45,13 @@ class TaskStateStoreModel(Base): dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False) task_id: Mapped[str] = mapped_column(StringID(), nullable=False) map_index: Mapped[int] = mapped_column(Integer, nullable=False, server_default="-1") + region_index = synonym("map_index") + region_id: Mapped[UUID] = mapped_column( + CompactUUID(), + nullable=False, + default=SENTINEL_REGION_ID, + server_default=compact_uuid_default(SENTINEL_REGION_ID), + ) key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False) dag_id: Mapped[str] = mapped_column(StringID(), nullable=False) @@ -57,7 +66,9 @@ class TaskStateStoreModel(Base): expires_at: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) __table_args__ = ( - UniqueConstraint("dag_run_id", "task_id", "map_index", "key", name="task_state_store_uq"), + UniqueConstraint( + "dag_run_id", "task_id", "region_id", "map_index", "key", name="task_state_store_uq" + ), ForeignKeyConstraint( ["dag_run_id"], ["dag_run.id"], diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 19934254e97e8..17f0432167f02 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -75,6 +75,7 @@ mapped_column, reconstructor, relationship, + synonym, with_loader_criteria, ) from sqlalchemy.orm.attributes import NO_VALUE, set_committed_value @@ -101,6 +102,7 @@ from airflow.models.dagbag import DBDagBag from airflow.models.deadline import Deadline, ReferenceModels from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel +from airflow.models.dynamic_region import SENTINEL_REGION_ID # Import HITLDetail at runtime so SQLAlchemy can resolve the relationship from airflow.models.hitl import HITLDetail # noqa: F401 @@ -120,7 +122,14 @@ from airflow.utils.platform import getuser from airflow.utils.retries import run_with_db_retries from airflow.utils.session import NEW_SESSION, create_session, provide_session -from airflow.utils.sqlalchemy import ExecutorConfigType, ExtendedJSON, UtcDateTime, with_row_locks +from airflow.utils.sqlalchemy import ( + CompactUUID, + ExecutorConfigType, + ExtendedJSON, + UtcDateTime, + compact_uuid_default, + with_row_locks, +) from airflow.utils.state import DagRunState, State, TaskInstanceState TR = TaskReschedule @@ -709,6 +718,13 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload): dag_id: Mapped[str] = mapped_column(StringID(), nullable=False) run_id: Mapped[str] = mapped_column(StringID(), nullable=False) map_index: Mapped[int] = mapped_column(Integer, nullable=False, server_default="-1") + region_index = synonym("map_index") + region_id: Mapped[UUID] = mapped_column( + CompactUUID(), + nullable=False, + default=SENTINEL_REGION_ID, + server_default=compact_uuid_default(SENTINEL_REGION_ID), + ) start_date: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) end_date: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) @@ -798,10 +814,22 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload): ), PrimaryKeyConstraint("id", name="task_instance_pkey"), UniqueConstraint( - "dag_id", "task_id", "run_id", "map_index", "working_set", name="task_instance_current_key" + "dag_id", + "task_id", + "run_id", + "region_id", + "map_index", + "working_set", + name="task_instance_current_key", ), UniqueConstraint( - "dag_id", "task_id", "run_id", "map_index", "try_number", name="task_instance_try_key" + "dag_id", + "task_id", + "run_id", + "region_id", + "map_index", + "try_number", + name="task_instance_try_key", ), ForeignKeyConstraint( [trigger_id], @@ -867,11 +895,14 @@ def __init__( run_id: str | None = None, state: str | None = None, map_index: int = -1, + *, + region_id: UUID = SENTINEL_REGION_ID, ): super().__init__() self.dag_id = task.dag_id self.task_id = task.task_id self.map_index = map_index + self.region_id = region_id if run_id is not None: self.run_id = run_id self.try_number = 0 @@ -907,7 +938,13 @@ def stats_tags(self) -> dict[str, str]: @staticmethod def insert_mapping( - run_id: str, task: Operator, map_index: int, *, dag_version_id: UUID | None, dag_run: DagRun + run_id: str, + task: Operator, + map_index: int, + *, + dag_version_id: UUID | None, + dag_run: DagRun, + region_id: UUID = SENTINEL_REGION_ID, ) -> dict[str, Any]: """ Insert mapping. @@ -918,7 +955,13 @@ def insert_mapping( if not hasattr(weight_rule, "get_weight"): weight_rule = validate_and_load_priority_weight_strategy(weight_rule) priority_weight = weight_rule.get_weight( - TaskInstance(task=task, run_id=run_id, map_index=map_index, dag_version_id=dag_version_id) + TaskInstance( + task=task, + run_id=run_id, + map_index=map_index, + region_id=region_id, + dag_version_id=dag_version_id, + ) ) context_carrier = new_task_run_carrier(dag_run.context_carrier) @@ -940,6 +983,7 @@ def insert_mapping( "operator": task.task_type, "custom_operator_name": getattr(task, "operator_name", None), "map_index": map_index, + "region_id": region_id, "_task_display_property_value": task.task_display_name, "dag_version_id": dag_version_id, "context_carrier": context_carrier, @@ -1017,6 +1061,7 @@ def get_task_instance( map_index: int, lock_for_update: bool = False, *, + region_id: UUID = SENTINEL_REGION_ID, session: Session = NEW_SESSION, ) -> TaskInstance | None: query = ( @@ -1027,6 +1072,7 @@ def get_task_instance( run_id=run_id, task_id=task_id, map_index=map_index, + region_id=region_id, ) ) @@ -1209,6 +1255,7 @@ def get_last_try_numbers( task_id: str, run_id: str, map_indexes: Collection[int] | None = None, + region_id: UUID = SENTINEL_REGION_ID, session: Session, ) -> dict[int, int]: """Return the highest try number per map index across current and historical task instances.""" @@ -1216,7 +1263,12 @@ def get_last_try_numbers( return {} statement = ( select(cls.map_index, func.max(cls.try_number)) - .where(cls.dag_id == dag_id, cls.task_id == task_id, cls.run_id == run_id) + .where( + cls.dag_id == dag_id, + cls.task_id == task_id, + cls.run_id == run_id, + cls.region_id == region_id, + ) .group_by(cls.map_index) .execution_options(include_all_attempts=True) ) @@ -2381,13 +2433,21 @@ def _get_num_task_instances_of_state( @staticmethod def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> ColumnElement[bool] | None: """Return SQLAlchemy filter to query selected task instances.""" - # DictKeys type, (what we often pass here from the scheduler) is not directly indexable :( - # Or it might be a generator, but we need to be able to iterate over it more than once - tis = list(tis) - - if not tis: + by_region: dict[UUID, list[TaskInstance | TaskInstanceKey]] = defaultdict(list) + # ``tis`` may be a one-shot iterable (dict keys view, generator); this loop is its only consumer. + for ti in tis: + by_region[SENTINEL_REGION_ID if isinstance(ti, TaskInstanceKey) else ti.region_id].append(ti) + if not by_region: return None + return or_( + *( + and_(TaskInstance.region_id == region_id, TaskInstance._filter_for_tis(group)) + for region_id, group in by_region.items() + ) + ) + @staticmethod + def _filter_for_tis(tis: list[TaskInstance | TaskInstanceKey]) -> ColumnElement[bool]: first = tis[0] dag_id = first.dag_id @@ -2568,12 +2628,17 @@ def expand_mapped_task(self, *, session: Session) -> tuple[Sequence[TaskInstance TaskInstance.task_id == task.task_id, TaskInstance.run_id == run_id, TaskInstance.map_index == -1, + TaskInstance.region_id == self.region_id, or_(TaskInstance.state.in_(State.unfinished), TaskInstance.state.is_(None)), ) ).one_or_none() last_tries = TaskInstance.get_last_try_numbers( - dag_id=task.dag_id, task_id=task.task_id, run_id=run_id, session=session + dag_id=task.dag_id, + task_id=task.task_id, + run_id=run_id, + region_id=self.region_id, + session=session, ) all_expanded_tis: list[TaskInstance] = [] @@ -2638,6 +2703,7 @@ def expand_mapped_task(self, *, session: Session) -> tuple[Sequence[TaskInstance TaskInstance.dag_id == task.dag_id, TaskInstance.task_id == task.task_id, TaskInstance.run_id == run_id, + TaskInstance.region_id == self.region_id, ) ) or 0 @@ -2667,6 +2733,7 @@ def expand_mapped_task(self, *, session: Session) -> tuple[Sequence[TaskInstance task, run_id=run_id, map_index=index, + region_id=self.region_id, state=state, dag_version_id=dag_version_id, ) @@ -2692,6 +2759,7 @@ def expand_mapped_task(self, *, session: Session) -> tuple[Sequence[TaskInstance TaskInstance.task_id == task.task_id, TaskInstance.run_id == run_id, TaskInstance.map_index >= total_expanded_ti_count, + TaskInstance.region_id == self.region_id, ) to_update = session.scalars(with_row_locks(query, of=TaskInstance, session=session, skip_locked=True)) for ti in to_update: diff --git a/airflow-core/src/airflow/state/metastore.py b/airflow-core/src/airflow/state/metastore.py index 7cc355b489f1d..7847a8f06438c 100644 --- a/airflow-core/src/airflow/state/metastore.py +++ b/airflow-core/src/airflow/state/metastore.py @@ -233,6 +233,7 @@ def _get_task_state_store(self, scope: TaskScope, key: str, *, session: Session) TaskStateStoreModel.dag_id == scope.dag_id, TaskStateStoreModel.run_id == scope.run_id, TaskStateStoreModel.task_id == scope.task_id, + TaskStateStoreModel.region_id == scope.region_id, TaskStateStoreModel.map_index == scope.map_index, TaskStateStoreModel.key == key, ) @@ -262,6 +263,7 @@ def _store_task_state( dag_id=scope.dag_id, run_id=scope.run_id, task_id=scope.task_id, + region_id=scope.region_id, map_index=scope.map_index, key=key, value=value, @@ -271,7 +273,7 @@ def _store_task_state( stmt = _build_upsert_stmt( get_dialect_name(session), TaskStateStoreModel, - ["dag_run_id", "task_id", "map_index", "key"], + ["dag_run_id", "task_id", "region_id", "map_index", "key"], values, dict(value=value, updated_at=now, expires_at=expires_at), ) @@ -283,6 +285,7 @@ def _delete_task_state_store(self, scope: TaskScope, key: str, *, session: Sessi TaskStateStoreModel.dag_id == scope.dag_id, TaskStateStoreModel.run_id == scope.run_id, TaskStateStoreModel.task_id == scope.task_id, + TaskStateStoreModel.region_id == scope.region_id, TaskStateStoreModel.map_index == scope.map_index, TaskStateStoreModel.key == key, ) @@ -295,6 +298,7 @@ def _clear_task_state_store( TaskStateStoreModel.dag_id == scope.dag_id, TaskStateStoreModel.run_id == scope.run_id, TaskStateStoreModel.task_id == scope.task_id, + TaskStateStoreModel.region_id == scope.region_id, ] if not all_map_indices: conditions.append(TaskStateStoreModel.map_index == scope.map_index) @@ -434,6 +438,7 @@ async def _aget_task_state_store( TaskStateStoreModel.dag_id == scope.dag_id, TaskStateStoreModel.run_id == scope.run_id, TaskStateStoreModel.task_id == scope.task_id, + TaskStateStoreModel.region_id == scope.region_id, TaskStateStoreModel.map_index == scope.map_index, TaskStateStoreModel.key == key, ) @@ -463,6 +468,7 @@ async def _aset_task_state_store( dag_id=scope.dag_id, run_id=scope.run_id, task_id=scope.task_id, + region_id=scope.region_id, map_index=scope.map_index, key=key, value=value, @@ -473,7 +479,7 @@ async def _aset_task_state_store( stmt = _build_upsert_stmt( get_dialect_name(session.sync_session), TaskStateStoreModel, - ["dag_run_id", "task_id", "map_index", "key"], + ["dag_run_id", "task_id", "region_id", "map_index", "key"], values, dict(value=value, updated_at=now, expires_at=expires_at), ) @@ -485,6 +491,7 @@ async def _adelete_task_state_store(self, scope: TaskScope, key: str, *, session TaskStateStoreModel.dag_id == scope.dag_id, TaskStateStoreModel.run_id == scope.run_id, TaskStateStoreModel.task_id == scope.task_id, + TaskStateStoreModel.region_id == scope.region_id, TaskStateStoreModel.map_index == scope.map_index, TaskStateStoreModel.key == key, ) @@ -497,6 +504,7 @@ async def _aclear_task_state_store( TaskStateStoreModel.dag_id == scope.dag_id, TaskStateStoreModel.run_id == scope.run_id, TaskStateStoreModel.task_id == scope.task_id, + TaskStateStoreModel.region_id == scope.region_id, ] if not all_map_indices: conditions.append(TaskStateStoreModel.map_index == scope.map_index) diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 7e276d02c8651..826bc4ae8877b 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "e7c2a91bd540", + "3.4.0": "54a27b6f9d01", } # Prefix used to identify tables holding data moved during migration. @@ -1760,6 +1760,9 @@ def _normalize_mysql_server_default(default: Any) -> str | None: return None normalized = str(default).strip().strip("'\"").lower() + # MySQL reflects expression defaults with escaped quotes and a charset introducer, + # e.g. unhex(_utf8mb4\'00ff\') for a model default of UNHEX('00ff'). + normalized = re.sub(r"\b_[a-z0-9]+(?=')", "", normalized.replace("\\'", "'")) while normalized.startswith("(") and normalized.endswith(")"): normalized = normalized[1:-1].strip() normalized = normalized.replace(" ", "") diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index 92cebbd236e6b..02ba0975d5c88 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/airflow-core/src/airflow/utils/db_cleanup.py @@ -401,7 +401,7 @@ def _dump_table_to_file(*, target_table: str, file_path: str, export_format: str "task_reschedule": "ti_id", "hitl_detail": "ti_id", } -_ATTEMPT_ARCHIVE_TABLES = {*_ATTEMPT_UUID_CHILDREN, "xcom_v1", "rtif_v1", "task_instance"} +_ATTEMPT_ARCHIVE_TABLES = {*_ATTEMPT_UUID_CHILDREN, "xcom_v1", "rtif_v1", "task_instance", "dynamic_region"} def _archive_cascading_attempt_data( @@ -435,6 +435,12 @@ def _archive_cascading_attempt_data( ) if source.name == "dag_run": selections["task_instance"] = select(ti).where(ti.c.id.in_(producer_ids)) + region = tables["dynamic_region"] + selections["dynamic_region"] = select(region).where( + tuple_(region.c.dag_id, region.c.run_id).in_( + select(parent_archive.c.dag_id, parent_archive.c.run_id) + ) + ) if not archives: archive_names = {} diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index 61ac586bda6cf..6b7db82b8efe4 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -24,8 +24,9 @@ import logging from collections.abc import Generator from typing import TYPE_CHECKING, Any +from uuid import UUID -from sqlalchemy import TIMESTAMP, PickleType, String, event, nullsfirst, text +from sqlalchemy import TIMESTAMP, PickleType, String, Uuid, event, nullsfirst, text from sqlalchemy.dialects import mysql from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.ext.compiler import compiles @@ -291,6 +292,60 @@ def load_dialect_impl(self, dialect): return super().load_dialect_impl(dialect) +_MYSQL_DIALECT_NAMES = ("mysql", "mariadb") + + +class CompactUUID(TypeDecorator): + """ + A UUID that MySQL stores as ``BINARY(16)``; other dialects use SQLAlchemy's :class:`~sqlalchemy.types.Uuid`. + + The default MySQL form, ``CHAR(32)``, takes 128 bytes in an index key under a utf8mb4 collation. + Keys that already approach InnoDB's 3072-byte limit, such as those holding several ``StringID`` columns, + overflow with it. + """ + + impl = Uuid + cache_ok = True + + def load_dialect_impl(self, dialect): + if dialect.name in _MYSQL_DIALECT_NAMES: + return dialect.type_descriptor(mysql.BINARY(16)) + return dialect.type_descriptor(Uuid()) + + def process_bind_param(self, value, dialect): + if value is None or dialect.name not in _MYSQL_DIALECT_NAMES: + return value + return (UUID(value) if isinstance(value, str) else value).bytes + + def process_result_value(self, value, dialect): + if value is None or dialect.name not in _MYSQL_DIALECT_NAMES: + return value + return UUID(bytes=bytes(value)) + + +class compact_uuid_default(ColumnElement): + """Server default for a :class:`CompactUUID` column, rendered in the form each dialect stores.""" + + inherit_cache = True + type = NullType() + + def __init__(self, value: UUID): + self.value = value + + +@compiles(compact_uuid_default) +def _compact_uuid_default(element, compiler, **kw): + return compiler.render_literal_value(element.value.hex, String()) + + +@compiles(compact_uuid_default, "mysql") +@compiles(compact_uuid_default, "mariadb") +def _compact_uuid_default_mysql(element, compiler, **kw): + # SQLAlchemy only parenthesizes expression defaults when it knows the server version, which offline SQL + # generation (--show-sql-only) does not; MySQL rejects an unparenthesized expression default. + return f"(UNHEX({compiler.render_literal_value(element.value.hex, String())}))" + + class ExtendedJSON(TypeDecorator): """ A version of the JSON column that uses the Airflow extended JSON serialization. diff --git a/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py b/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py new file mode 100644 index 0000000000000..16fc5206149bc --- /dev/null +++ b/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py @@ -0,0 +1,231 @@ +# 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 the dynamic region migration, written against the schema as it stood at that revision. + +The tables are declared here instead of taken from ``airflow.models``, and alembic moves the database +to the revision, so later migrations and model changes cannot change what these tests exercise. +""" + +from __future__ import annotations + +from importlib import import_module +from uuid import UUID, uuid4 + +import pytest +import sqlalchemy as sa +from sqlalchemy import inspect +from sqlalchemy.exc import DBAPIError + +from airflow import settings +from airflow._shared.timezones import timezone +from airflow.utils.db import downgrade, upgradedb +from airflow.utils.sqlalchemy import CompactUUID, UtcDateTime + +pytestmark = pytest.mark.db_test + +_migration = import_module("airflow.migrations.versions.0143_3_4_0_add_dynamic_region_storage") +REVISION = _migration.revision +PREVIOUS_REVISION = _migration.down_revision +DAG_ID = "migration_0143" +RUN_ID = "run" +SENTINEL = UUID(int=0) + +metadata = sa.MetaData() +dag_run = sa.Table( + "dag_run", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("dag_id", sa.String(250)), + sa.Column("run_id", sa.String(250)), + sa.Column("run_type", sa.String(50)), + sa.Column("run_after", UtcDateTime), +) +task_instance = sa.Table( + "task_instance", + metadata, + sa.Column("id", sa.Uuid, primary_key=True), + sa.Column("dag_id", sa.String(250)), + sa.Column("task_id", sa.String(250)), + sa.Column("run_id", sa.String(250)), + sa.Column("map_index", sa.Integer), + sa.Column("region_id", CompactUUID), + sa.Column("working_set", sa.Boolean), + sa.Column("try_number", sa.Integer), + sa.Column("pool", sa.String(256)), + sa.Column("pool_slots", sa.Integer), +) +task_state_store = sa.Table( + "task_state_store", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("dag_run_id", sa.Integer), + sa.Column("dag_id", sa.String(250)), + sa.Column("task_id", sa.String(250)), + sa.Column("run_id", sa.String(250)), + sa.Column("map_index", sa.Integer), + sa.Column("region_id", CompactUUID), + sa.Column("key", sa.String(512)), + sa.Column("value", sa.Text), + sa.Column("updated_at", UtcDateTime), +) +dynamic_region = sa.Table( + "dynamic_region", + metadata, + sa.Column("id", CompactUUID, primary_key=True), + sa.Column("dag_id", sa.String(250)), + sa.Column("run_id", sa.String(250)), + sa.Column("node_id", sa.String(250)), + sa.Column("created_at", UtcDateTime), +) +TABLES = (task_instance, task_state_store) + + +def _delete_rows(connection): + connection.execute(dynamic_region.delete().where(dynamic_region.c.dag_id == DAG_ID)) + connection.execute(task_state_store.delete().where(task_state_store.c.dag_id == DAG_ID)) + connection.execute(task_instance.delete().where(task_instance.c.dag_id == DAG_ID)) + connection.execute(dag_run.delete().where(dag_run.c.dag_id == DAG_ID)) + + +@pytest.fixture +def database_at_revision(): + """Move the database to the migration's revision for one test, then back to the current head.""" + original_async_conn = settings.SQL_ALCHEMY_CONN_ASYNC + settings.SQL_ALCHEMY_CONN_ASYNC = "" + settings.dispose_orm(do_log=False) + settings.configure_orm(disable_connection_pool=True) + try: + downgrade(to_revision=REVISION) + with settings.engine.begin() as connection: + _delete_rows(connection) + yield settings.engine + finally: + upgradedb(to_revision=REVISION) + with settings.engine.begin() as connection: + _delete_rows(connection) + upgradedb() + settings.SQL_ALCHEMY_CONN_ASYNC = original_async_conn + settings.reconfigure_orm() + + +@pytest.fixture +def stored_rows(database_at_revision): + """One task instance and one state entry, both in the sentinel region, in a Dag run with no regions.""" + with database_at_revision.begin() as connection: + dag_run_id = connection.execute( + dag_run.insert().values( + dag_id=DAG_ID, run_id=RUN_ID, run_type="manual", run_after=timezone.utcnow() + ) + ).inserted_primary_key[0] + connection.execute( + task_instance.insert().values( + id=uuid4(), + dag_id=DAG_ID, + task_id="task", + run_id=RUN_ID, + map_index=-1, + working_set=True, + try_number=0, + pool="default_pool", + pool_slots=1, + ) + ) + connection.execute( + task_state_store.insert().values( + dag_run_id=dag_run_id, + dag_id=DAG_ID, + task_id="task", + run_id=RUN_ID, + map_index=-1, + key="state", + value="value", + updated_at=timezone.utcnow(), + ) + ) + return dag_run_id + + +def _add_region_row(connection, table, **overrides): + """Copy the sentinel-region row of ``table`` into a new region, changing the given columns.""" + row = connection.execute(sa.select(table).where(table.c.dag_id == DAG_ID)).mappings().first() + region_id = uuid4() + connection.execute( + dynamic_region.insert().values( + id=region_id, dag_id=DAG_ID, run_id=RUN_ID, node_id="task", created_at=timezone.utcnow() + ) + ) + values = {name: value for name, value in row.items() if name != "id" or table is task_instance} + if table is task_instance: + values["id"] = uuid4() + connection.execute(table.insert().values(**{**values, "region_id": region_id, **overrides})) + + +def _rows(engine, table): + with engine.connect() as connection: + return connection.execute(sa.select(table).where(table.c.dag_id == DAG_ID)).mappings().all() + + +def test_downgrade_then_upgrade_keeps_existing_rows_in_the_sentinel_region(stored_rows, database_at_revision): + downgrade(to_revision=PREVIOUS_REVISION) + upgradedb(to_revision=REVISION) + + for table in TABLES: + (row,) = _rows(database_at_revision, table) + assert row["region_id"] == SENTINEL + assert row["map_index"] == -1 + + +@pytest.mark.parametrize("table", TABLES, ids=lambda table: table.name) +def test_downgrade_refuses_each_old_coordinate_collision_before_ddl(stored_rows, database_at_revision, table): + with database_at_revision.begin() as connection: + _add_region_row(connection, table) + + with pytest.raises((RuntimeError, DBAPIError), match=table.name): + downgrade(to_revision=PREVIOUS_REVISION) + + tables = inspect(database_at_revision).get_table_names() + assert "dynamic_region" in tables + for retained in TABLES: + assert "region_id" in {c["name"] for c in inspect(database_at_revision).get_columns(retained.name)} + + +@pytest.mark.parametrize( + ("table", "region_rows"), + [ + # Archived tries leave working_set NULL, which the old current-key does not constrain. + pytest.param( + task_instance, + [{"try_number": 1, "working_set": None}, {"try_number": 2, "working_set": None}], + id="task_instance", + ), + pytest.param(task_state_store, [{"key": "state-1"}, {"key": "state-2"}], id="task_state_store"), + ], +) +def test_downgrade_keeps_region_rows_that_differ_on_an_old_key_column( + stored_rows, database_at_revision, table, region_rows +): + with database_at_revision.begin() as connection: + for overrides in region_rows: + _add_region_row(connection, table, **overrides) + + downgrade(to_revision=PREVIOUS_REVISION) + upgradedb(to_revision=REVISION) + + rows = _rows(database_at_revision, table) + assert len(rows) == 3 + assert {row["region_id"] for row in rows} == {SENTINEL} diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 29f98da417a39..c0e6debf6731e 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -2617,6 +2617,7 @@ def test_refresh_from_db(self, create_task_instance): "task_id": "test_refresh_from_db_task", "dag_id": "test_refresh_from_db_dag", "run_id": "test", + "region_id": uuid6.uuid7(), "map_index": -1, "start_date": run_date + datetime.timedelta(days=1), "end_date": run_date + datetime.timedelta(days=1, seconds=1, milliseconds=234), diff --git a/airflow-core/tests/unit/state/test_metastore.py b/airflow-core/tests/unit/state/test_metastore.py index ee0c578c21b48..27048c70a000f 100644 --- a/airflow-core/tests/unit/state/test_metastore.py +++ b/airflow-core/tests/unit/state/test_metastore.py @@ -21,6 +21,7 @@ from datetime import timedelta from typing import TYPE_CHECKING from unittest.mock import patch +from uuid import UUID, uuid4 import pytest import pytest_asyncio @@ -118,6 +119,36 @@ def asset_committed() -> AssetModel: class TestMetastoreBackendTaskScope: + @pytest.mark.parametrize("operation", ["update", "delete", "clear_index", "clear_region"]) + def test_operations_isolate_colliding_regions(self, session, backend, dag_run, operation): + regions = [UUID(int=0), uuid4(), uuid4()] + expected = { + (TaskScope(DAG_ID, RUN_ID, TASK_ID, index, region_id=region), key): f"{region}:{index}:{key}" + for region in regions + for index in (0, 1) + for key in ("value", "other") + } + for (scope, key), value in expected.items(): + backend.set(scope, key, value, session=session) + target = TaskScope(DAG_ID, RUN_ID, TASK_ID, 0, region_id=regions[1]) + missing = TaskScope(DAG_ID, RUN_ID, TASK_ID, 0, region_id=uuid4()) + assert backend.get(missing, "value", session=session) is None + + if operation == "update": + backend.set(target, "value", "updated", session=session) + expected[target, "value"] = "updated" + elif operation == "delete": + backend.delete(target, "value", session=session) + expected[target, "value"] = None + else: + all_indices = operation == "clear_region" + backend.clear(target, all_map_indices=all_indices, session=session) + for scope, key in expected: + if scope.region_id == target.region_id and (all_indices or scope.map_index == 0): + expected[scope, key] = None + for (scope, key), value in expected.items(): + assert backend.get(scope, key, session=session) == value + def test_get_returns_none_for_missing_key( self, session: Session, backend: MetastoreBackend, dag_run: DagRun ): @@ -586,6 +617,36 @@ async def dispose_async_engine(): @pytest.mark.usefixtures("dispose_async_engine") @pytest.mark.asyncio(loop_scope="class") class TestMetastoreBackendAsync: + @pytest.mark.parametrize("operation", ["update", "delete", "clear_index", "clear_region"]) + async def test_operations_isolate_colliding_regions(self, backend, dag_run_committed, operation): + regions = [UUID(int=0), uuid4(), uuid4()] + expected = { + (TaskScope(DAG_ID, RUN_ID, TASK_ID, index, region_id=region), key): f"{region}:{index}:{key}" + for region in regions + for index in (0, 1) + for key in ("value", "other") + } + for (scope, key), value in expected.items(): + await backend.aset(scope, key, value) + target = TaskScope(DAG_ID, RUN_ID, TASK_ID, 0, region_id=regions[1]) + missing = TaskScope(DAG_ID, RUN_ID, TASK_ID, 0, region_id=uuid4()) + assert await backend.aget(missing, "value") is None + + if operation == "update": + await backend.aset(target, "value", "updated") + expected[target, "value"] = "updated" + elif operation == "delete": + await backend.adelete(target, "value") + expected[target, "value"] = None + else: + all_indices = operation == "clear_region" + await backend.aclear(target, all_map_indices=all_indices) + for scope, key in expected: + if scope.region_id == target.region_id and (all_indices or scope.map_index == 0): + expected[scope, key] = None + for (scope, key), value in expected.items(): + assert await backend.aget(scope, key) == value + async def test_aset_and_aget_task_roundtrip(self, backend: MetastoreBackend, dag_run_committed: DagRun): scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) await backend.aset(scope, "job_id", "app_async") diff --git a/airflow-core/tests/unit/utils/test_db.py b/airflow-core/tests/unit/utils/test_db.py index f4e2901fa3024..83a1e65304576 100644 --- a/airflow-core/tests/unit/utils/test_db.py +++ b/airflow-core/tests/unit/utils/test_db.py @@ -40,6 +40,7 @@ LazySelectSequence, _get_alembic_config, _get_current_revision, + _normalize_mysql_server_default, check_migrations, check_team_names_can_be_lower_cased, compare_server_default, @@ -227,6 +228,25 @@ def test_has_pending_upgrade_ops(self, initialized_db): config = _get_alembic_config() check(config) + @pytest.mark.parametrize( + ("reflected", "rendered"), + [ + pytest.param( + "unhex(_utf8mb4\\'00000000000000000000000000000000\\')", + "UNHEX('00000000000000000000000000000000')", + id="mysql8-charset-introducer-and-escaped-quotes", + ), + pytest.param( + "unhex('00000000000000000000000000000000')", + "UNHEX('00000000000000000000000000000000')", + id="plain-expression", + ), + pytest.param("(0)", "false", id="boolean"), + ], + ) + def test_normalize_mysql_server_default_matches_rendered_default(self, reflected, rendered): + assert _normalize_mysql_server_default(reflected) == _normalize_mysql_server_default(rendered) + def test_default_connections_sort(self): conn_ids = [c.conn_id for c in get_default_connections()] assert conn_ids diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index ddf7dc5ae4bab..2649b6768674f 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -53,6 +53,7 @@ from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel from airflow.models.deadline import Deadline +from airflow.models.dynamic_region import DynamicRegion from airflow.models.serialized_dag import SerializedDagModel from airflow.models.task_state_store import TaskStateStoreModel from airflow.models.taskreschedule import TaskReschedule @@ -78,9 +79,11 @@ run_cleanup, ) from airflow.utils.session import create_session +from airflow.utils.sqlalchemy import CompactUUID from airflow.utils.types import DagRunType from tests_common.test_utils.attempt_ownership import ( + COORDINATES, CURRENT_ID, HISTORY_ID, NOW, @@ -1221,6 +1224,7 @@ def test_no_models_missing(self): "asset_watcher", # cascade from trigger "dag_favorite", # cascade from dag "deadline_alert", # cascade from serialized_dag, which cascades from dag_version + "dynamic_region", # cascade from dag_run; archived with it, never aged out on its own "hitl_detail", # cascade from task_instance "hitl_detail_history", # cascade from task_instance_history "job_team", # cascade from job @@ -1669,6 +1673,43 @@ def test_cleanup_reuses_child_archives_across_parent_batches(self, ownership_ses == count ) + @pytest.mark.execution_timeout(10) + def test_dag_run_cleanup_archives_its_regions(self, ownership_session): + session = ownership_session + region_id = uuid4() + session.execute( + sa.insert(DynamicRegion).values( + id=region_id, + dag_id=COORDINATES["dag_id"], + run_id=COORDINATES["run_id"], + node_id="loop", + created_at=NOW - timedelta(days=100), + ) + ) + session.execute(sa.update(TaskInstance).values(region_id=region_id)) + session.commit() + + run_cleanup( + clean_before_timestamp=NOW + timedelta(days=1), + table_names=["dag_run"], + confirm=False, + session=session, + error_on_cleanup_failure=True, + ) + + assert session.scalar(sa.select(sa.func.count()).select_from(DynamicRegion)) == 0 + archived = [ + archive + for archive in _get_archived_table_names(["dag_run"], session) + if archive.startswith(f"{ARCHIVE_TABLE_PREFIX}dynamic_region__") + ] + archived_ids = [ + row[0] + for archive in archived + for row in session.execute(sa.text(f"SELECT id FROM {archive}").columns(id=CompactUUID())) + ] + assert archived_ids == [region_id] + @pytest.mark.parametrize("include", [True, False]) @pytest.mark.execution_timeout(10) def test_cleanup_scopes_attempt_cascade_to_selected_dag(self, ownership_session, include): diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index 8dbcd41313055..e5e883eac8509 100644 --- a/airflow-core/tests/unit/utils/test_sqlalchemy.py +++ b/airflow-core/tests/unit/utils/test_sqlalchemy.py @@ -21,11 +21,14 @@ import pickle from copy import deepcopy from unittest import mock +from uuid import UUID import pytest from kubernetes.client import Configuration, models as k8s -from sqlalchemy import text +from sqlalchemy import Column, MetaData, Table, text +from sqlalchemy.dialects import mysql, postgresql, sqlite from sqlalchemy.exc import StatementError +from sqlalchemy.schema import CreateTable from airflow import settings from airflow.sdk import DAG @@ -34,8 +37,10 @@ from airflow.serialization.serialized_objects import BaseSerialization from airflow.settings import Session from airflow.utils.sqlalchemy import ( + CompactUUID, ExecutorConfigType, apply_regex_query_timeout, + compact_uuid_default, ensure_pod_is_valid_after_unpickling, get_dialect_name, prohibit_commit, @@ -443,3 +448,58 @@ def test_noop_when_timeout_disabled(self): with apply_regex_query_timeout(session): pass session.execute.assert_not_called() + + +class TestCompactUUID: + value = UUID("0190a1b2-c3d4-7e5f-8a9b-0c1d2e3f4a5b") + + @pytest.mark.parametrize( + ("dialect", "expected"), + [ + pytest.param(mysql.dialect(), "BINARY(16)", id="mysql"), + pytest.param(sqlite.dialect(), "CHAR(32)", id="sqlite"), + pytest.param(postgresql.dialect(), "UUID", id="postgresql"), + ], + ) + def test_column_type_per_dialect(self, dialect, expected): + assert CompactUUID().compile(dialect=dialect) == expected + + @pytest.mark.parametrize("given", [value, str(value)], ids=["uuid", "str"]) + def test_mysql_stores_the_sixteen_raw_bytes(self, given): + assert CompactUUID().process_bind_param(given, mysql.dialect()) == self.value.bytes + + def test_mysql_reads_the_bytes_back_as_a_uuid(self): + assert CompactUUID().process_result_value(self.value.bytes, mysql.dialect()) == self.value + + @pytest.mark.parametrize( + "dialect", [sqlite.dialect(), postgresql.dialect()], ids=["sqlite", "postgresql"] + ) + def test_other_dialects_leave_the_value_to_the_uuid_type(self, dialect): + assert CompactUUID().process_bind_param(self.value, dialect) == self.value + assert CompactUUID().process_result_value(self.value, dialect) == self.value + + def test_none_passes_through(self): + assert CompactUUID().process_bind_param(None, mysql.dialect()) is None + assert CompactUUID().process_result_value(None, mysql.dialect()) is None + + @pytest.mark.parametrize( + ("dialect", "expected"), + [ + pytest.param(mysql.dialect(), "(UNHEX('0190a1b2c3d47e5f8a9b0c1d2e3f4a5b'))", id="mysql"), + pytest.param(sqlite.dialect(), "'0190a1b2c3d47e5f8a9b0c1d2e3f4a5b'", id="sqlite"), + pytest.param(postgresql.dialect(), "'0190a1b2c3d47e5f8a9b0c1d2e3f4a5b'", id="postgresql"), + ], + ) + def test_server_default_per_dialect(self, dialect, expected): + assert str(compact_uuid_default(self.value).compile(dialect=dialect)) == expected + + def test_mysql_column_default_is_parenthesized_without_a_server_version(self): + table = Table( + "t", + MetaData(), + Column( + "region_id", CompactUUID(), nullable=False, server_default=compact_uuid_default(self.value) + ), + ) + ddl = str(CreateTable(table).compile(dialect=mysql.dialect())) + assert "DEFAULT (UNHEX('0190a1b2c3d47e5f8a9b0c1d2e3f4a5b'))" in ddl diff --git a/providers/common/io/docs/state_store_backend.rst b/providers/common/io/docs/state_store_backend.rst index 9f56a2a4a99b2..8abe992ec5824 100644 --- a/providers/common/io/docs/state_store_backend.rst +++ b/providers/common/io/docs/state_store_backend.rst @@ -27,7 +27,7 @@ section to ``airflow.providers.common.io.state_store.backend.StateStoreObjectSto ``state_store_objectstorage_path`` to the desired base location. The connection id is obtained from the user part of the URL, e.g. ``state_store_objectstorage_path = s3://conn_id@mybucket/task-state/``. -Task state is stored under ``////`` and asset state under +Task state is stored under ``//[region=/]//`` and asset state under ``assets//`` beneath the configured base path. By default (``state_store_objectstorage_threshold = 0``) all serialized values are offloaded to object storage. @@ -76,6 +76,6 @@ Using the local filesystem (useful for development):: .. note:: - Task state paths are keyed on ``(dag_id, run_id, task_id, map_index)`` and are stable across + Task state paths are keyed on ``(dag_id, run_id, region_id, task_id, map_index)`` and are stable across task retries. This makes this backend suitable for operators that use :class:`~airflow.sdk.ResumableJobMixin` to reconnect to external jobs after a retry. diff --git a/providers/common/io/provider.yaml b/providers/common/io/provider.yaml index 57782cfc895bc..a210285ab0c8b 100644 --- a/providers/common/io/provider.yaml +++ b/providers/common/io/provider.yaml @@ -122,8 +122,8 @@ config: description: | Base path on object storage for the task/asset state store backend, in URL format. When set, StateStoreObjectStorageBackend will persist task and asset state under this - prefix, organized as //// for tasks and - assets// for assets. + prefix, organized as //[region=/]// + for tasks and assets// for assets. version_added: 1.8.0 type: string example: "s3://conn_id@bucket/task-state/" diff --git a/providers/common/io/src/airflow/providers/common/io/get_provider_info.py b/providers/common/io/src/airflow/providers/common/io/get_provider_info.py index e7388270c4d0c..caba2983f09e4 100644 --- a/providers/common/io/src/airflow/providers/common/io/get_provider_info.py +++ b/providers/common/io/src/airflow/providers/common/io/get_provider_info.py @@ -85,7 +85,7 @@ def get_provider_info(): "default": "", }, "state_store_objectstorage_path": { - "description": "Base path on object storage for the task/asset state store backend, in URL format.\nWhen set, StateStoreObjectStorageBackend will persist task and asset state under this\nprefix, organized as //// for tasks and\nassets// for assets.\n", + "description": "Base path on object storage for the task/asset state store backend, in URL format.\nWhen set, StateStoreObjectStorageBackend will persist task and asset state under this\nprefix, organized as //[region=/]//\nfor tasks and assets// for assets.\n", "version_added": "1.8.0", "type": "string", "example": "s3://conn_id@bucket/task-state/", diff --git a/providers/common/io/src/airflow/providers/common/io/state_store/backend.py b/providers/common/io/src/airflow/providers/common/io/state_store/backend.py index f74d04673336c..00e29dff74a8a 100644 --- a/providers/common/io/src/airflow/providers/common/io/state_store/backend.py +++ b/providers/common/io/src/airflow/providers/common/io/state_store/backend.py @@ -25,6 +25,7 @@ import fsspec.utils from airflow.providers.common.compat.sdk import conf +from airflow.providers.common.io.version_compat import AIRFLOW_V_3_4_PLUS if TYPE_CHECKING: from datetime import datetime @@ -79,16 +80,16 @@ def _sanitise_segment(value: str) -> str: return quote(value, safe="") +def _build_task_prefix(scope: TaskScope) -> ObjectStoragePath: + prefix = _get_base_path() / _sanitise_segment(scope.dag_id) / _sanitise_segment(scope.run_id) + if AIRFLOW_V_3_4_PLUS and scope.region_id.int: + prefix /= f"region={scope.region_id}" + return prefix / _sanitise_segment(scope.task_id) + + def _build_task_path(scope: TaskScope, key: str) -> ObjectStoragePath: suffix = _get_compression_suffix() - return ( - _get_base_path() - / _sanitise_segment(scope.dag_id) - / _sanitise_segment(scope.run_id) - / _sanitise_segment(scope.task_id) - / str(scope.map_index) - / f"{_sanitise_segment(key)}{suffix}" - ) + return _build_task_prefix(scope) / str(scope.map_index) / f"{_sanitise_segment(key)}{suffix}" def _build_asset_path(scope: AssetScope, key: str) -> ObjectStoragePath: @@ -163,23 +164,12 @@ def clear( ) -> None: match scope: case TaskScope(): + prefix = _build_task_prefix(scope) if all_map_indices: - prefix = ( - _get_base_path() - / _sanitise_segment(scope.dag_id) - / _sanitise_segment(scope.run_id) - / _sanitise_segment(scope.task_id) - ) for p in prefix.glob("*/*"): p.unlink(missing_ok=True) else: - prefix = ( - _get_base_path() - / _sanitise_segment(scope.dag_id) - / _sanitise_segment(scope.run_id) - / _sanitise_segment(scope.task_id) - / str(scope.map_index) - ) + prefix /= str(scope.map_index) for p in prefix.glob("*"): p.unlink(missing_ok=True) case AssetScope(): diff --git a/providers/common/io/src/airflow/providers/common/io/version_compat.py b/providers/common/io/src/airflow/providers/common/io/version_compat.py index 48d122b669696..a999450478472 100644 --- a/providers/common/io/src/airflow/providers/common/io/version_compat.py +++ b/providers/common/io/src/airflow/providers/common/io/version_compat.py @@ -33,3 +33,4 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) +AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0) diff --git a/providers/common/io/tests/unit/common/io/state_store/test_backend.py b/providers/common/io/tests/unit/common/io/state_store/test_backend.py index f59c2e964413a..131f1c1f92789 100644 --- a/providers/common/io/tests/unit/common/io/state_store/test_backend.py +++ b/providers/common/io/tests/unit/common/io/state_store/test_backend.py @@ -18,10 +18,11 @@ import threading from unittest import mock +from uuid import UUID, uuid4 import pytest -from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS, AIRFLOW_V_3_4_PLUS if not AIRFLOW_V_3_3_PLUS: pytest.skip("Store backend requires Airflow >= 3.3", allow_module_level=True) @@ -73,6 +74,14 @@ def conf_overrides(base_path): @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task state store requires Airflow >= 3.3") class TestPathBuilders: + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Regions require Airflow >= 3.4") + def test_region_path_cannot_collide_with_task_id(self, conf_overrides): + region_id = uuid4() + regional = TaskScope("d", "r", "t", 0, region_id=region_id) + sentinel = TaskScope("d", "r", f"region={region_id}", 0) + assert str(_build_task_path(regional, "key")).endswith(f"d/r/region={region_id}/t/0/key") + assert str(_build_task_path(sentinel, "key")).endswith(f"d/r/region%3D{region_id}/0/key") + def test_build_task_path_segments(self, conf_overrides): scope = TaskScope(dag_id="my_dag", run_id="run_1", task_id="my_task", map_index=-1) path = _build_task_path(scope, "job_id") @@ -161,6 +170,44 @@ def test_set_and_get_task(self, store, task_scope): store.set(task_scope, "k", "hello") assert store.get(task_scope, "k") == "hello" + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Regions require Airflow >= 3.4") + def test_region_writes_and_serialized_references_are_isolated(self, store): + scopes = [TaskScope("d", "r", "t", 0, region_id=region) for region in (UUID(int=0), uuid4())] + refs = [ + store.serialize_task_state_store_to_ref(value={"region": index}, key="k", scope=scope) + for index, scope in enumerate(scopes) + ] + assert refs[0] != refs[1] + assert [store.deserialize_task_state_store_from_ref(ref) for ref in refs] == [ + {"region": 0}, + {"region": 1}, + ] + store.set(scopes[1], "k", "updated") + assert store.get(scopes[1], "k") == "updated" + assert store.deserialize_task_state_store_from_ref(refs[0]) == {"region": 0} + store.delete(scopes[1], "k") + assert store.get(scopes[1], "k") is None + assert store.deserialize_task_state_store_from_ref(refs[0]) == {"region": 0} + + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Regions require Airflow >= 3.4") + @pytest.mark.asyncio + @pytest.mark.parametrize("async_clear", [False, True]) + @pytest.mark.parametrize("all_map_indices", [False, True]) + @pytest.mark.parametrize("clear_sentinel", [False, True]) + async def test_clear_isolates_region_and_index(self, store, async_clear, all_map_indices, clear_sentinel): + regions = [UUID(int=0), uuid4(), uuid4()] + scopes = [TaskScope("d", "r", "t", index, region_id=region) for region in regions for index in (0, 1)] + for scope in scopes: + store.set(scope, "key", str(scope)) + target = TaskScope("d", "r", "t", 0, region_id=regions[0 if clear_sentinel else 1]) + if async_clear: + await store.aclear(target, all_map_indices=all_map_indices) + else: + store.clear(target, all_map_indices=all_map_indices) + for scope in scopes: + removed = scope.region_id == target.region_id and (all_map_indices or scope.map_index == 0) + assert store.get(scope, "key") == (None if removed else str(scope)) + def test_get_missing_returns_none(self, store, task_scope): assert store.get(task_scope, "missing") is None diff --git a/shared/state/src/airflow_shared/state/__init__.py b/shared/state/src/airflow_shared/state/__init__.py index 37f289aedfd1b..524c5fd5f4156 100644 --- a/shared/state/src/airflow_shared/state/__init__.py +++ b/shared/state/src/airflow_shared/state/__init__.py @@ -21,6 +21,7 @@ from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING +from uuid import UUID if TYPE_CHECKING: from datetime import datetime @@ -37,13 +38,19 @@ class TaskScope: ``map_index`` defaults to ``-1`` for non-mapped tasks. For mapped tasks, set it to the actual mapped index. ``get``/``set``/``delete`` always match - on ``(dag_id, run_id, task_id, map_index)`` exactly. + on ``(dag_id, run_id, task_id, region_id, map_index)`` exactly. The zero + region UUID preserves the namespace of legacy callers. """ dag_id: str run_id: str task_id: str map_index: int = -1 + region_id: UUID = UUID(int=0) + + @property + def region_index(self) -> int: + return self.map_index @dataclass(frozen=True) @@ -182,8 +189,8 @@ def clear( For ``TaskScope``: by default, only keys for the exact ``map_index`` on the scope are cleared. When ``all_map_indices=True``, the ``map_index`` filter is - dropped and state is wiped across every mapped instance — for use by external - callers (UI, CLI) only, not from within a running task. + dropped and state is wiped across every mapped instance in the selected region. + This flag is for external callers (UI, CLI), not a running task. For ``AssetScope`` the flag has no effect. """ @@ -231,8 +238,8 @@ async def aclear( For ``TaskScope``: by default, only keys for the exact ``map_index`` on the scope are cleared. When ``all_map_indices=True``, the ``map_index`` filter is - dropped and state is wiped across every mapped instance — for use by external - callers (UI, CLI) only, not from within a running task. + dropped and state is wiped across every mapped instance in the selected region. + This flag is for external callers (UI, CLI), not a running task. For ``AssetScope`` the flag has no effect. ``session`` is optional. If provided, implementations should use it directly. diff --git a/shared/state/tests/state/test_state.py b/shared/state/tests/state/test_state.py index c4b49c7bb3f97..149127725f737 100644 --- a/shared/state/tests/state/test_state.py +++ b/shared/state/tests/state/test_state.py @@ -16,11 +16,26 @@ # under the License. from __future__ import annotations +from uuid import UUID, uuid4 + import pytest from airflow_shared.state import AssetScope, BaseStoreBackend, StoreScope, TaskScope +class TestTaskScope: + def test_legacy_scope_uses_sentinel_region(self): + scope = TaskScope("dag", "run", "task", 3) + assert scope.region_id == UUID(int=0) + assert scope.region_index == 3 + + def test_region_participates_in_scope_identity(self): + first = TaskScope("dag", "run", "task", 3, region_id=uuid4()) + second = TaskScope("dag", "run", "task", 3, region_id=uuid4()) + assert len({first, second}) == 2 + assert first.region_index == second.region_index == 3 + + class TestAssetScope: def test_requires_at_least_one_identifier(self): with pytest.raises(ValueError, match="at least one of"):