diff --git a/changelog.d/tsk-kwtvfq-schema-column-follow-helpers.md b/changelog.d/tsk-kwtvfq-schema-column-follow-helpers.md new file mode 100644 index 000000000..d850eff8b --- /dev/null +++ b/changelog.d/tsk-kwtvfq-schema-column-follow-helpers.md @@ -0,0 +1,7 @@ +### Fixed + +- `scripts/check_schema_column_migrations.py` now follows one level of same-file + module-level helper calls from `_post_init` when checking for ALTER TABLE + migrations. A guarded migration that lives in a module-level coroutine called + by `_post_init` (the `agent_registry_store.py` pattern) no longer produces a + false violation. diff --git a/scripts/check_schema_column_migrations.py b/scripts/check_schema_column_migrations.py index e6b6c8181..f51485210 100644 --- a/scripts/check_schema_column_migrations.py +++ b/scripts/check_schema_column_migrations.py @@ -129,7 +129,8 @@ class Violation: def __str__(self) -> str: fix = ( "add a guarded _post_init coroutine that ALTERs this column " - "into place after a PRAGMA table_info check" + "into place after a PRAGMA table_info check, either inline " + "or via a module-level helper that _post_init calls" ) return ( f"{self.path}: table '{self.table}', column '{self.column}' " @@ -449,8 +450,30 @@ def _post_init_added_columns(tree: ast.AST) -> set[tuple[str, str]]: SQL is read out of the AST's string constants rather than out of stripped source text, so ``#`` inside a SQL literal cannot chop the statement and a triple-quoted SQL literal is not mistaken for a docstring. + + One level of same-file call indirection is also followed: if ``_post_init`` + calls a module-level ``FunctionDef``/``AsyncFunctionDef`` by plain name + (e.g. ``await _migration_v1_add_status(self._db)``), that helper's SQL + literals are collected too. A visited set prevents cycles. """ added: set[tuple[str, str]] = set() + module_functions: dict[str, ast.FunctionDef | ast.AsyncFunctionDef] = {} + for node in ast.walk(tree): + if isinstance(node, ast.Module): + for item in node.body: + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)): + module_functions[item.name] = item + + def _called_names(fn: ast.AST) -> set[str]: + names: set[str] = set() + for child in ast.walk(fn): + if ( + isinstance(child, ast.Call) + and isinstance(child.func, ast.Name) + ): + names.add(child.func.id) + return names + for node in ast.walk(tree): if not isinstance(node, ast.ClassDef): continue @@ -459,9 +482,23 @@ def _post_init_added_columns(tree: ast.AST) -> set[tuple[str, str]]: continue if item.name != "_post_init": continue - for literal in _method_sql_literals(item): - for m in _ADD_COLUMN_RE.finditer(literal): - added.add((m.group(1), m.group(2))) + visited: set[str] = set() + queue = [item] + while queue: + fn = queue.pop(0) + if fn.name in visited: + continue + visited.add(fn.name) + for literal in _method_sql_literals(fn): + for m in _ADD_COLUMN_RE.finditer(literal): + added.add((m.group(1), m.group(2))) + for name in _called_names(fn): + helper = module_functions.get(name) + if ( + helper is not None + and helper.name not in visited + ): + queue.append(helper) return added diff --git a/tests/scripts/test_check_schema_column_migrations.py b/tests/scripts/test_check_schema_column_migrations.py index 72c8a628d..d81b6f380 100644 --- a/tests/scripts/test_check_schema_column_migrations.py +++ b/tests/scripts/test_check_schema_column_migrations.py @@ -1036,3 +1036,130 @@ class Store: schemas = guard_mod._extract_schemas(ast.parse(src)) assert len(schemas) == 1 assert "CREATE TABLE shared" in schemas[0] + + +class TestPostInitFollowsModuleHelpers: + """RED-FIRST: _post_init_added_columns must follow one level of same-file + module-level function calls. + + Case (a) RED: SCHEMA gains a column, no ALTER anywhere -> violation. + Case (b) RED: ALTER in a module-level helper that _post_init does NOT call + -> violation (a helper merely existing in the file must not silence). + Case (c) GREEN: ALTER in a module-level helper that _post_init DOES call + -> clean (this is the agent_registry_store.py shape that was wrongly + red before the fix). + """ + + def _baseline(self, guard_mod, monkeypatch, baselines: dict) -> None: + monkeypatch.setattr( + guard_mod, "_baseline_columns", lambda p, ref: baselines.get(p.name, {}) + ) + + def test_case_a_no_alter_anywhere_remains_red( + self, guard_mod, tmp_path: Path, monkeypatch + ) -> None: + body = ''' +SCHEMA = """ +CREATE TABLE IF NOT EXISTS gadgets ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL DEFAULT '' +); +""" +''' + path = _write_store(tmp_path, "no_alter.py", body) + self._baseline(guard_mod, monkeypatch, {"no_alter.py": {"gadgets": {"id"}}}) + violations = guard_mod.find_violations(path, "origin/dev") + assert [(v.table, v.column) for v in violations] == [("gadgets", "kind")] + + def test_case_b_helper_not_called_remains_red( + self, guard_mod, tmp_path: Path, monkeypatch + ) -> None: + body = ''' +async def _migration_v99_add_kind(conn) -> None: + """Module-level helper, but _post_init never calls it.""" + existing_cols = {row[1] for row in await conn.execute("PRAGMA table_info(gadgets)")} + if "kind" not in existing_cols: + await conn.execute("ALTER TABLE gadgets ADD COLUMN kind TEXT") + + +class GadgetStore: + SCHEMA = """ + CREATE TABLE IF NOT EXISTS gadgets ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL DEFAULT '' + ); + """ + + async def _post_init(self) -> None: + # Deliberately does NOT call _migration_v99_add_kind. + pass +''' + path = _write_store(tmp_path, "uncalled_helper.py", body) + self._baseline(guard_mod, monkeypatch, {"uncalled_helper.py": {"gadgets": {"id"}}}) + violations = guard_mod.find_violations(path, "origin/dev") + assert [(v.table, v.column) for v in violations] == [("gadgets", "kind")] + + def test_case_c_called_helper_goes_green( + self, guard_mod, tmp_path: Path, monkeypatch + ) -> None: + """Mirrors agent_registry_store.py: ALTER in a module-level helper + that _post_init DOES call.""" + body = ''' +async def _migration_v99_add_kind(conn) -> None: + """Module-level helper called by _post_init.""" + existing_cols = {row[1] for row in await conn.execute("PRAGMA table_info(gadgets)")} + if "kind" not in existing_cols: + await conn.execute("ALTER TABLE gadgets ADD COLUMN kind TEXT") + + +class GadgetStore: + SCHEMA = """ + CREATE TABLE IF NOT EXISTS gadgets ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL DEFAULT '' + ); + """ + + async def _post_init(self) -> None: + await _migration_v99_add_kind(self._db) +''' + path = _write_store(tmp_path, "called_helper.py", body) + self._baseline(guard_mod, monkeypatch, {"called_helper.py": {"gadgets": {"id"}}}) + assert guard_mod.find_violations(path, "origin/dev") == [] + + def test_recursive_helper_terminates( + self, guard_mod, tmp_path: Path, monkeypatch + ) -> None: + """A helper that calls itself (or a cycle between two helpers) must not + recurse forever.""" + body = ''' +async def _migration_v99_add_kind(conn) -> None: + await _migration_v99_add_kind(conn) + + +class GadgetStore: + SCHEMA = """ + CREATE TABLE IF NOT EXISTS gadgets ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL DEFAULT '' + ); + """ + + async def _post_init(self) -> None: + await _migration_v99_add_kind(self._db) +''' + path = _write_store(tmp_path, "recursive_helper.py", body) + self._baseline(guard_mod, monkeypatch, {"recursive_helper.py": {"gadgets": {"id"}}}) + violations = guard_mod.find_violations(path, "origin/dev") + assert [(v.table, v.column) for v in violations] == [("gadgets", "kind")] + + def test_fix_message_names_both_shapes(self, guard_mod) -> None: + v = guard_mod.Violation( + path=Path("x.py"), + table="t", + column="c", + detail="new column 'c' in CREATE TABLE t", + ) + msg = str(v) + assert "inline" in msg + assert "module-level helper" in msg