fixup! Introduce the (internal/db level) concept of "dynamic regions" for TIs

Normalise MySQL's reflected expression defaults in compare_server_default.

MySQL 8 reports the task_instance.region_id and task_state_store.region_id
defaults through information_schema as
unhex(_utf8mb4\'00000000000000000000000000000000\'), with escaped quotes and
a charset introducer, while the model default renders as
UNHEX('00000000000000000000000000000000'). The two never compared equal, so
Alembic autogenerate reported a modify_default for both columns and
test_has_pending_upgrade_ops failed on MySQL.

Strip the escaping and the introducer in _normalize_mysql_server_default so
both spellings compare equal. Migration 0143 and the model default are
unchanged. MariaDB is expected to reflect the plain unhex('...') form, which
already normalised equal.

Author: ashb

airflow-core/docs/core-concepts/task-state-store.rst

@@ -32,6 +32,8 @@ Task State Store
 
 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 <concepts-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.
 

airflow-core/src/airflow/migrations/utils.py

@@ -31,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"""

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")

airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py

@@ -31,7 +31,7 @@
 import sqlalchemy as sa
 from alembic import op
 
-from airflow.migrations.utils import raise_if_rows_exist
+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
 
@@ -79,15 +79,9 @@ def _replace_unique(table_name, constraint_name, columns):
             """
         op.execute(dedent(sql))
     else:
-        with op.get_context().autocommit_block():
-            foreign_keys = op.get_bind().exec_driver_sql("PRAGMA foreign_keys").scalar_one()
-            op.execute("PRAGMA foreign_keys=OFF")
-            try:
-                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)
-            finally:
-                op.execute(f"PRAGMA foreign_keys={foreign_keys}")
+        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):
@@ -153,63 +147,65 @@ def _configure_index_builds():
 
 def upgrade():
     _configure_index_builds()
-    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)
+    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"),
         )
-    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)
+        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()
-    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")
+    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")

airflow-core/src/airflow/utils/db.py

@@ -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(" ", "")

airflow-core/src/airflow/utils/sqlalchemy.py

@@ -341,7 +341,9 @@ def _compact_uuid_default(element, compiler, **kw):
 @compiles(compact_uuid_default, "mysql")
 @compiles(compact_uuid_default, "mariadb")
 def _compact_uuid_default_mysql(element, compiler, **kw):
-    return f"UNHEX({compiler.render_literal_value(element.value.hex, String())})"
+    # 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):

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

airflow-core/tests/unit/utils/test_sqlalchemy.py

@@ -25,9 +25,10 @@
 
 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
@@ -484,10 +485,21 @@ def test_none_passes_through(self):
     @pytest.mark.parametrize(
         ("dialect", "expected"),
         [
-            pytest.param(mysql.dialect(), "UNHEX('0190a1b2c3d47e5f8a9b0c1d2e3f4a5b')", id="mysql"),
+            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