Lee-W commented on code in PR #64972: URL: https://github.com/apache/airflow/pull/64972#discussion_r3152661222
########## airflow-core/tests/unit/migrations/test_migration_patterns.py: ########## @@ -0,0 +1,71 @@ +# +# 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. +"""Pytest harness for migration anti-pattern checks (MIG001/MIG002/MIG003).""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Import check functions from the prek hook script (single source of truth). +# --------------------------------------------------------------------------- +_PREK_SCRIPT = Path(__file__).resolve().parents[4] / "scripts" / "ci" / "prek" / "check_migration_patterns.py" +_spec = importlib.util.spec_from_file_location("check_migration_patterns", _PREK_SCRIPT) +if _spec is None or _spec.loader is None: + raise ImportError(f"Could not load prek script: {_PREK_SCRIPT}") +_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_mod) # type: ignore[union-attr] # typeshed: Loader lacks exec_module +check_mig001 = _mod.check_mig001 +check_mig002 = _mod.check_mig002 +check_mig003 = _mod.check_mig003 + +MIGRATIONS_DIR = Path(__file__).resolve().parents[3] / "src" / "airflow" / "migrations" / "versions" + + +def _migration_files() -> list[Path]: Review Comment: ```suggestion @pytest.fixture(scope=module) def _migration_files() -> list[Path]: ``` not a good function name, probably make it a fixture? ########## airflow-core/tests/unit/migrations/test_migration_utils.py: ########## @@ -0,0 +1,176 @@ +# +# 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 ``airflow.migrations.utils``.""" + +from __future__ import annotations + +import pytest +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy import text + +from airflow import settings +from airflow.migrations.utils import ( + disable_sqlite_fkeys, + ignore_sqlite_value_error, + mysql_drop_foreignkey_if_exists, +) + +pytestmark = pytest.mark.db_test + + +class TestDisableSqliteFkeys: + """Tests for :func:`disable_sqlite_fkeys`.""" + + def test_yields_op(self): + """The context manager yields the *op* object on every backend.""" + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + with disable_sqlite_fkeys(op) as yielded: + assert yielded is op + + @pytest.mark.backend("sqlite") + def test_sqlite_turns_off_fkeys(self): + """On SQLite, ``PRAGMA foreign_keys`` is OFF inside the context.""" + + with settings.engine.connect() as conn: + # Ensure foreign_keys is ON before entering the context manager. + conn.execute(text("PRAGMA foreign_keys=ON")) + conn.commit() + + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + + with disable_sqlite_fkeys(op): + result = conn.execute(text("PRAGMA foreign_keys")).scalar() + assert result == 0, "foreign_keys should be OFF inside disable_sqlite_fkeys" + + result = conn.execute(text("PRAGMA foreign_keys")).scalar() + assert result == 1, "foreign_keys should be restored to ON after disable_sqlite_fkeys" + + @pytest.mark.backend("postgres", "mysql") + def test_non_sqlite_is_noop(self): + """On non-SQLite backends, the context manager simply yields.""" + + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + with disable_sqlite_fkeys(op) as yielded: + assert yielded is op + + +class TestMysqlDropForeignkeyIfExists: + """Tests for :func:`mysql_drop_foreignkey_if_exists`.""" + + @pytest.mark.backend("mysql") + def test_drops_existing_fk(self): + """On MySQL, an existing foreign key is dropped.""" + + with settings.engine.connect() as conn: + conn.execute(text("CREATE TABLE IF NOT EXISTS _test_mig_parent (id INT PRIMARY KEY)")) + conn.execute( + text( + """ + CREATE TABLE IF NOT EXISTS _test_mig_child ( + id INT PRIMARY KEY, + parent_id INT, + CONSTRAINT _test_mig_fk FOREIGN KEY (parent_id) + REFERENCES _test_mig_parent(id) + ) + """ + ) + ) + conn.commit() + + try: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + mysql_drop_foreignkey_if_exists("_test_mig_fk", "_test_mig_child", op) + conn.commit() + + count = conn.execute( + text( + """ + SELECT COUNT(*) FROM information_schema.TABLE_CONSTRAINTS + WHERE CONSTRAINT_SCHEMA = DATABASE() + AND TABLE_NAME = '_test_mig_child' + AND CONSTRAINT_NAME = '_test_mig_fk' + AND CONSTRAINT_TYPE = 'FOREIGN KEY' + """ + ) + ).scalar() + assert count == 0, "Foreign key should have been dropped" + finally: + conn.execute(text("DROP TABLE IF EXISTS _test_mig_child")) + conn.execute(text("DROP TABLE IF EXISTS _test_mig_parent")) + conn.commit() + + @pytest.mark.backend("mysql") + def test_noop_for_nonexistent_fk(self): + """On MySQL, dropping a non-existent FK does not raise.""" + + with settings.engine.connect() as conn: + conn.execute(text("CREATE TABLE IF NOT EXISTS _test_mig_nofk (id INT PRIMARY KEY)")) + conn.commit() + + try: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + mysql_drop_foreignkey_if_exists("nonexistent_fk", "_test_mig_nofk", op) + conn.commit() + finally: + conn.execute(text("DROP TABLE IF EXISTS _test_mig_nofk")) + conn.commit() + + +class TestIgnoreSqliteValueError: + """Tests for :func:`ignore_sqlite_value_error`.""" + + @pytest.mark.backend("sqlite") + def test_sqlite_suppresses_value_error(self): + """On SQLite, ``ValueError`` is suppressed.""" + + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + cm = ignore_sqlite_value_error() + with cm: + raise ValueError("should be suppressed on SQLite") + + @pytest.mark.backend("postgres", "mysql") + def test_non_sqlite_does_not_suppress(self): + """On non-SQLite backends, ``ValueError`` propagates.""" + + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + cm = ignore_sqlite_value_error() + with pytest.raises(ValueError, match="should propagate"): + with cm: + raise ValueError("should propagate") + + def test_does_not_suppress_other_exceptions(self): + """Even on SQLite, only ``ValueError`` is suppressed — other exceptions propagate.""" + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + cm = ignore_sqlite_value_error() + with pytest.raises(TypeError): + with cm: Review Comment: ```suggestion with pytest.raises(TypeError): with ignore_sqlite_value_error(): ``` ########## scripts/ci/prek/check_migration_patterns.py: ########## @@ -0,0 +1,383 @@ +#!/usr/bin/env python +# +# 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. +# /// script +# requires-python = ">=3.10,<3.11" Review Comment: 3.10 EOL happens this year. ```suggestion # requires-python = ">=3.10" ``` ########## airflow-core/tests/unit/migrations/test_migration_utils.py: ########## @@ -0,0 +1,176 @@ +# +# 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 ``airflow.migrations.utils``.""" + +from __future__ import annotations + +import pytest +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy import text + +from airflow import settings +from airflow.migrations.utils import ( + disable_sqlite_fkeys, + ignore_sqlite_value_error, + mysql_drop_foreignkey_if_exists, +) + +pytestmark = pytest.mark.db_test + + +class TestDisableSqliteFkeys: + """Tests for :func:`disable_sqlite_fkeys`.""" + + def test_yields_op(self): + """The context manager yields the *op* object on every backend.""" + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + with disable_sqlite_fkeys(op) as yielded: + assert yielded is op + + @pytest.mark.backend("sqlite") + def test_sqlite_turns_off_fkeys(self): + """On SQLite, ``PRAGMA foreign_keys`` is OFF inside the context.""" + + with settings.engine.connect() as conn: + # Ensure foreign_keys is ON before entering the context manager. + conn.execute(text("PRAGMA foreign_keys=ON")) + conn.commit() + + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + + with disable_sqlite_fkeys(op): + result = conn.execute(text("PRAGMA foreign_keys")).scalar() + assert result == 0, "foreign_keys should be OFF inside disable_sqlite_fkeys" + + result = conn.execute(text("PRAGMA foreign_keys")).scalar() + assert result == 1, "foreign_keys should be restored to ON after disable_sqlite_fkeys" + + @pytest.mark.backend("postgres", "mysql") + def test_non_sqlite_is_noop(self): + """On non-SQLite backends, the context manager simply yields.""" + + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + with disable_sqlite_fkeys(op) as yielded: + assert yielded is op + + +class TestMysqlDropForeignkeyIfExists: + """Tests for :func:`mysql_drop_foreignkey_if_exists`.""" + + @pytest.mark.backend("mysql") + def test_drops_existing_fk(self): + """On MySQL, an existing foreign key is dropped.""" + + with settings.engine.connect() as conn: + conn.execute(text("CREATE TABLE IF NOT EXISTS _test_mig_parent (id INT PRIMARY KEY)")) + conn.execute( + text( + """ + CREATE TABLE IF NOT EXISTS _test_mig_child ( + id INT PRIMARY KEY, + parent_id INT, + CONSTRAINT _test_mig_fk FOREIGN KEY (parent_id) + REFERENCES _test_mig_parent(id) + ) + """ + ) + ) + conn.commit() + + try: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + mysql_drop_foreignkey_if_exists("_test_mig_fk", "_test_mig_child", op) + conn.commit() + + count = conn.execute( + text( + """ + SELECT COUNT(*) FROM information_schema.TABLE_CONSTRAINTS + WHERE CONSTRAINT_SCHEMA = DATABASE() + AND TABLE_NAME = '_test_mig_child' + AND CONSTRAINT_NAME = '_test_mig_fk' + AND CONSTRAINT_TYPE = 'FOREIGN KEY' + """ + ) + ).scalar() + assert count == 0, "Foreign key should have been dropped" + finally: + conn.execute(text("DROP TABLE IF EXISTS _test_mig_child")) + conn.execute(text("DROP TABLE IF EXISTS _test_mig_parent")) + conn.commit() + + @pytest.mark.backend("mysql") + def test_noop_for_nonexistent_fk(self): + """On MySQL, dropping a non-existent FK does not raise.""" + + with settings.engine.connect() as conn: + conn.execute(text("CREATE TABLE IF NOT EXISTS _test_mig_nofk (id INT PRIMARY KEY)")) + conn.commit() + + try: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + mysql_drop_foreignkey_if_exists("nonexistent_fk", "_test_mig_nofk", op) + conn.commit() + finally: + conn.execute(text("DROP TABLE IF EXISTS _test_mig_nofk")) + conn.commit() + + +class TestIgnoreSqliteValueError: + """Tests for :func:`ignore_sqlite_value_error`.""" + + @pytest.mark.backend("sqlite") + def test_sqlite_suppresses_value_error(self): + """On SQLite, ``ValueError`` is suppressed.""" + + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + cm = ignore_sqlite_value_error() + with cm: + raise ValueError("should be suppressed on SQLite") + + @pytest.mark.backend("postgres", "mysql") + def test_non_sqlite_does_not_suppress(self): + """On non-SQLite backends, ``ValueError`` propagates.""" + + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + cm = ignore_sqlite_value_error() + with pytest.raises(ValueError, match="should propagate"): + with cm: Review Comment: ```suggestion with pytest.raises(ValueError, match="should propagate"): with ignore_sqlite_value_error(): ``` ########## scripts/ci/prek/check_migration_patterns.py: ########## @@ -0,0 +1,383 @@ +#!/usr/bin/env python +# +# 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. +# /// script +# requires-python = ">=3.10,<3.11" +# dependencies = [ +# "rich>=13.6.0", +# ] +# /// +""" +Static analysis checks for Alembic migration anti-patterns. + +These checks use AST analysis to detect common migration mistakes that can cause +silent failures, particularly on SQLite where ``PRAGMA foreign_keys`` has specific +requirements about transaction state. + +Rules +===== + +MIG001 -- DML before ``disable_sqlite_fkeys`` +---------------------------------------------- + +**What it does:** +Detects ``op.execute()`` calls containing DML keywords (UPDATE, INSERT, DELETE) that +appear before a ``with disable_sqlite_fkeys(op):`` block in the same function. + +**Why is this bad:** +On SQLite, any DML statement triggers *autobegin*, starting a transaction. Once a +transaction is active, ``PRAGMA foreign_keys=off`` (issued by ``disable_sqlite_fkeys``) +becomes a no-op. This means foreign key checks remain enabled, and +``batch_alter_table`` (which recreates tables) may fail with foreign key constraint +violations. + +**Example (bad):** + +.. code-block:: python + + def upgrade(): + op.execute("UPDATE dag SET col = '' WHERE col IS NULL") # triggers autobegin + with disable_sqlite_fkeys(op): # PRAGMA is now a no-op! + with op.batch_alter_table("dag") as batch_op: + batch_op.alter_column("col", nullable=False) + +**Use instead:** + +.. code-block:: python + + def upgrade(): + with disable_sqlite_fkeys(op): + op.execute("UPDATE dag SET col = '' WHERE col IS NULL") + with op.batch_alter_table("dag") as batch_op: + batch_op.alter_column("col", nullable=False) + + +MIG002 -- DDL before ``disable_sqlite_fkeys`` +---------------------------------------------- + +**What it does:** +Detects any Alembic ``op.*`` call (DDL operations like ``op.add_column``, +``op.drop_column``, ``op.create_table``, ``op.drop_table``, ``op.create_index``, +``op.batch_alter_table``, etc.) that appears before a +``with disable_sqlite_fkeys(op):`` block in the same function. Excludes DML calls +covered by MIG001. + +**Why is this bad:** +DDL operations also trigger *autobegin* on SQLite, making the subsequent +``PRAGMA foreign_keys=off`` a no-op, exactly like DML. + +**Example (bad):** + +.. code-block:: python + + def upgrade(): + op.add_column("dag_run", sa.Column("created_at", ...)) # triggers autobegin + with disable_sqlite_fkeys(op): # PRAGMA is now a no-op! + with op.batch_alter_table("backfill") as batch_op: + batch_op.alter_column("col", nullable=True) + +**Use instead:** + +.. code-block:: python + + def upgrade(): + with disable_sqlite_fkeys(op): + op.add_column("dag_run", sa.Column("created_at", ...)) + with op.batch_alter_table("backfill") as batch_op: + batch_op.alter_column("col", nullable=True) + + +MIG003 -- DML without offline-mode guard +------------------------------------------ + +**What it does:** +Detects ``op.execute()`` calls containing DML keywords (UPDATE, INSERT, DELETE) in +``upgrade()`` or ``downgrade()`` functions that do not contain a +``context.is_offline_mode()`` check anywhere in the same function. + +**Why is this bad:** +Alembic's offline mode generates SQL scripts instead of executing against a live +database. DML statements like ``UPDATE ... WHERE col IS NULL`` depend on actual data +and produce no useful output in offline mode. Without an ``is_offline_mode()`` guard, +the generated script includes DML that may fail or silently do nothing when applied +later. + +**Example (bad):** + +.. code-block:: python + + def upgrade(): + op.execute("UPDATE dag SET col = '' WHERE col IS NULL") + with op.batch_alter_table("dag") as batch_op: + batch_op.alter_column("col", nullable=False) + +**Use instead:** + +.. code-block:: python + + def upgrade(): + if not context.is_offline_mode(): + op.execute("UPDATE dag SET col = '' WHERE col IS NULL") + with op.batch_alter_table("dag") as batch_op: + batch_op.alter_column("col", nullable=False) + + +Known Limitation -- Future MIG004 Candidate +--------------------------------------------- + +Migrations that perform DML and use ``batch_alter_table`` but never call +``disable_sqlite_fkeys`` at all are NOT detected by MIG001/MIG002. These migrations +may silently break SQLite foreign key handling. Detecting this pattern requires +understanding whether foreign key relationships exist on the affected tables, which +is beyond AST-level analysis. This is tracked as a future MIG004 candidate. + + +Suppression +=========== + +Add ``# noqa: MIG0XX`` to a line to suppress the corresponding check for that line. +Multiple codes can be comma-separated: ``# noqa: MIG001, MIG003``. + +Include a brief reason after ``--``:: + + # noqa: MIG003 -- simple UPDATE safe in offline mode +""" + +from __future__ import annotations + +import ast +import re +import sys +from pathlib import Path + +from rich.console import Console + +console = Console(color_system="standard", width=200) + +DML_KEYWORDS = ("UPDATE", "INSERT", "DELETE") + + +def _get_noqa_codes(line: str) -> set[str]: + """Extract noqa codes from a source line.""" + match = re.search(r"#\s*noqa:\s*([\w,\s]+)", line) + if match: + return {code.strip() for code in match.group(1).split(",")} + return set() + + +def _is_dml_string(node: ast.expr) -> bool: + """Check if an AST expression is a string literal starting with a DML keyword.""" + if isinstance(node, ast.Constant) and isinstance(node.value, str): + parts = node.value.strip().upper().split(None, 1) + return len(parts) > 1 and parts[0] in DML_KEYWORDS + # f-strings: check the first string fragment + if isinstance(node, ast.JoinedStr): + for value in node.values: + if isinstance(value, ast.Constant) and isinstance(value.value, str): + parts = value.value.strip().upper().split(None, 1) + if len(parts) > 1 and parts[0] in DML_KEYWORDS: + return True + break # only the leading fragment matters + return False + + +def _is_op_execute_with_dml(node: ast.Call) -> bool: + """Check if a Call node is ``op.execute(<DML string>)``.""" + if not ( + isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "op" + and node.func.attr == "execute" + ): + return False + if not node.args: + return False + arg = node.args[0] + # Direct string: op.execute("UPDATE ...") + if _is_dml_string(arg): + return True + # Wrapped: op.execute(text("UPDATE ...")) or op.execute(sa.text("UPDATE ...")) + if isinstance(arg, ast.Call) and arg.args: + return _is_dml_string(arg.args[0]) + return False + + +def _is_op_call(node: ast.Call) -> bool: + """Check if a Call node is ``op.<something>(...)``.""" + return ( + isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "op" + ) + + +def _find_disable_sqlite_fkeys_line(func_node: ast.FunctionDef) -> int | None: + """Find the line number of the first ``with disable_sqlite_fkeys(op):`` in a function.""" + # Collect all matches then take min() — ast.walk BFS order does not guarantee + # line-number order when guards appear at different nesting depths. + found: list[int] = [] + for node in ast.walk(func_node): + if not isinstance(node, ast.With): + continue + for item in node.items: + ctx = item.context_expr + if isinstance(ctx, ast.Call): + func = ctx.func + if (isinstance(func, ast.Name) and func.id == "disable_sqlite_fkeys") or ( + isinstance(func, ast.Attribute) and func.attr == "disable_sqlite_fkeys" + ): + found.append(node.lineno) + return min(found) if found else None + + +def _has_offline_mode_check(func_node: ast.FunctionDef) -> bool: + """Check if function contains ``context.is_offline_mode()`` call.""" + for node in ast.walk(func_node): + if isinstance(node, ast.Call): + func = node.func + if isinstance(func, ast.Attribute) and func.attr == "is_offline_mode": + return True + return False + + +def _get_upgrade_downgrade_functions(tree: ast.Module) -> list[ast.FunctionDef]: + """Get ``upgrade()`` and ``downgrade()`` function definitions from module.""" + return [ + node + for node in ast.iter_child_nodes(tree) + if isinstance(node, ast.FunctionDef) and node.name in ("upgrade", "downgrade") + ] + + +def _line_has_noqa(lines: list[str], lineno: int, code: str) -> bool: + """Check if a source line has a noqa suppression for the given code.""" + if lineno < 1 or lineno > len(lines): + return False + return code in _get_noqa_codes(lines[lineno - 1]) + + +# --------------------------------------------------------------------------- +# Check functions +# --------------------------------------------------------------------------- + + +def check_mig001(filepath: Path) -> list[str]: + """MIG001: DML before ``disable_sqlite_fkeys``.""" + source = filepath.read_text() + lines = source.splitlines() + tree = ast.parse(source, filename=str(filepath)) + errors: list[str] = [] + + for func in _get_upgrade_downgrade_functions(tree): + guard_line = _find_disable_sqlite_fkeys_line(func) + if guard_line is None: + continue # MIG001 only fires when disable_sqlite_fkeys IS present + + for node in ast.walk(func): + if not isinstance(node, ast.Call): + continue + if node.lineno >= guard_line: + continue + if _is_op_execute_with_dml(node) and not _line_has_noqa(lines, node.lineno, "MIG001"): + errors.append( + f"MIG001 {func.name}(): op.execute() with DML at line {node.lineno} " + f"before disable_sqlite_fkeys at line {guard_line}" + ) + + return errors + + +def check_mig002(filepath: Path) -> list[str]: + """MIG002: DDL before ``disable_sqlite_fkeys``.""" + source = filepath.read_text() + lines = source.splitlines() + tree = ast.parse(source, filename=str(filepath)) + errors: list[str] = [] + + for func in _get_upgrade_downgrade_functions(tree): + guard_line = _find_disable_sqlite_fkeys_line(func) + if guard_line is None: + continue # MIG002 only fires when disable_sqlite_fkeys IS present + + for node in ast.walk(func): + if not isinstance(node, ast.Call): + continue + if node.lineno >= guard_line: + continue + # Skip DML calls — covered by MIG001 + if _is_op_execute_with_dml(node): + continue + if _is_op_call(node) and not _line_has_noqa(lines, node.lineno, "MIG002"): + attr = node.func.attr if isinstance(node.func, ast.Attribute) else "?" + errors.append( + f"MIG002 {func.name}(): op.{attr}() at line {node.lineno} " + f"before disable_sqlite_fkeys at line {guard_line}" + ) + + return errors + + +def check_mig003(filepath: Path) -> list[str]: + """MIG003: DML without offline-mode guard.""" + source = filepath.read_text() + lines = source.splitlines() + tree = ast.parse(source, filename=str(filepath)) + errors: list[str] = [] + + for func in _get_upgrade_downgrade_functions(tree): + if _has_offline_mode_check(func): + continue # function has the guard + + for node in ast.walk(func): + if not isinstance(node, ast.Call): + continue + if _is_op_execute_with_dml(node) and not _line_has_noqa(lines, node.lineno, "MIG003"): + errors.append( + f"MIG003 {func.name}(): op.execute() with DML at line {node.lineno} " + f"without context.is_offline_mode() guard" + ) + + return errors + + +# --------------------------------------------------------------------------- +# CLI entry point +# --------------------------------------------------------------------------- + + +def main() -> int: + """Check migration files for anti-patterns. Exit 0 if clean, 1 if violations found.""" + checks = [check_mig001, check_mig002, check_mig003] Review Comment: ```suggestion check_fns = [check_mig001, check_mig002, check_mig003] ``` nit: easier to read IMO. i was not sure what this `checks` are when readding line 371. but non-block ########## scripts/tests/ci/prek/test_check_migration_patterns.py: ########## @@ -0,0 +1,333 @@ +# +# 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 ``scripts/ci/prek/check_migration_patterns.py``.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +from ci.prek.check_migration_patterns import ( + _get_noqa_codes, + _line_has_noqa, + check_mig001, + check_mig002, + check_mig003, + main, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def write_migration(tmp_path: Path, source: str) -> Path: + """Write *source* to a temp file that looks like a migration file.""" + p = tmp_path / "0099_test_migration.py" + p.write_text(source) + return p + + +# --------------------------------------------------------------------------- +# _get_noqa_codes +# --------------------------------------------------------------------------- + + +class TestGetNoqaCodes: + def test_single_code(self): + assert _get_noqa_codes("op.execute(...) # noqa: MIG001") == {"MIG001"} + + def test_multiple_codes(self): + assert _get_noqa_codes("op.execute(...) # noqa: MIG001, MIG003") == {"MIG001", "MIG003"} + + def test_code_with_reason_suffix(self): + # Reason after '--' should NOT be included in the code set. + codes = _get_noqa_codes("op.execute(...) # noqa: MIG003 -- safe in offline mode") + assert "MIG003" in codes + assert "--" not in "".join(codes) + assert "safe" not in "".join(codes) + + def test_no_noqa_comment(self): + assert _get_noqa_codes("op.execute('UPDATE foo SET x=1')") == set() + + def test_empty_string(self): + assert _get_noqa_codes("") == set() + + +# --------------------------------------------------------------------------- +# _line_has_noqa +# --------------------------------------------------------------------------- + + +class TestLineHasNoqa: + def test_line_has_matching_code(self): + lines = ["op.execute('UPDATE ...') # noqa: MIG001"] + assert _line_has_noqa(lines, 1, "MIG001") is True + + def test_line_has_different_code(self): + lines = ["op.execute('UPDATE ...') # noqa: MIG002"] + assert _line_has_noqa(lines, 1, "MIG001") is False + + def test_out_of_bounds_zero(self): + assert _line_has_noqa(["some line"], 0, "MIG001") is False + + def test_out_of_bounds_beyond_end(self): + assert _line_has_noqa(["some line"], 2, "MIG001") is False Review Comment: looks like something we can parameterize ########## airflow-core/tests/unit/migrations/test_migration_utils.py: ########## @@ -0,0 +1,176 @@ +# +# 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 ``airflow.migrations.utils``.""" + +from __future__ import annotations + +import pytest +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy import text + +from airflow import settings +from airflow.migrations.utils import ( + disable_sqlite_fkeys, + ignore_sqlite_value_error, + mysql_drop_foreignkey_if_exists, +) + +pytestmark = pytest.mark.db_test + + +class TestDisableSqliteFkeys: + """Tests for :func:`disable_sqlite_fkeys`.""" + + def test_yields_op(self): + """The context manager yields the *op* object on every backend.""" + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + with disable_sqlite_fkeys(op) as yielded: + assert yielded is op + + @pytest.mark.backend("sqlite") + def test_sqlite_turns_off_fkeys(self): + """On SQLite, ``PRAGMA foreign_keys`` is OFF inside the context.""" + + with settings.engine.connect() as conn: + # Ensure foreign_keys is ON before entering the context manager. + conn.execute(text("PRAGMA foreign_keys=ON")) + conn.commit() + + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + + with disable_sqlite_fkeys(op): + result = conn.execute(text("PRAGMA foreign_keys")).scalar() + assert result == 0, "foreign_keys should be OFF inside disable_sqlite_fkeys" + + result = conn.execute(text("PRAGMA foreign_keys")).scalar() + assert result == 1, "foreign_keys should be restored to ON after disable_sqlite_fkeys" + + @pytest.mark.backend("postgres", "mysql") + def test_non_sqlite_is_noop(self): + """On non-SQLite backends, the context manager simply yields.""" + + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + with disable_sqlite_fkeys(op) as yielded: + assert yielded is op + + +class TestMysqlDropForeignkeyIfExists: + """Tests for :func:`mysql_drop_foreignkey_if_exists`.""" + + @pytest.mark.backend("mysql") + def test_drops_existing_fk(self): + """On MySQL, an existing foreign key is dropped.""" + + with settings.engine.connect() as conn: + conn.execute(text("CREATE TABLE IF NOT EXISTS _test_mig_parent (id INT PRIMARY KEY)")) + conn.execute( + text( + """ + CREATE TABLE IF NOT EXISTS _test_mig_child ( + id INT PRIMARY KEY, + parent_id INT, + CONSTRAINT _test_mig_fk FOREIGN KEY (parent_id) + REFERENCES _test_mig_parent(id) + ) + """ + ) + ) + conn.commit() + + try: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + mysql_drop_foreignkey_if_exists("_test_mig_fk", "_test_mig_child", op) + conn.commit() + + count = conn.execute( + text( + """ + SELECT COUNT(*) FROM information_schema.TABLE_CONSTRAINTS + WHERE CONSTRAINT_SCHEMA = DATABASE() + AND TABLE_NAME = '_test_mig_child' + AND CONSTRAINT_NAME = '_test_mig_fk' + AND CONSTRAINT_TYPE = 'FOREIGN KEY' + """ + ) + ).scalar() + assert count == 0, "Foreign key should have been dropped" + finally: + conn.execute(text("DROP TABLE IF EXISTS _test_mig_child")) + conn.execute(text("DROP TABLE IF EXISTS _test_mig_parent")) + conn.commit() + + @pytest.mark.backend("mysql") + def test_noop_for_nonexistent_fk(self): + """On MySQL, dropping a non-existent FK does not raise.""" + + with settings.engine.connect() as conn: + conn.execute(text("CREATE TABLE IF NOT EXISTS _test_mig_nofk (id INT PRIMARY KEY)")) + conn.commit() + + try: + ctx = MigrationContext.configure(conn) + op = Operations(ctx) + mysql_drop_foreignkey_if_exists("nonexistent_fk", "_test_mig_nofk", op) + conn.commit() + finally: + conn.execute(text("DROP TABLE IF EXISTS _test_mig_nofk")) + conn.commit() + + +class TestIgnoreSqliteValueError: + """Tests for :func:`ignore_sqlite_value_error`.""" + + @pytest.mark.backend("sqlite") + def test_sqlite_suppresses_value_error(self): + """On SQLite, ``ValueError`` is suppressed.""" + + with settings.engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + cm = ignore_sqlite_value_error() + with cm: Review Comment: ```suggestion with ignore_sqlite_value_error(): ``` ########## scripts/tests/ci/prek/test_check_migration_patterns.py: ########## @@ -0,0 +1,333 @@ +# +# 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 ``scripts/ci/prek/check_migration_patterns.py``.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +from ci.prek.check_migration_patterns import ( + _get_noqa_codes, + _line_has_noqa, + check_mig001, + check_mig002, + check_mig003, + main, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def write_migration(tmp_path: Path, source: str) -> Path: + """Write *source* to a temp file that looks like a migration file.""" + p = tmp_path / "0099_test_migration.py" + p.write_text(source) + return p + + +# --------------------------------------------------------------------------- +# _get_noqa_codes +# --------------------------------------------------------------------------- + + +class TestGetNoqaCodes: + def test_single_code(self): + assert _get_noqa_codes("op.execute(...) # noqa: MIG001") == {"MIG001"} + + def test_multiple_codes(self): + assert _get_noqa_codes("op.execute(...) # noqa: MIG001, MIG003") == {"MIG001", "MIG003"} + + def test_code_with_reason_suffix(self): + # Reason after '--' should NOT be included in the code set. + codes = _get_noqa_codes("op.execute(...) # noqa: MIG003 -- safe in offline mode") + assert "MIG003" in codes + assert "--" not in "".join(codes) + assert "safe" not in "".join(codes) + + def test_no_noqa_comment(self): + assert _get_noqa_codes("op.execute('UPDATE foo SET x=1')") == set() + + def test_empty_string(self): + assert _get_noqa_codes("") == set() + + +# --------------------------------------------------------------------------- +# _line_has_noqa +# --------------------------------------------------------------------------- + + +class TestLineHasNoqa: + def test_line_has_matching_code(self): + lines = ["op.execute('UPDATE ...') # noqa: MIG001"] + assert _line_has_noqa(lines, 1, "MIG001") is True + + def test_line_has_different_code(self): + lines = ["op.execute('UPDATE ...') # noqa: MIG002"] + assert _line_has_noqa(lines, 1, "MIG001") is False + + def test_out_of_bounds_zero(self): + assert _line_has_noqa(["some line"], 0, "MIG001") is False + + def test_out_of_bounds_beyond_end(self): + assert _line_has_noqa(["some line"], 2, "MIG001") is False + + +# --------------------------------------------------------------------------- +# check_mig001 +# --------------------------------------------------------------------------- + + +class TestCheckMig001: Review Comment: we probably can parametrize this as well. not all, but test cases can be grouped ########## scripts/tests/ci/prek/test_check_migration_patterns.py: ########## @@ -0,0 +1,333 @@ +# +# 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 ``scripts/ci/prek/check_migration_patterns.py``.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +from ci.prek.check_migration_patterns import ( + _get_noqa_codes, + _line_has_noqa, + check_mig001, + check_mig002, + check_mig003, + main, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def write_migration(tmp_path: Path, source: str) -> Path: + """Write *source* to a temp file that looks like a migration file.""" + p = tmp_path / "0099_test_migration.py" + p.write_text(source) + return p + + +# --------------------------------------------------------------------------- +# _get_noqa_codes +# --------------------------------------------------------------------------- Review Comment: I don't think we need these ```suggestion ``` we can move them test docstring if we really want it -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
