This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6cd2f6f36d9 Remove migration tests for migrations that have already
landed. (#73953)
6cd2f6f36d9 is described below
commit 6cd2f6f36d9eef481bb8ef7a3c16beb2b8ce1c02
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Wed Sep 30 13:31:26 2026 +0100
Remove migration tests for migrations that have already landed. (#73953)
These tests exist to make sure that the migration was doing the right
thing at the time, but now those have been merged, the value of these
tests is minimal-to-nothing for a couple of reasons.
1. These migrations are almost always write-once. By which we mean that
now they are merged, those migrations are essentially read only. This
isn't always true, but it's the rule.
2. Some (but not all) of these migration tests are badly written, in
that they use the _current_ ORM table definition to test against,
which means these tests need correcting when future ORM changes occur.
Both of these add up to my view that these tests are low value in keeping.
---
...est_0049_remove_pickled_data_from_xcom_table.py | 237 -----------
...t_0055_remove_pickled_data_from_dagrun_table.py | 143 -------
.../test_0080_deadline_callback_migration.py | 268 ------------
.../test_0094_deadline_callback_migration.py | 178 --------
.../test_0117_deadline_interval_json_migration.py | 91 -----
...9_add_indexes_on_serialized_dag_and_dag_code.py | 55 ---
.../migrations/test_0131_lower_case_team_names.py | 138 -------
.../test_0136_fold_task_map_into_xcom.py | 184 ---------
...t_0138_allocate_pending_task_attempt_numbers.py | 448 ---------------------
9 files changed, 1742 deletions(-)
diff --git
a/airflow-core/tests/unit/migrations/test_0049_remove_pickled_data_from_xcom_table.py
b/airflow-core/tests/unit/migrations/test_0049_remove_pickled_data_from_xcom_table.py
deleted file mode 100644
index 66fbd25e31a..00000000000
---
a/airflow-core/tests/unit/migrations/test_0049_remove_pickled_data_from_xcom_table.py
+++ /dev/null
@@ -1,237 +0,0 @@
-#
-# 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.
-
-"""
-Regression tests for migration 0049 (eed27faa34e3) value sanitization.
-
-The 2.x -> 3.x conversion of ``xcom.value`` from pickled bytea to JSON/JSONB
must not choke on
-values that are legal in the pickled blob but illegal in strict JSON/JSONB:
non-finite floats
-(NaN/Infinity/-Infinity) and the U+0000 (NUL) escape. It must leave a value
that already parses
-as JSON untouched, including one that wraps another JSON document with its
interior quotes
-escaped, and it must not corrupt a literal backslash-u-0000 in the data. These
tests run the
-migration's own per-dialect SQL against an isolated table.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import json
-from pathlib import Path
-
-import pytest
-import sqlalchemy as sa
-
-from airflow import settings
-
-from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
-
-# A single backslash, built via chr() so no literal escape appears in the
source.
-_BS = chr(92)
-
-# Row 1: every value class the sanitizer must clean. chr(0) is a real embedded
null byte;
-# json.dumps serializes it to the 6-char NUL escape, which is what the
migration must strip.
-_RAW = json.dumps(
- {"d": "F" + chr(0) + "oo", "a": float("nan"), "b": float("inf"), "c":
float("-inf"), "ok": 1.5}
-)
-_EXPECTED = {"d": "Foo", "a": None, "b": None, "c": None, "ok": 1.5}
-
-# Row 2: a string that literally contains backslash-u-0000 (no null byte). It
serializes to an
-# escaped backslash sequence and must survive unchanged.
-_LITERAL_VALUE = "x" + _BS + "u0000y"
-_LITERAL_RAW = json.dumps({"k": _LITERAL_VALUE})
-_LITERAL_EXPECTED = {"k": _LITERAL_VALUE}
-
-# Rows 3 and 4: a task pushed already-serialized JSON, so the value is a JSON
string wrapping
-# another document with its interior quotes escaped. Both already parse, so
the non-finite
-# rewrite must skip them: quoting the token would close the wrapping string
early and abort the
-# cast, and nulling it would rewrite data the migration has no reason to touch.
-_INNER = json.dumps({"amount": 604441.0, "commission": float("nan"), "rate":
float("-inf")})
-_ESCAPED_RAW = json.dumps(_INNER)
-_NESTED_RAW = json.dumps({"report": _INNER})
-
-# Row 5: invalid at the top level and wrapping an escaped document. The
rewrite has to run, so
-# the inner tokens go too; ``null`` keeps the result parseable where a quote
would not.
-_MIXED_RAW = json.dumps({"top": float("nan"), "report": _INNER})
-_MIXED_EXPECTED = {
- "top": None,
- "report": json.dumps({"amount": 604441.0, "commission": None, "rate":
None}),
-}
-
-# Row 6: invalid only because of the NUL escape, and wrapping an escaped
document. Proves the
-# strip runs before the rewrite: stripping first makes the value parse, so the
inner NaN is
-# preserved. Guarding before stripping would send this row through the rewrite
and corrupt it.
-_ORDER_RAW = json.dumps({"n": "x" + chr(0) + "y", "report": _INNER})
-_ORDER_EXPECTED = {"n": "xy", "report": _INNER}
-
-# Migration filenames start with a digit so they cannot be imported via the
normal import
-# system; load the module by file path instead.
-_MIGRATION_PATH = (
- Path(AIRFLOW_CORE_SOURCES_PATH)
- /
"airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py"
-)
-_spec = importlib.util.spec_from_file_location("migration_0049",
_MIGRATION_PATH)
-_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
-_spec.loader.exec_module(_migration) # type: ignore[union-attr]
-
-_TABLE = "_test_xcom_sanitize"
-
-# id -> serialized value, inserted into the isolated table by every dialect
test.
-_ROWS = {
- 1: _RAW,
- 2: _LITERAL_RAW,
- 3: _ESCAPED_RAW,
- 4: _NESTED_RAW,
- 5: _MIXED_RAW,
- 6: _ORDER_RAW,
-}
-
-
-def _assert_sanitized(rows: dict[int, str]) -> None:
- """Check the sanitized text of every row. Rows 3, 4 and 6 compare the
inner document as an
- exact string, so dropped whitespace inside it fails here rather than
passing a loads() check.
- """
- assert json.loads(rows[1]) == _EXPECTED
- assert json.loads(rows[2]) == _LITERAL_EXPECTED
- assert json.loads(rows[3]) == _INNER
- assert json.loads(rows[4]) == {"report": _INNER}
- assert json.loads(rows[5]) == _MIXED_EXPECTED
- assert json.loads(rows[6]) == _ORDER_EXPECTED
- # Rows 3 and 4 already parsed, so they must be byte-identical to what was
stored.
- assert rows[3] == _ESCAPED_RAW
- assert rows[4] == _NESTED_RAW
-
-
-def _sqlite_sanitized(json1: bool = True) -> dict[int, str]:
- engine = sa.create_engine("sqlite://")
- with engine.begin() as conn:
- conn.execute(sa.text(f"CREATE TABLE {_TABLE} (id INTEGER PRIMARY KEY,
value BLOB)"))
- for row_id, value in _ROWS.items():
- conn.execute(
- sa.text(f"INSERT INTO {_TABLE} (id, value) VALUES (:i, :v)"),
- {"i": row_id, "v": value.encode("utf-8")},
- )
- for stmt in _migration._xcom_sqlite_sanitize_statements(_TABLE,
json1=json1):
- conn.execute(sa.text(stmt))
- # json(...) mirrors the migration's own conversion and raises if still
invalid JSON.
- # It also re-serializes, so the assertions read the stored text
instead.
- rows = conn.execute(
- sa.text(f"SELECT id, CAST(value AS TEXT), json(CAST(value AS
TEXT)) FROM {_TABLE}")
- ).all()
- return {row[0]: row[1] for row in rows}
-
-
-def test_sqlite_sanitize():
- """SQLite branch: real sanitize SQL on an in-memory db.
Backend-independent."""
- _assert_sanitized(_sqlite_sanitized())
-
-
-def test_sqlite_sanitize_without_json1():
- """Without JSON1 there is no guard, so already-valid values are rewritten
too. The result
- still has to be valid JSON, which is what keeps the migration completing
on old builds.
- """
- rows = _sqlite_sanitized(json1=False)
- assert json.loads(rows[1]) == _EXPECTED
- assert json.loads(rows[3]) == _MIXED_EXPECTED["report"]
-
-
-def test_sqlite_has_json1_probe():
- """The probe reports True on a build with JSON1 and swallows only the
missing-function error."""
- engine = sa.create_engine("sqlite://")
- with engine.connect() as conn:
- assert _migration._sqlite_has_json1(conn) is True
-
- class _NoJson1:
- def execute(self, *args, **kwargs):
- raise sa.exc.OperationalError(
- "SELECT json_valid('{}')", {}, Exception("no such function:
json_valid")
- )
-
- assert _migration._sqlite_has_json1(_NoJson1()) is False
-
-
-def test_sqlite_has_json1_probe_propagates_other_errors():
- """A failure that is not a missing function must surface instead of
downgrading the sanitize."""
-
- class _Broken:
- def execute(self, *args, **kwargs):
- raise sa.exc.InterfaceError("SELECT json_valid('{}')", {},
Exception("connection gone"))
-
- with pytest.raises(sa.exc.InterfaceError):
- _migration._sqlite_has_json1(_Broken())
-
-
[email protected]_test
-class TestPostgresSanitize:
- @pytest.mark.backend("postgres")
- def test_nul_and_nan_block_jsonb_cast_until_sanitized(self):
- drop = f"DROP TABLE IF EXISTS {_TABLE}"
- cast = f"SELECT CAST(CONVERT_FROM(value, 'UTF8') AS JSONB) FROM
{_TABLE}"
- with settings.engine.begin() as conn:
- conn.execute(sa.text(drop))
- conn.execute(sa.text(f"CREATE TABLE {_TABLE} (id int PRIMARY KEY,
value bytea)"))
- for row_id, value in _ROWS.items():
- conn.execute(
- sa.text(f"INSERT INTO {_TABLE} VALUES (:i, convert_to(:v,
'UTF8'))"),
- {"i": row_id, "v": value},
- )
- try:
- # Before sanitizing, the JSONB cast fails (the reported upgrade
failure).
- with settings.engine.connect() as conn:
- with pytest.raises(sa.exc.DataError):
- conn.execute(sa.text(cast)).all()
- conn.rollback()
- # pg_temp is per-session, so the helper and the UPDATE share one
connection.
- with settings.engine.begin() as conn:
- for stmt in _migration._xcom_pg_sanitize_statements(_TABLE):
- conn.execute(sa.text(stmt))
- conn.execute(sa.text(cast)).all()
- rows = dict(
- conn.execute(sa.text(f"SELECT id, CONVERT_FROM(value,
'UTF8') FROM {_TABLE}")).all()
- )
- _assert_sanitized(rows)
- finally:
- with settings.engine.begin() as conn:
- conn.execute(sa.text(drop))
-
-
[email protected]_test
-class TestMysqlSanitize:
- @pytest.mark.backend("mysql")
- def test_sanitize_allows_json_cast(self):
- drop = f"DROP TABLE IF EXISTS {_TABLE}"
- cast = f"SELECT CAST(CONVERT(value USING utf8mb4) AS JSON) FROM
{_TABLE}"
- with settings.engine.begin() as conn:
- conn.execute(sa.text(drop))
- conn.execute(sa.text(f"CREATE TABLE {_TABLE} (id int PRIMARY KEY,
value LONGBLOB)"))
- for row_id, value in _ROWS.items():
- conn.execute(
- sa.text(f"INSERT INTO {_TABLE} VALUES (:i, CONVERT(:v
USING utf8mb4))"),
- {"i": row_id, "v": value},
- )
- try:
- with settings.engine.begin() as conn:
- for stmt in _migration._xcom_mysql_sanitize_statements(_TABLE):
- conn.execute(sa.text(stmt))
- conn.execute(sa.text(cast)).all() # must not raise (bare NaN
would be rejected)
- rows = dict(
- conn.execute(sa.text(f"SELECT id, CONVERT(value USING
utf8mb4) FROM {_TABLE}")).all()
- )
- _assert_sanitized(rows)
- finally:
- with settings.engine.begin() as conn:
- conn.execute(sa.text(drop))
diff --git
a/airflow-core/tests/unit/migrations/test_0055_remove_pickled_data_from_dagrun_table.py
b/airflow-core/tests/unit/migrations/test_0055_remove_pickled_data_from_dagrun_table.py
deleted file mode 100644
index aeca8563a01..00000000000
---
a/airflow-core/tests/unit/migrations/test_0055_remove_pickled_data_from_dagrun_table.py
+++ /dev/null
@@ -1,143 +0,0 @@
-#
-# 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.
-
-"""
-Unit tests for migration 0055 (e39a26ac59f6) conf sanitization.
-
-The 2.x -> 3.x conversion of ``dag_run.conf`` from pickled bytea to JSON/JSONB
happens
-Python-side (``json.dumps`` + a per-row insert). ``_json_safe`` quotes
non-finite floats
-and strips embedded NUL characters so confs carrying those values are
preserved instead of
-being dropped by the migration's per-row error handler. NUL is handled on the
object (not
-on the dumped text) so a genuinely escaped backslash sequence is not
corrupted. These are
-pure-Python tests; no database is required.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import json
-from collections import OrderedDict
-from pathlib import Path
-
-import pytest
-
-from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
-
-# A single backslash, built via chr() so no literal escape appears in the
source.
-_BS = chr(92)
-# The 6-char escape json.dumps emits for an embedded null byte.
-_NUL_ESCAPE = _BS + "u0000"
-
-_MIGRATION_PATH = (
- Path(AIRFLOW_CORE_SOURCES_PATH)
- /
"airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py"
-)
-_spec = importlib.util.spec_from_file_location("migration_0055",
_MIGRATION_PATH)
-_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
-_spec.loader.exec_module(_migration) # type: ignore[union-attr]
-
-_json_safe = _migration._json_safe
-
-
[email protected](
- ("value", "expected"),
- [
- (float("nan"), "NaN"),
- (float("inf"), "Infinity"),
- (float("-inf"), "-Infinity"),
- (1.5, 1.5),
- (0.0, 0.0),
- (-2.0, -2.0),
- ("plain", "plain"),
- (42, 42),
- (None, None),
- (True, True),
- ],
-)
-def test_json_safe_scalars(value, expected):
- assert _json_safe(value) == expected
-
-
-def test_json_safe_strips_null_bytes_in_strings():
- assert _json_safe("foo" + chr(0) + "bar") == "foobar"
- assert _json_safe(chr(0)) == ""
-
-
-def test_json_safe_preserves_literal_backslash_u_text():
- """A string literally containing backslash-u-0000 (no null byte) must
survive intact."""
- literal = "foo" + _NUL_ESCAPE + "bar"
- assert _json_safe(literal) == literal
- # and round-trips through json without corruption
- assert json.loads(json.dumps(_json_safe({"k": literal}))) == {"k": literal}
-
-
-def test_json_safe_recurses_into_mappings_and_sequences():
- data = OrderedDict(
- [
- ("f", float("nan")),
- ("lst", [float("inf"), 1, {"deep": float("-inf")}]),
- ("tpl", (float("nan"), 2)),
- ("nul" + chr(0), "v" + chr(0)),
- ("keep", 3.14),
- ]
- )
- assert _json_safe(data) == {
- "f": "NaN",
- "lst": ["Infinity", 1, {"deep": "-Infinity"}],
- "tpl": ["NaN", 2], # tuples normalize to lists, like json.dumps would
- "nul": "v", # NUL stripped from both key and value
- "keep": 3.14,
- }
-
-
-def test_json_safe_does_not_explode_strings_into_chars():
- assert _json_safe("hello") == "hello"
-
-
-def _reject_constant(token):
- raise AssertionError(f"non-finite token survived sanitization: {token!r}")
-
-
-def test_full_pipeline_yields_strict_valid_json():
- """Mirror the migration's exact serialization:
json.dumps(_json_safe(...))."""
- original = {
- "d": "F" + chr(0) + "oo", # real embedded null byte
- "lit": "x" + _NUL_ESCAPE + "y", # literal backslash-u-0000 text, must
survive
- "a": float("nan"),
- "b": float("inf"),
- "c": float("-inf"),
- "ok": 1.5,
- }
- json_data = json.dumps(_json_safe(original))
-
- # parse_constant fires on any surviving bare NaN/Infinity/-Infinity token.
- parsed = json.loads(json_data, parse_constant=_reject_constant)
- assert parsed == {
- "d": "Foo",
- "lit": "x" + _NUL_ESCAPE + "y",
- "a": "NaN",
- "b": "Infinity",
- "c": "-Infinity",
- "ok": 1.5,
- }
-
-
-def test_finite_floats_are_untouched():
- original = {"x": 1.25, "y": [0.0, -3.5], "z": 1000000.0}
- json_data = json.dumps(_json_safe(original))
- assert json.loads(json_data, parse_constant=_reject_constant) == original
diff --git
a/airflow-core/tests/unit/migrations/test_0080_deadline_callback_migration.py
b/airflow-core/tests/unit/migrations/test_0080_deadline_callback_migration.py
deleted file mode 100644
index c777d8418b1..00000000000
---
a/airflow-core/tests/unit/migrations/test_0080_deadline_callback_migration.py
+++ /dev/null
@@ -1,268 +0,0 @@
-#
-# 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.
-
-"""
-Regression tests for migration 0080 (808787349f22):
-upgrade() and downgrade() must correctly migrate existing deadline rows
-without raising NotNullViolation.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import json
-import uuid
-from pathlib import Path
-from unittest import mock
-
-import sqlalchemy as sa
-from alembic.migration import MigrationContext
-from alembic.operations import Operations
-
-from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
-
-# Migration filenames start with a digit so they cannot be imported via the
-# normal import system; load the module by file path instead.
-_MIGRATION_PATH = (
- Path(AIRFLOW_CORE_SOURCES_PATH)
- /
"airflow/migrations/versions/0080_3_1_0_modify_deadline_callback_schema.py"
-)
-_spec = importlib.util.spec_from_file_location("migration_0080",
_MIGRATION_PATH)
-_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
-_spec.loader.exec_module(_migration) # type: ignore[union-attr]
-
-upgrade = _migration.upgrade
-downgrade = _migration.downgrade
-_ASYNC_CALLBACK_CLASSNAME = _migration._ASYNC_CALLBACK_CLASSNAME
-
-_PRE_0080_DDL = """
-CREATE TABLE deadline (
- id TEXT PRIMARY KEY,
- dagrun_id INTEGER NOT NULL,
- deadline_time TEXT NOT NULL,
- callback TEXT NOT NULL,
- callback_kwargs TEXT,
- created_at TEXT,
- last_updated_at TEXT
-)
-"""
-
-_POST_0080_DDL = """
-CREATE TABLE deadline (
- id TEXT PRIMARY KEY,
- dagrun_id INTEGER NOT NULL,
- deadline_time TEXT NOT NULL,
- callback TEXT NOT NULL,
- created_at TEXT,
- last_updated_at TEXT
-)
-"""
-
-
-def _make_engine_pre_0080():
- """Return an in-memory SQLite engine with the pre-0080 deadline schema."""
- engine = sa.create_engine("sqlite:///:memory:")
- with engine.connect() as conn:
- conn.execute(sa.text(_PRE_0080_DDL))
- conn.commit()
- return engine
-
-
-def _make_engine_post_0080():
- """Return an in-memory SQLite engine with the post-0080 deadline schema."""
- engine = sa.create_engine("sqlite:///:memory:")
- with engine.connect() as conn:
- conn.execute(sa.text(_POST_0080_DDL))
- conn.commit()
- return engine
-
-
-def _run_upgrade(engine):
- # alembic.context is a proxy that is only populated when running through
- # Alembic's full migration runner (alembic upgrade). When calling the
- # migration function directly in a test we must mock it so that the
- # is_offline_mode() guard does not raise AttributeError.
- with engine.begin() as conn:
- with Operations.context(MigrationContext.configure(conn)):
- with mock.patch.object(_migration, "context") as mock_ctx:
- mock_ctx.is_offline_mode.return_value = False
- upgrade()
-
-
-def _run_downgrade(engine):
- with engine.begin() as conn:
- with Operations.context(MigrationContext.configure(conn)):
- with mock.patch.object(_migration, "context") as mock_ctx:
- mock_ctx.is_offline_mode.return_value = False
- downgrade()
-
-
-def _read_deadline(engine):
- with engine.connect() as conn:
- return conn.execute(sa.text("SELECT * FROM deadline")).mappings().all()
-
-
-class TestMigration0080Upgrade:
- def test_upgrade_empty_table(self):
- """Upgrade on an empty table must not raise."""
- engine = _make_engine_pre_0080()
- _run_upgrade(engine)
- rows = _read_deadline(engine)
- assert rows == []
-
- def test_upgrade_migrates_existing_row(self):
- """Upgrade converts VARCHAR callback + JSON kwargs to the expected
JSON envelope."""
- engine = _make_engine_pre_0080()
- row_id = str(uuid.uuid4())
- with engine.begin() as conn:
- conn.execute(
- sa.text(
- "INSERT INTO deadline (id, dagrun_id, deadline_time,
callback, callback_kwargs)"
- " VALUES (:id, 1, '2025-01-01', :cb, :kw)"
- ),
- {"id": row_id, "cb": "mymodule.my_callback", "kw":
json.dumps({"key": "val"})},
- )
-
- _run_upgrade(engine)
-
- rows = _read_deadline(engine)
- assert len(rows) == 1
- cb = rows[0]["callback"]
- if isinstance(cb, str):
- cb = json.loads(cb)
- assert cb["__classname__"] == _ASYNC_CALLBACK_CLASSNAME
- assert cb["__version__"] == 0
- assert cb["__data__"]["path"] == "mymodule.my_callback"
- assert cb["__data__"]["kwargs"] == {"key": "val"}
- assert "callback_kwargs" not in rows[0]
-
- def test_upgrade_null_kwargs_defaults_to_empty_dict(self):
- """Upgrade with NULL callback_kwargs must produce an empty dict in the
envelope."""
- engine = _make_engine_pre_0080()
- row_id = str(uuid.uuid4())
- with engine.begin() as conn:
- conn.execute(
- sa.text(
- "INSERT INTO deadline (id, dagrun_id, deadline_time,
callback, callback_kwargs)"
- " VALUES (:id, 1, '2025-01-01', :cb, NULL)"
- ),
- {"id": row_id, "cb": "mymodule.my_callback"},
- )
-
- _run_upgrade(engine)
-
- rows = _read_deadline(engine)
- cb = rows[0]["callback"]
- if isinstance(cb, str):
- cb = json.loads(cb)
- assert cb["__data__"]["kwargs"] == {}
-
- def test_upgrade_exact_batch_boundary(self, monkeypatch):
- """Rows == batch_size must force a second iteration that returns 0
rows and exits cleanly."""
- # Force a small batch_size so 2 inserted rows == batch_size exactly.
- monkeypatch.setattr(_migration.conf, "getint", lambda *a, **kw: 2)
- engine = _make_engine_pre_0080()
- with engine.begin() as conn:
- for i in range(2):
- conn.execute(
- sa.text(
- "INSERT INTO deadline (id, dagrun_id, deadline_time,
callback, callback_kwargs)"
- " VALUES (:id, 1, '2025-01-01', :cb, :kw)"
- ),
- {"id": str(uuid.uuid4()), "cb": f"mod.cb_{i}", "kw":
json.dumps({"i": i})},
- )
-
- _run_upgrade(engine)
-
- rows = _read_deadline(engine)
- assert len(rows) == 2
- paths = sorted(
- (json.loads(r["callback"]) if isinstance(r["callback"], str) else
r["callback"])["__data__"][
- "path"
- ]
- for r in rows
- )
- assert paths == ["mod.cb_0", "mod.cb_1"]
-
-
-class TestMigration0080Downgrade:
- def test_downgrade_empty_table(self):
- """Downgrade on an empty table must not raise."""
- engine = _make_engine_post_0080()
- _run_downgrade(engine)
- rows = _read_deadline(engine)
- assert rows == []
-
- def test_downgrade_restores_existing_row(self):
- """Downgrade extracts path and kwargs back from the JSON envelope."""
- engine = _make_engine_post_0080()
- row_id = str(uuid.uuid4())
- callback_json = json.dumps(
- {
- "__data__": {"path": "mymodule.my_callback", "kwargs": {"key":
"val"}},
- "__classname__": _ASYNC_CALLBACK_CLASSNAME,
- "__version__": 0,
- }
- )
- with engine.begin() as conn:
- conn.execute(
- sa.text(
- "INSERT INTO deadline (id, dagrun_id, deadline_time,
callback)"
- " VALUES (:id, 1, '2025-01-01', :cb)"
- ),
- {"id": row_id, "cb": callback_json},
- )
-
- _run_downgrade(engine)
-
- rows = _read_deadline(engine)
- assert len(rows) == 1
- assert rows[0]["callback"] == "mymodule.my_callback"
- kw = rows[0]["callback_kwargs"]
- if isinstance(kw, str):
- kw = json.loads(kw)
- assert kw == {"key": "val"}
-
-
-class TestMigration0080RoundTrip:
- def test_round_trip_preserves_data(self):
- """Upgrade followed by downgrade preserves the original callback
path."""
- engine = _make_engine_pre_0080()
- row_id = str(uuid.uuid4())
- original_path = "mymodule.my_callback"
- original_kwargs = {"x": 1}
-
- with engine.begin() as conn:
- conn.execute(
- sa.text(
- "INSERT INTO deadline (id, dagrun_id, deadline_time,
callback, callback_kwargs)"
- " VALUES (:id, 1, '2025-01-01', :cb, :kw)"
- ),
- {"id": row_id, "cb": original_path, "kw":
json.dumps(original_kwargs)},
- )
-
- _run_upgrade(engine)
- _run_downgrade(engine)
-
- rows = _read_deadline(engine)
- assert len(rows) == 1
- assert rows[0]["callback"] == original_path
- kw = rows[0]["callback_kwargs"]
- if isinstance(kw, str):
- kw = json.loads(kw)
- assert kw == original_kwargs
diff --git
a/airflow-core/tests/unit/migrations/test_0094_deadline_callback_migration.py
b/airflow-core/tests/unit/migrations/test_0094_deadline_callback_migration.py
deleted file mode 100644
index 3d848bf843a..00000000000
---
a/airflow-core/tests/unit/migrations/test_0094_deadline_callback_migration.py
+++ /dev/null
@@ -1,178 +0,0 @@
-#
-# 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.
-
-"""
-Regression tests for migration 0094 (e812941398f4).
-
-These tests focus on the defensive NULL-callback path: legacy MySQL
-deployments that ran the original (buggy) 0080 left ``deadline.callback``
-rows as NULL. 0094 must heal those rows instead of crashing on
-``json.loads(None)``.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import json
-import uuid
-from pathlib import Path
-
-import sqlalchemy as sa
-
-from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
-
-_MIGRATION_PATH = (
- Path(AIRFLOW_CORE_SOURCES_PATH)
- /
"airflow/migrations/versions/0094_3_2_0_replace_deadline_inline_callback_with_fkey.py"
-)
-_spec = importlib.util.spec_from_file_location("migration_0094",
_MIGRATION_PATH)
-_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
-_spec.loader.exec_module(_migration) # type: ignore[union-attr]
-
-
-# Minimal post-0080 / pre-0094 schema. 0094 adds ``missed`` and ``callback_id``
-# itself before invoking ``_upgrade_mysql_sqlite``; we mimic that here so we
-# can call the inner helper directly without driving the full alembic chain.
-_POST_0080_DDL = [
- """
- CREATE TABLE dag_run (
- id INTEGER PRIMARY KEY,
- dag_id TEXT NOT NULL
- )
- """,
- """
- CREATE TABLE deadline (
- id TEXT PRIMARY KEY,
- dagrun_id INTEGER NOT NULL,
- deadline_time TEXT NOT NULL,
- callback TEXT,
- callback_state TEXT,
- trigger_id INTEGER,
- callback_id TEXT,
- missed BOOLEAN
- )
- """,
- """
- CREATE TABLE callback (
- id TEXT PRIMARY KEY,
- type TEXT NOT NULL,
- fetch_method TEXT NOT NULL,
- data TEXT NOT NULL,
- state TEXT NOT NULL,
- priority_weight INTEGER NOT NULL,
- created_at TEXT NOT NULL
- )
- """,
-]
-
-
-def _make_engine():
- engine = sa.create_engine("sqlite:///:memory:")
- with engine.connect() as conn:
- for ddl in _POST_0080_DDL:
- conn.execute(sa.text(ddl))
- conn.commit()
- return engine
-
-
-def _insert_dagrun(conn, dagrun_id: int = 1, dag_id: str = "test_dag"):
- conn.execute(
- sa.text("INSERT INTO dag_run (id, dag_id) VALUES (:id, :dag_id)"),
- {"id": dagrun_id, "dag_id": dag_id},
- )
-
-
-def _insert_deadline(conn, deadline_id: str, callback, callback_state: str |
None = None):
- conn.execute(
- sa.text(
- "INSERT INTO deadline (id, dagrun_id, deadline_time, callback,
callback_state)"
- " VALUES (:id, 1, '2025-01-01', :cb, :state)"
- ),
- {"id": deadline_id, "cb": callback, "state": callback_state},
- )
-
-
-class TestMigration0094NullCallbackRepair:
- """A NULL callback row from a buggy 0080 must not crash 0094's upgrade."""
-
- def test_null_callback_does_not_crash(self):
- engine = _make_engine()
- # `_upgrade_mysql_sqlite` declares ``id`` as ``sa.Uuid()``; on SQLite
the
- # write path emits the hex (no-dash) form. Insert IDs in that same
form so
- # the UPDATE in the migration loop matches the row we created.
- deadline_id = uuid.uuid4().hex
- with engine.begin() as conn:
- _insert_dagrun(conn)
- _insert_deadline(conn, deadline_id, callback=None)
-
- # _upgrade_mysql_sqlite reads from `deadline` and writes to `callback`;
- # it does not depend on alembic's batch_alter_table prelude.
- with engine.begin() as conn:
- _migration._upgrade_mysql_sqlite(conn, batch_size=10)
-
- with engine.connect() as conn:
- deadline_rows = conn.execute(sa.text("SELECT * FROM
deadline")).mappings().all()
- callback_rows = conn.execute(sa.text("SELECT * FROM
callback")).mappings().all()
-
- assert len(deadline_rows) == 1
- assert len(callback_rows) == 1
- assert deadline_rows[0]["callback_id"] == callback_rows[0]["id"]
- assert deadline_rows[0]["missed"] == 0 # SQLite: False -> 0
-
- cb_data = json.loads(callback_rows[0]["data"])
- assert cb_data["path"] == ""
- assert cb_data["kwargs"] == {}
- assert cb_data["dag_id"] == "test_dag"
-
- def test_mixed_null_and_valid_callbacks(self):
- """A batch with both NULL and well-formed rows must migrate both."""
- engine = _make_engine()
- null_id = uuid.uuid4().hex
- valid_id = uuid.uuid4().hex
- valid_callback = json.dumps(
- {
- "__data__": {"path": "mymodule.cb", "kwargs": {"k": "v"}},
- "__classname__":
"airflow.sdk.definitions.deadline.AsyncCallback",
- "__version__": 0,
- }
- )
- with engine.begin() as conn:
- _insert_dagrun(conn)
- _insert_deadline(conn, null_id, callback=None)
- _insert_deadline(conn, valid_id, callback=valid_callback)
-
- with engine.begin() as conn:
- _migration._upgrade_mysql_sqlite(conn, batch_size=10)
-
- with engine.connect() as conn:
- rows = (
- conn.execute(
- sa.text(
- "SELECT d.id AS deadline_id, c.data"
- " FROM deadline d JOIN callback c ON d.callback_id =
c.id"
- )
- )
- .mappings()
- .all()
- )
-
- by_id = {r["deadline_id"]: json.loads(r["data"]) for r in rows}
- assert by_id[null_id]["path"] == ""
- assert by_id[null_id]["kwargs"] == {}
- assert by_id[valid_id]["path"] == "mymodule.cb"
- assert by_id[valid_id]["kwargs"] == {"k": "v"}
diff --git
a/airflow-core/tests/unit/migrations/test_0117_deadline_interval_json_migration.py
b/airflow-core/tests/unit/migrations/test_0117_deadline_interval_json_migration.py
deleted file mode 100644
index 289e83797c3..00000000000
---
a/airflow-core/tests/unit/migrations/test_0117_deadline_interval_json_migration.py
+++ /dev/null
@@ -1,91 +0,0 @@
-#
-# 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.
-
-"""
-Regression test for migration 0117 (8812eb67b63c) on MySQL.
-
-The downgrade must convert ``deadline_alert.interval`` from JSON back to FLOAT
without
-raising ``ER_INVALID_JSON_TEXT`` (3140). The failure only reproduces with at
least one
-row present, so this seeds rows and runs the migration's own conversion SQL.
It is a
-no-op on SQLite/PostgreSQL, which take different code paths.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-from pathlib import Path
-
-import pytest
-import sqlalchemy as sa
-
-from airflow import settings
-
-from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
-
-pytestmark = pytest.mark.db_test
-
-# Migration filenames start with a digit so they cannot be imported via the
-# normal import system; load the module by file path instead.
-_MIGRATION_PATH = (
- Path(AIRFLOW_CORE_SOURCES_PATH)
- /
"airflow/migrations/versions/0117_3_3_0_change_deadline_interval_to_json.py"
-)
-_spec = importlib.util.spec_from_file_location("migration_0117",
_MIGRATION_PATH)
-_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
-_spec.loader.exec_module(_migration) # type: ignore[union-attr]
-
-# Isolated table so we run the real conversion SQL without seeding the live
deadline_alert
-# table (which has FK/NOT NULL columns).
-_TABLE = "_test_deadline_interval_dg"
-
-# A serialized timedelta as written by the 0117 upgrade.
-_WRAPPED_TIMEDELTA = '{"__classname__": "datetime.timedelta", "__version__":
2, "__data__": 300.0}'
-
-
-class TestMigration0117Downgrade:
- @pytest.mark.backend("mysql")
- def
test_mysql_downgrade_interval_value_update_does_not_reject_on_json_column(self):
- """The downgrade value-conversion UPDATE must not raise 3140, and must
round-trip to FLOAT."""
- create = f"CREATE TABLE {_TABLE} (id INT PRIMARY KEY, `interval` JSON
NOT NULL)"
- drop = f"DROP TABLE IF EXISTS {_TABLE}"
-
- with settings.engine.begin() as conn:
- conn.execute(sa.text(drop))
- conn.execute(sa.text(create))
- conn.execute(
- sa.text(f"INSERT INTO {_TABLE} (id, `interval`) VALUES (1,
:v)"),
- {"v": _WRAPPED_TIMEDELTA},
- )
- conn.execute(
- sa.text(f"INSERT INTO {_TABLE} (id, `interval`) VALUES (2,
CAST(:v AS JSON))"),
- {"v": "60.0"},
- )
-
- try:
- with settings.engine.begin() as conn:
- # Column is still JSON here; this UPDATE must not raise 3140.
The retype casts.
-
conn.execute(sa.text(_migration._mysql_downgrade_interval_value_sql(_TABLE)))
- conn.execute(sa.text(f"ALTER TABLE {_TABLE} MODIFY `interval`
FLOAT NOT NULL"))
-
- with settings.engine.connect() as conn:
- rows = dict(conn.execute(sa.text(f"SELECT id, `interval` FROM
{_TABLE}")).all())
-
- assert rows == {1: 300.0, 2: 60.0}
- finally:
- with settings.engine.begin() as conn:
- conn.execute(sa.text(drop))
diff --git
a/airflow-core/tests/unit/migrations/test_0129_add_indexes_on_serialized_dag_and_dag_code.py
b/airflow-core/tests/unit/migrations/test_0129_add_indexes_on_serialized_dag_and_dag_code.py
deleted file mode 100644
index 4774a7711dd..00000000000
---
a/airflow-core/tests/unit/migrations/test_0129_add_indexes_on_serialized_dag_and_dag_code.py
+++ /dev/null
@@ -1,55 +0,0 @@
-#
-# 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
-
-import pytest
-import sqlalchemy as sa
-
-from airflow import settings
-from airflow.utils.db import downgrade, upgradedb
-
-pytestmark = pytest.mark.db_test
-
-_REVISION = "3c525f44bea8"
-_DOWN_REVISION = "b2f1a9c7d4e0"
-
-_EXPECTED_INDEXES = {
- "idx_serialized_dag_dag_id_created_at": ("serialized_dag", ["dag_id",
"created_at"]),
- "idx_dag_code_dag_id_last_updated": ("dag_code", ["dag_id",
"last_updated"]),
-}
-
-
-class TestMigration0124AddDagIdIndexes:
- @pytest.fixture(autouse=True)
- def _restore_head(self):
- yield
- upgradedb()
-
- @staticmethod
- def _get_indexes(table):
- with settings.engine.connect() as conn:
- return {ix["name"]: ix["column_names"] for ix in
sa.inspect(conn).get_indexes(table)}
-
- def test_upgrade_creates_indexes_and_downgrade_drops_them(self):
- downgrade(to_revision=_DOWN_REVISION)
- for index, (table, _) in _EXPECTED_INDEXES.items():
- assert index not in self._get_indexes(table)
-
- upgradedb(to_revision=_REVISION)
- for index, (table, columns) in _EXPECTED_INDEXES.items():
- assert self._get_indexes(table).get(index) == columns
diff --git
a/airflow-core/tests/unit/migrations/test_0131_lower_case_team_names.py
b/airflow-core/tests/unit/migrations/test_0131_lower_case_team_names.py
deleted file mode 100644
index 15008679853..00000000000
--- a/airflow-core/tests/unit/migrations/test_0131_lower_case_team_names.py
+++ /dev/null
@@ -1,138 +0,0 @@
-#
-# 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 migration 0131 (c7f0a5d2e9b4), which lower-cases stored team names.
-
-``team.name`` is a primary key whose foreign keys do not cascade on update, so
the rename has
-to carry every referring row with it. A row left behind would either break the
constraint or
-strand that team's Connections, Variables and Pools under a name nothing
resolves.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-from pathlib import Path
-
-import pytest
-import sqlalchemy as sa
-from alembic import command
-
-from airflow import settings
-from airflow.models import Connection, Pool, Variable
-from airflow.models.team import Team
-from airflow.utils.db import _get_alembic_config
-from airflow.utils.session import create_session
-
-from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
-
-pytestmark = pytest.mark.db_test
-
-# Migration filenames start with a digit so they cannot be imported via the
normal import
-# system; load the module by file path instead.
-_MIGRATION_PATH = (
- Path(AIRFLOW_CORE_SOURCES_PATH) /
"airflow/migrations/versions/0131_3_4_0_lower_case_team_names.py"
-)
-_spec = importlib.util.spec_from_file_location("migration_0131",
_MIGRATION_PATH)
-_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
-_spec.loader.exec_module(_migration) # type: ignore[union-attr]
-
-_OLD = "Data-Eng"
-_NEW = "data-eng"
-
-
-# Addressed through the mapped columns rather than by name: ``variable.key``
and
-# ``slot_pool.pool`` are reserved words on MySQL, which only quotes them
correctly when
-# SQLAlchemy renders the identifier for the dialect.
-_REFERRING_ROWS = (
- (Connection, Connection.conn_id, "c_0131"),
- (Variable, Variable.key, "v_0131"),
- (Pool, Pool.pool, "p_0131"),
-)
-
-
-def _team_names(conn) -> set[str]:
- return set(conn.scalars(sa.select(Team.name)))
-
-
-class TestMigration0131:
- def test_rename_carries_referring_rows_with_the_team(self):
- with create_session() as session:
- session.add(Team(name=_OLD))
- session.flush()
- session.add_all(
- [
- Connection(conn_id="c_0131", conn_type="http",
team_name=_OLD),
- Variable(key="v_0131", val="x", team_name=_OLD),
- Pool(pool="p_0131", slots=1, description="",
include_deferred=False, team_name=_OLD),
- ]
- )
- session.commit()
-
- try:
- with settings.engine.begin() as conn:
- for statement in
_migration.build_lower_casing_statements(conn.dialect.name):
- conn.execute(statement)
-
- with settings.engine.connect() as conn:
- assert _OLD not in _team_names(conn)
- assert _NEW in _team_names(conn)
- for model, key_column, key in _REFERRING_ROWS:
- team_name =
conn.scalar(sa.select(model.team_name).where(key_column == key))
- assert team_name == _NEW, model.__tablename__
- finally:
- with settings.engine.begin() as conn:
- for model, key_column, key in _REFERRING_ROWS:
- conn.execute(sa.delete(model).where(key_column == key))
- conn.execute(sa.delete(Team).where(Team.name.in_([_OLD,
_NEW])))
-
- def test_upgrade_emits_sql_offline(self, capsys):
- """``db migrate --show-sql-only`` has no rows to read, so the
migration must not read any."""
- command.upgrade(_get_alembic_config(),
f"{_migration.down_revision}:{_migration.revision}", sql=True)
-
- emitted = capsys.readouterr().out
- assert "INSERT INTO team (name) SELECT lower(team.name)" in emitted
- assert "UPDATE team SET name=lower(team.name)" in emitted
-
- @pytest.mark.parametrize(
- ("dialect_name", "closing"),
- [
- pytest.param(
- "mysql",
- [
- "SET FOREIGN_KEY_CHECKS=0",
- "UPDATE team SET name=lower(team.name)",
- "SET FOREIGN_KEY_CHECKS=1",
- ],
- id="mysql-suspends-the-check-over-the-rename",
- ),
- pytest.param(
- "postgresql",
- ["UPDATE team SET name=lower(team.name)"],
- id="others-rename-under-the-check",
- ),
- ],
- )
- def test_only_mysql_suspends_foreign_key_checks(self, dialect_name,
closing):
- """InnoDB refuses the closing rename while the foreign keys still
guard the old name."""
- rendered = [
- str(statement).strip() for statement in
_migration.build_lower_casing_statements(dialect_name)
- ]
-
- assert rendered[-len(closing) :] == closing
- assert "SET FOREIGN_KEY_CHECKS=0" not in rendered[: -len(closing)]
diff --git
a/airflow-core/tests/unit/migrations/test_0136_fold_task_map_into_xcom.py
b/airflow-core/tests/unit/migrations/test_0136_fold_task_map_into_xcom.py
deleted file mode 100644
index 1b2e575a502..00000000000
--- a/airflow-core/tests/unit/migrations/test_0136_fold_task_map_into_xcom.py
+++ /dev/null
@@ -1,184 +0,0 @@
-#
-# 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 migration 0136 (3b7a91c5df20), which folds task_map into
xcom.mapped_length.
-
-An in-flight DagRun's expansion length either survives the backfill or quietly
disappears,
-so the real statements run against isolated tables on whichever backend the
suite is using.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-from pathlib import Path
-
-import pytest
-import sqlalchemy as sa
-
-from airflow import settings
-
-from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
-
-pytestmark = pytest.mark.db_test
-
-# Migration filenames start with a digit so they cannot be imported via the
normal import
-# system; load the module by file path instead.
-_MIGRATION_PATH = (
- Path(AIRFLOW_CORE_SOURCES_PATH)
- /
"airflow/migrations/versions/0136_3_4_0_fold_task_map_into_xcom_mapped_length.py"
-)
-_spec = importlib.util.spec_from_file_location("migration_0134",
_MIGRATION_PATH)
-_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
-_spec.loader.exec_module(_migration) # type: ignore[union-attr]
-
-# Isolated because the live xcom table has FK and NOT NULL columns, and
task_map is dropped.
-_XCOM = "_test_xcom_0134"
-_TASK_MAP = "_test_task_map_0134"
-
-_RETURN_VALUE = "return_value"
-
-_metadata = sa.MetaData()
-_xcom = sa.Table(
- _XCOM,
- _metadata,
- 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),
- # Declared rather than raw DDL so SQLAlchemy quotes it: reserved on MySQL.
- sa.Column("key", sa.String(512)),
- sa.Column("mapped_length", sa.Integer),
-)
-_task_map = sa.Table(
- _TASK_MAP,
- _metadata,
- sa.Column("dag_id", sa.String(250), nullable=False),
- sa.Column("task_id", sa.String(250), nullable=False),
- sa.Column("run_id", sa.String(250), nullable=False),
- sa.Column("map_index", sa.Integer, nullable=False),
- sa.Column("length", sa.Integer, nullable=False),
- sa.Column("keys", sa.String(512)),
- # The key-less PK the downgrade recreates, so an unscoped restore collides
here.
- sa.PrimaryKeyConstraint("dag_id", "task_id", "run_id", "map_index"),
-)
-
-_FEEDS_MAPPED = ("d", "feeds_mapped", "r", -1)
-_PLAIN = ("d", "plain", "r", -1)
-
-
-def _xcom_row(coords, key, mapped_length=None):
- dag_id, task_id, run_id, map_index = coords
- return {
- "dag_id": dag_id,
- "task_id": task_id,
- "run_id": run_id,
- "map_index": map_index,
- "key": key,
- "mapped_length": mapped_length,
- }
-
-
-def _task_map_row(coords, length):
- dag_id, task_id, run_id, map_index = coords
- return {
- "dag_id": dag_id,
- "task_id": task_id,
- "run_id": run_id,
- "map_index": map_index,
- "length": length,
- "keys": None,
- }
-
-
[email protected]
-def conn():
- _metadata.drop_all(settings.engine)
- _metadata.create_all(settings.engine)
- try:
- with settings.engine.begin() as connection:
- yield connection
- finally:
- _metadata.drop_all(settings.engine)
-
-
-def _lengths(conn) -> dict[tuple[str, str], int | None]:
- rows = conn.execute(sa.select(_xcom.c.task_id, _xcom.c.key,
_xcom.c.mapped_length)).all()
- return {(r.task_id, r.key): r.mapped_length for r in rows}
-
-
-def test_backfill_copies_the_length_onto_the_return_value_row(conn):
- conn.execute(
- _xcom.insert(),
- [
- _xcom_row(_FEEDS_MAPPED, _RETURN_VALUE),
- _xcom_row(_FEEDS_MAPPED, "side_output"),
- _xcom_row(_PLAIN, _RETURN_VALUE),
- ],
- )
- conn.execute(_task_map.insert(), [_task_map_row(_FEEDS_MAPPED, 3)])
-
- conn.execute(_migration.build_backfill_statement(_XCOM, _TASK_MAP))
-
- assert _lengths(conn) == {
- ("feeds_mapped", _RETURN_VALUE): 3,
- ("feeds_mapped", "side_output"): None,
- ("plain", _RETURN_VALUE): None,
- }
-
-
-def test_backfill_is_idempotent(conn):
- conn.execute(_xcom.insert(), [_xcom_row(_FEEDS_MAPPED, _RETURN_VALUE)])
- conn.execute(_task_map.insert(), [_task_map_row(_FEEDS_MAPPED, 3)])
-
- conn.execute(_migration.build_backfill_statement(_XCOM, _TASK_MAP))
- conn.execute(_migration.build_backfill_statement(_XCOM, _TASK_MAP))
-
- assert _lengths(conn) == {("feeds_mapped", _RETURN_VALUE): 3}
-
-
-def test_restore_rebuilds_task_map_from_the_return_value_length(conn):
- conn.execute(
- _xcom.insert(),
- [
- _xcom_row(_FEEDS_MAPPED, _RETURN_VALUE, mapped_length=3),
- _xcom_row(_PLAIN, _RETURN_VALUE),
- ],
- )
-
- conn.execute(_migration.build_restore_statement(_XCOM, _TASK_MAP))
-
- # Subscripted because ``.c.keys`` would resolve to ColumnCollection.keys,
the method.
- assert conn.execute(sa.select(_task_map.c.task_id, _task_map.c.length,
_task_map.c["keys"])).all() == [
- ("feeds_mapped", 3, None)
- ]
-
-
-def test_restore_ignores_a_length_recorded_under_another_key(conn):
- """Restoring both keys would collide on task_map's key-less primary key."""
- conn.execute(
- _xcom.insert(),
- [
- _xcom_row(_FEEDS_MAPPED, _RETURN_VALUE, mapped_length=3),
- _xcom_row(_FEEDS_MAPPED, "side_output", mapped_length=9),
- ],
- )
-
- conn.execute(_migration.build_restore_statement(_XCOM, _TASK_MAP))
-
- assert conn.execute(sa.select(_task_map.c.task_id,
_task_map.c.length)).all() == [("feeds_mapped", 3)]
diff --git
a/airflow-core/tests/unit/migrations/test_0138_allocate_pending_task_attempt_numbers.py
b/airflow-core/tests/unit/migrations/test_0138_allocate_pending_task_attempt_numbers.py
deleted file mode 100644
index 861892f537b..00000000000
---
a/airflow-core/tests/unit/migrations/test_0138_allocate_pending_task_attempt_numbers.py
+++ /dev/null
@@ -1,448 +0,0 @@
-#
-# 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 timedelta
-from importlib import import_module
-from uuid import uuid4
-
-import pytest
-import sqlalchemy as sa
-from alembic import command
-from alembic.migration import MigrationContext
-from alembic.operations import Operations
-
-from airflow import settings
-from airflow.models.hitl import HITLDetail
-from airflow.models.hitl_history import HITLDetailHistory
-from airflow.models.taskinstance import TaskInstance, TaskInstanceNote
-from airflow.models.taskinstancehistory import TaskInstanceHistory
-from airflow.models.taskreschedule import TaskReschedule
-from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.utils.db import _get_alembic_config
-from airflow.utils.state import TaskInstanceState
-
-_migration =
import_module("airflow.migrations.versions.0138_3_4_0_allocate_pending_task_attempt_numbers")
-
-
[email protected]
-def run_migration(session):
- def run(direction):
- with
Operations.context(MigrationContext.configure(session.connection())):
- direction()
-
- return run
-
-
[email protected]_test
[email protected](
- ("state", "try_number", "allocated_try_number", "with_history"),
- [
- (state, previous, allocated, with_history)
- for state, previous, allocated in [
- (None, 0, 0),
- (None, 1, 2),
- (None, 4, 5),
- (TaskInstanceState.UP_FOR_RETRY, 0, 1),
- (TaskInstanceState.UP_FOR_RETRY, 2, 3),
- *[(state, 2, 2) for state in TaskInstanceState if state !=
TaskInstanceState.UP_FOR_RETRY],
- ]
- for with_history in (False, True)
- if with_history or state != TaskInstanceState.UP_FOR_RETRY
- ],
-)
-def test_pending_attempt_migration_preserves_other_data(
- dag_maker, session, run_migration, state, try_number,
allocated_try_number, with_history
-):
- with dag_maker():
- EmptyOperator(task_id="task", retries=5)
- dr = dag_maker.create_dagrun()
- ti = dr.get_task_instance("task", session=session)
- ti.state = state
- ti.try_number = try_number
- ti.end_date = dr.start_date + timedelta(seconds=10)
- ti.retry_delay_override = 123
- ti.retry_reason = "Retry after the service recovers"
- if with_history:
- history = TaskInstanceHistory(ti, state=TaskInstanceState.FAILED)
- history.task_instance_id = uuid4()
- session.add(history)
- session.flush()
- live_query = sa.select(TaskInstance.__table__).where(TaskInstance.id ==
ti.id)
- history_query =
sa.select(TaskInstanceHistory.__table__).where(TaskInstanceHistory.dag_id ==
ti.dag_id)
- live_before = dict(session.execute(live_query).mappings().one())
- history_before = session.execute(history_query).mappings().all()
-
- run_migration(_migration.upgrade)
-
- assert dict(session.execute(live_query).mappings().one()) == {
- **live_before,
- "try_number": allocated_try_number,
- }
- assert session.execute(history_query).mappings().all() == history_before
-
- run_migration(_migration.downgrade)
-
- assert dict(session.execute(live_query).mappings().one()) == live_before
- assert session.execute(history_query).mappings().all() == history_before
-
-
[email protected]_test
[email protected]("try_number", [0, 2])
[email protected]("map_index", [-1, 0])
-def test_pending_retry_migration_archives_missing_attempt(
- dag_maker, session, run_migration, try_number, map_index
-):
- """Retire a legacy retry UUID without losing attempt data or breaking
UUID-based dependents.
-
- A queued/scheduled failure in older releases did not archive or rotate its
UUID. The migration
- must do both before advancing its try number. Downgrade keeps that archive
and fresh UUID,
- matching the old representation of a pending retry, so upgrading again
must not allocate twice.
- """
- with dag_maker():
- EmptyOperator(task_id="task", retries=5)
- dr = dag_maker.create_dagrun()
- ti = dr.get_task_instance("task", session=session)
- ti.state = TaskInstanceState.UP_FOR_RETRY
- ti.try_number = try_number
- ti.map_index = map_index
- ti.start_date = dr.start_date
- ti.end_date = dr.start_date + timedelta(seconds=10)
- ti.duration = 10
- ti.retry_delay_override = 123
- ti.retry_reason = "Retry after the service recovers"
- ti.executor_config = {"test": "keep executor configuration"}
- ti.next_kwargs = {"test": "keep continuation arguments"}
- session.flush()
- old_id = ti.id
- session.add(
- HITLDetail(
- ti_id=old_id,
- options=["approve", "reject"],
- subject="Review",
- body="Review the submitted request",
- defaults=["approve"],
- multiple=True,
- params={"reason": {"type": "string"}},
- assignees=[{"id": "reviewer", "name": "Reviewer"}],
- created_at=ti.start_date,
- responded_at=ti.end_date,
- responded_by={"id": "reviewer", "name": "Reviewer"},
- chosen_options=["approve"],
- params_input={"reason": "Approved after review"},
- )
- )
- note = TaskInstanceNote("Retain this note")
- note.ti_id = old_id
- session.add(note)
- session.add(TaskReschedule(old_id, ti.start_date, ti.end_date,
ti.end_date))
- session.flush()
- live_query = sa.select(TaskInstance.__table__).where(
- TaskInstance.dag_id == ti.dag_id,
- TaskInstance.task_id == ti.task_id,
- TaskInstance.run_id == ti.run_id,
- TaskInstance.map_index == map_index,
- )
- live_before = dict(session.execute(live_query).mappings().one())
- hitl_before =
dict(session.execute(sa.select(HITLDetail.__table__)).mappings().one())
-
- run_migration(_migration.upgrade)
-
- live_after = dict(session.execute(live_query).mappings().one())
- new_id = live_after["id"]
- assert new_id != old_id
- assert new_id.version == 7
- assert live_after == {**live_before, "id": new_id, "try_number":
try_number + 1}
- history_query = sa.select(TaskInstanceHistory.__table__).where(
- TaskInstanceHistory.task_instance_id == old_id
- )
- history = dict(session.execute(history_query).mappings().one())
- assert history == {
- **{column: live_before[column] for column in history if column !=
"task_instance_id"},
- "task_instance_id": old_id,
- "state": TaskInstanceState.FAILED,
- }
- assert (
-
session.scalar(sa.select(TaskInstanceNote.ti_id).where(TaskInstanceNote.content
== note.content))
- == new_id
- )
- assert
session.scalar(sa.select(sa.func.count()).select_from(TaskReschedule)) == 0
- assert
dict(session.execute(sa.select(HITLDetail.__table__)).mappings().one()) == {
- **hitl_before,
- "ti_id": new_id,
- }
- assert
dict(session.execute(sa.select(HITLDetailHistory.__table__)).mappings().one())
== {
- **{column: value for column, value in hitl_before.items() if column !=
"ti_id"},
- "ti_history_id": old_id,
- }
-
- run_migration(_migration.downgrade)
- assert dict(session.execute(live_query).mappings().one()) ==
{**live_before, "id": new_id}
- assert dict(session.execute(history_query).mappings().one()) == history
- run_migration(_migration.upgrade)
- assert dict(session.execute(live_query).mappings().one()) == live_after
- assert dict(session.execute(history_query).mappings().one()) == history
-
-
[email protected]_test
-def test_pending_retry_migration_only_retires_unarchived_tries_in_mixed_batch(
- dag_maker, session, run_migration
-):
- """Only missing retry histories retire UUIDs; their dependent rows must
follow the same selection.
-
- Older tries and other map indexes must not prevent archival. Already
archived retries and
- cleared tasks advance their try numbers without rotating UUIDs or changing
dependent rows.
- Unstarted, running, failed, and rescheduled tasks must retain all their
data.
- """
- with dag_maker() as dag:
- task = EmptyOperator(task_id="task", retries=5)
- dr = dag_maker.create_dagrun()
- session.execute(
- sa.delete(TaskInstance).where(TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id == dr.run_id)
- )
- states_and_tries = [
- (TaskInstanceState.UP_FOR_RETRY, 2),
- (TaskInstanceState.UP_FOR_RETRY, 2),
- (TaskInstanceState.UP_FOR_RETRY, 2),
- (None, 2),
- (None, 0),
- (TaskInstanceState.RUNNING, 2),
- (TaskInstanceState.FAILED, 2),
- (TaskInstanceState.UP_FOR_RESCHEDULE, 2),
- ]
- tasks = [
- TaskInstance(
- task=task,
- run_id=dr.run_id,
- map_index=index,
- dag_version_id=dr.created_dag_version_id,
- state=state,
- )
- for index, (state, _) in enumerate(states_and_tries)
- ]
- for ti, (_, try_number) in zip(tasks, states_and_tries):
- ti.try_number = try_number
- session.add_all(tasks)
- session.flush()
-
- dependents = {}
- for index in (0, 1, 2, 3, 5, 6, 7):
- ti = tasks[index]
- hitl = HITLDetail(
- ti_id=ti.id,
- options=["approve"],
- subject=f"Review index {index}",
- params={},
- params_input={},
- )
- dependents[index] = hitl
- note = TaskInstanceNote(f"Note for index {index}")
- note.ti_id = ti.id
- session.add_all(
- [
- hitl,
- note,
- TaskReschedule(ti.id, dr.start_date, dr.start_date,
dr.start_date),
- TaskReschedule(ti.id, dr.start_date, dr.start_date,
dr.start_date + timedelta(minutes=1)),
- ]
- )
- session.flush()
-
- for index, try_number in ((0, 2), (1, 1), (3, 2)):
- history = TaskInstanceHistory(tasks[index],
state=TaskInstanceState.FAILED)
- history.task_instance_id = uuid4()
- history.try_number = try_number
- session.add(history)
- session.flush()
- hitl_history = HITLDetailHistory(dependents[index])
- hitl_history.ti_history_id = history.task_instance_id
- hitl_history.subject = f"Previous review for index {index}"
- session.add(hitl_history)
- session.flush()
-
- def rows_by(model, key):
- return {row[key]: dict(row) for row in
session.execute(sa.select(model.__table__)).mappings()}
-
- live_before = rows_by(TaskInstance, "map_index")
- history_before = rows_by(TaskInstanceHistory, "task_instance_id")
- hitl_before = rows_by(HITLDetail, "ti_id")
- hitl_history_before = rows_by(HITLDetailHistory, "ti_history_id")
- notes_before = rows_by(TaskInstanceNote, "ti_id")
- reschedules_before = rows_by(TaskReschedule, "id")
- old_ids = [ti.id for ti in tasks]
-
- run_migration(_migration.upgrade)
-
- live_after = rows_by(TaskInstance, "map_index")
- new_ids = [live_after[index]["id"] for index in range(len(tasks))]
- assert len(set(new_ids)) == len(tasks)
- for index in (1, 2):
- assert new_ids[index] not in old_ids
- assert new_ids[index].version == 7
- for index in (0, 3, 4, 5, 6, 7):
- assert new_ids[index] == old_ids[index]
- assert live_after == {
- index: {**live_before[index], "id": new_ids[index], "try_number":
try_number}
- for index, try_number in enumerate([3, 3, 3, 3, 0, 2, 2, 2])
- }
- expected_history = {
- **history_before,
- **{
- old_ids[index]: {
- **{
- column.name: live_before[index][column.name]
- for column in TaskInstanceHistory.__table__.columns
- if column.name != "task_instance_id"
- },
- "task_instance_id": old_ids[index],
- "state": TaskInstanceState.FAILED,
- }
- for index in (1, 2)
- },
- }
- expected_hitl = {
- new_ids[index]: {**hitl_before[old_ids[index]], "ti_id":
new_ids[index]} for index in dependents
- }
- expected_hitl_history = {
- **hitl_history_before,
- **{
- old_ids[index]: {
- **{key: value for key, value in
hitl_before[old_ids[index]].items() if key != "ti_id"},
- "ti_history_id": old_ids[index],
- }
- for index in (1, 2)
- },
- }
- expected_notes = {
- new_ids[index]: {**notes_before[old_ids[index]], "ti_id":
new_ids[index]} for index in dependents
- }
- expected_reschedules = {
- key: row for key, row in reschedules_before.items() if row["ti_id"]
not in (old_ids[1], old_ids[2])
- }
-
- assert rows_by(TaskInstanceHistory, "task_instance_id") == expected_history
- assert rows_by(HITLDetail, "ti_id") == expected_hitl
- assert rows_by(HITLDetailHistory, "ti_history_id") == expected_hitl_history
- assert rows_by(TaskInstanceNote, "ti_id") == expected_notes
- assert rows_by(TaskReschedule, "id") == expected_reschedules
-
- run_migration(_migration.downgrade)
- run_migration(_migration.upgrade)
-
- assert rows_by(TaskInstance, "map_index") == live_after
- assert rows_by(TaskInstanceHistory, "task_instance_id") == expected_history
- assert rows_by(HITLDetail, "ti_id") == expected_hitl
- assert rows_by(HITLDetailHistory, "ti_history_id") == expected_hitl_history
- assert rows_by(TaskInstanceNote, "ti_id") == expected_notes
- assert rows_by(TaskReschedule, "id") == expected_reschedules
-
-
[email protected]_test
-def
test_pending_retry_migration_does_not_match_history_from_other_dags_tasks_or_runs(
- dag_maker, session, run_migration
-):
- with dag_maker(dag_id="first_dag"):
- EmptyOperator(task_id="task")
- EmptyOperator(task_id="other_task")
- first_run = dag_maker.create_dagrun(run_id="first_run")
- later_run = dag_maker.create_dagrun_after(first_run, run_id="later_run")
- with dag_maker(dag_id="other_dag"):
- EmptyOperator(task_id="task")
- other_dag_run = dag_maker.create_dagrun(run_id="first_run")
- target = first_run.get_task_instance("task", session=session)
- other_tasks = [
- first_run.get_task_instance("other_task", session=session),
- later_run.get_task_instance("task", session=session),
- other_dag_run.get_task_instance("task", session=session),
- ]
- for ti in [target, *other_tasks]:
- ti.state = TaskInstanceState.UP_FOR_RETRY
- ti.try_number = 2
- for ti in other_tasks:
- history = TaskInstanceHistory(ti, state=TaskInstanceState.FAILED)
- history.task_instance_id = uuid4()
- session.add(history)
- session.flush()
- old_id = target.id
- other_ids = [ti.id for ti in other_tasks]
- history_before = {
- row.task_instance_id: dict(row)
- for row in
session.execute(sa.select(TaskInstanceHistory.__table__)).mappings()
- }
-
- run_migration(_migration.upgrade)
- session.expire_all()
-
- migrated = first_run.get_task_instance("task", session=session)
- assert migrated.id != old_id
- assert migrated.try_number == 3
- assert session.get(TaskInstanceHistory, old_id).try_number == 2
- for ti_id in other_ids:
- assert session.get(TaskInstance, ti_id).try_number == 3
- assert {
- row.task_instance_id: dict(row)
- for row in session.execute(
-
sa.select(TaskInstanceHistory.__table__).where(TaskInstanceHistory.task_instance_id
!= old_id)
- ).mappings()
- } == history_before
-
-
[email protected]_test
-def test_pending_try_migration_handles_empty_tables(dag_maker, session,
run_migration):
- with dag_maker():
- EmptyOperator(task_id="task")
- dag_maker.create_dagrun()
- session.execute(sa.delete(TaskInstance))
-
- run_migration(_migration.upgrade)
- run_migration(_migration.downgrade)
- run_migration(_migration.upgrade)
-
- for model in (
- TaskInstance,
- TaskInstanceHistory,
- HITLDetail,
- HITLDetailHistory,
- TaskInstanceNote,
- TaskReschedule,
- ):
- assert session.scalar(sa.select(sa.func.count()).select_from(model))
== 0
-
-
[email protected]("dialect_name", ["postgresql", "mysql", "sqlite"])
[email protected]("direction", ["upgrade", "downgrade"])
-def test_pending_attempt_migration_emits_sql_offline(capsys, monkeypatch,
dialect_name, direction):
- config = _get_alembic_config()
- monkeypatch.setattr(settings, "SQL_ALCHEMY_CONN", f"{dialect_name}://")
- revisions = f"{_migration.down_revision}:{_migration.revision}"
- if direction == "downgrade":
- revisions = f"{_migration.revision}:{_migration.down_revision}"
- getattr(command, direction)(config, revisions, sql=True)
-
- emitted = capsys.readouterr().out
- assert "UPDATE task_instance SET" in emitted
- assert "try_number=(task_instance.try_number " in emitted
- assert "up_for_retry" in emitted
- assert "task_instance.state IS NULL" in emitted
- if direction == "upgrade":
- assert "INSERT INTO task_instance_history" in emitted
- assert "INSERT INTO hitl_detail_history" in emitted
- assert "DELETE FROM task_reschedule" in emitted
- else:
- assert "task_instance_history" not in emitted
