scripts/ci/prek/check_migration_patterns.py

#!/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.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)


**Exception -- DML tied to DDL on the same table:**

When an ``op.execute()`` DML call targets a table that also has DDL operations
(``batch_alter_table``, ``create_table``, ``drop_table``, ``rename_table``,
``add_column``, ``drop_column``, ``alter_column``, ``create_index``, ``drop_index``,
``drop_constraint``) in the same ``upgrade()`` or ``downgrade()`` function, MIG003
is not raised.  This covers the common pattern of filling NULL values before adding
a NOT NULL constraint, deleting rows before dropping a table, etc.  The DML is
considered a safe prerequisite for the DDL.


MIG004 -- Migration description style
--------------------------------------

**What it does:**
Checks the first paragraph of the migration module docstring -- the description
Alembic exposes as ``Script.doc``. It must start with a capital letter, use the
imperative mood ("Add", not "Added"/"Adding"), and end with a period.

**Why is this bad:**
The description is user-facing: it is rendered into the ``migrations-ref.rst``
documentation table (by the ``update-migration-references`` hook) and printed by
``alembic history``.  Inconsistent capitalization and verb tense make those
listings needlessly messy.

**Example (bad):** a migration docstring starting with
``Added DagPriorityParsingRequest table.`` or
``add timetable_type to dag table for filtering.``

**Use instead:** ``Add DagPriorityParsingRequest table.`` and
``Add timetable_type to dag table for filtering.``


Known Limitation -- Future MIG005 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 MIG005 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 dataclasses
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")

_DML_TABLE_RE = re.compile(
    r"^\s*(?:"
    r"(?:UPDATE)\s+[`\"]?(\w+)[`\"]?"
    r"|(?:DELETE\s+FROM)\s+[`\"]?(\w+)[`\"]?"
    r"|(?:INSERT\s+(?:(?:OR\s+)?IGNORE\s+)?INTO)\s+[`\"]?(\w+)[`\"]?"
    r")",
    re.IGNORECASE,
)


@dataclasses.dataclass(frozen=True)
class MigrationFile:
    path: Path
    source: str
    lines: list[str]
    tree: ast.Module

    @classmethod
    def from_path(cls, filepath: Path) -> MigrationFile:
        source = filepath.read_text()
        return cls(
            path=filepath,
            source=source,
            lines=source.splitlines(),
            tree=ast.parse(source, filename=str(filepath)),
        )


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 _get_sql_string_value(node: ast.expr) -> str | None:
    """Extract the leading SQL string from an AST expression."""
    if isinstance(node, ast.Constant) and isinstance(node.value, str):
        return node.value
    if isinstance(node, ast.JoinedStr):
        for value in node.values:
            if isinstance(value, ast.Constant) and isinstance(value.value, str):
                return value.value
            break
    return None


def _extract_dml_table_name(node: ast.Call) -> str | None:
    """Extract the target table name from a DML ``op.execute()`` call.

    Returns the lowercase table name, or None if extraction fails.
    """
    if not node.args:
        return None
    arg = node.args[0]
    if isinstance(arg, ast.Call) and arg.args:
        arg = arg.args[0]
    sql = _get_sql_string_value(arg)
    if sql is None:
        return None
    m = _DML_TABLE_RE.match(sql)
    if m is None:
        return None
    table = m.group(1) or m.group(2) or m.group(3)
    return table.lower() if table else None


_DDL_TABLE_FIRST_ARG = frozenset(
    {
        "batch_alter_table",
        "drop_table",
        "create_table",
        "add_column",
        "drop_column",
        "alter_column",
    }
)


def _collect_ddl_table_names(func_node: ast.FunctionDef) -> set[str]:
    """Collect all table names referenced by DDL operations in a function."""
    tables: set[str] = set()

    def _str_arg(call: ast.Call, index: int) -> str | None:
        if index < len(call.args):
            a = call.args[index]
            if isinstance(a, ast.Constant) and isinstance(a.value, str):
                return a.value.lower()
        return None

    for node in ast.walk(func_node):
        if not isinstance(node, ast.Call):
            continue
        if not (
            isinstance(node.func, ast.Attribute)
            and isinstance(node.func.value, ast.Name)
            and node.func.value.id == "op"
        ):
            continue
        attr = node.func.attr
        if attr in _DDL_TABLE_FIRST_ARG:
            name = _str_arg(node, 0)
            if name:
                tables.add(name)
        elif attr == "rename_table":
            for i in range(2):
                name = _str_arg(node, i)
                if name:
                    tables.add(name)
        elif attr in ("create_index", "drop_constraint"):
            name = _str_arg(node, 1)
            if name:
                tables.add(name)
        elif attr == "drop_index":
            name = _str_arg(node, 1)
            if name is None:
                for kw in node.keywords:
                    if kw.arg == "table_name":
                        if isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
                            name = kw.value.value.lower()
                        break
            if name:
                tables.add(name)

    return tables


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(mf: MigrationFile) -> list[str]:
    """MIG001: DML before ``disable_sqlite_fkeys``."""
    errors: list[str] = []

    for func in _get_upgrade_downgrade_functions(mf.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(mf.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(mf: MigrationFile) -> list[str]:
    """MIG002: DDL before ``disable_sqlite_fkeys``."""
    errors: list[str] = []

    for func in _get_upgrade_downgrade_functions(mf.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(mf.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(mf: MigrationFile) -> list[str]:
    """MIG003: DML without offline-mode guard."""
    errors: list[str] = []

    for func in _get_upgrade_downgrade_functions(mf.tree):
        if _has_offline_mode_check(func):
            continue  # function has the guard

        ddl_tables = _collect_ddl_table_names(func)

        for node in ast.walk(func):
            if not isinstance(node, ast.Call):
                continue
            if not _is_op_execute_with_dml(node):
                continue
            dml_table = _extract_dml_table_name(node)
            if dml_table is not None and dml_table in ddl_tables:
                continue
            if not _line_has_noqa(mf.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


def check_mig004(mf: MigrationFile) -> list[str]:
    """MIG004: description must be capitalized, imperative, and end with a period."""
    if mf.path.name == "__init__.py":
        return []
    doc = ast.get_docstring(mf.tree)
    # Same extraction as alembic's Script.doc.
    description = re.split(r"\n\n", doc.strip())[0].strip() if doc else ""
    if not description:
        return ["MIG004 module docstring must start with a migration description"]

    errors: list[str] = []
    if not description[0].isupper():
        errors.append(f"MIG004 description {description!r} must start with a capital letter")
    first_word = re.match(r"[A-Za-z]+", description)
    if first_word and first_word.group().lower().endswith(("ed", "ing")):
        errors.append(
            f"MIG004 description {description!r} must use the imperative mood "
            f"('Add ...', not 'Added ...' or 'Adding ...')"
        )
    if not description.endswith("."):
        errors.append(f"MIG004 description {description!r} must end with a period")
    return errors


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------


def main() -> int:
    """Check migration files for anti-patterns. Exit 0 if clean, 1 if violations found."""
    check_fns = [check_mig001, check_mig002, check_mig003, check_mig004]
    has_errors = False

    for filepath_str in sys.argv[1:]:
        filepath = Path(filepath_str)
        try:
            mf = MigrationFile.from_path(filepath)
            for check_fn in check_fns:
                for error in check_fn(mf):
                    console.print(f"{filepath}: {error}")
                    has_errors = True
        except (OSError, SyntaxError, UnicodeDecodeError) as exc:
            console.print(f"{filepath}: {exc}")
            has_errors = True

    return 1 if has_errors else 0


if __name__ == "__main__":
    sys.exit(main())