schema_check.py
python
sha256:969ddc5e88776e70af33c016b8777bb721356508ae5f304b3fb46c451494ece7
test(mwp3): Phase 4 — fix adjacent suite regressions, lock …
Sonnet 4.6
22 days ago
| 1 | |
| 2 | """Schema integrity check — compares ORM metadata against the live DB. |
| 3 | |
| 4 | Called from ``init_db()`` after the engine is created to catch Alembic |
| 5 | migration drift before the application starts serving requests. A schema |
| 6 | mismatch is a hard startup failure: the server raises ``RuntimeError`` and |
| 7 | refuses to boot rather than silently serving incorrect data. |
| 8 | |
| 9 | Design |
| 10 | ------ |
| 11 | - Uses SQLAlchemy ``inspect()`` run synchronously inside an async connection |
| 12 | via ``conn.run_sync()``. |
| 13 | - Checks: every ORM-declared table exists, every column exists, nullability |
| 14 | matches, JSONB/JSON type is correct, VARCHAR lengths match, all declared |
| 15 | indexes exist, and all declared FK constraints exist. |
| 16 | - Skipped automatically for SQLite engines (in-memory test databases are |
| 17 | created via ``create_all`` and are by definition correct). |
| 18 | """ |
| 19 | |
| 20 | import logging |
| 21 | from typing import TYPE_CHECKING |
| 22 | |
| 23 | from sqlalchemy import Table |
| 24 | from sqlalchemy.engine import Connection |
| 25 | from musehub.types.json_types import StrDict |
| 26 | |
| 27 | type AliasMapDict = dict[str, StrDict] |
| 28 | |
| 29 | if TYPE_CHECKING: |
| 30 | from sqlalchemy.engine import Inspector |
| 31 | from sqlalchemy.ext.asyncio import AsyncEngine |
| 32 | from sqlalchemy.orm import DeclarativeBase |
| 33 | |
| 34 | logger = logging.getLogger(__name__) |
| 35 | |
| 36 | |
| 37 | def _build_alias_map(base: type[DeclarativeBase]) -> AliasMapDict: |
| 38 | """Build {table_name: {db_col_name: attr_key}} for aliased columns.""" |
| 39 | alias_map: AliasMapDict = {} |
| 40 | for mapper in base.registry.mappers: |
| 41 | if not isinstance(mapper.local_table, Table): |
| 42 | continue |
| 43 | tname: str = mapper.local_table.name |
| 44 | alias_map[tname] = { |
| 45 | attr.columns[0].name: attr.key |
| 46 | for attr in mapper.column_attrs |
| 47 | if attr.key != attr.columns[0].name |
| 48 | } |
| 49 | return alias_map |
| 50 | |
| 51 | |
| 52 | def _inspect_schema( |
| 53 | inspector: Inspector, |
| 54 | base: type[DeclarativeBase], |
| 55 | alias_map: AliasMapDict, |
| 56 | ) -> list[str]: |
| 57 | """Run all schema integrity checks against a live inspector. |
| 58 | |
| 59 | Checks (in order): |
| 60 | 1. Table existence |
| 61 | 2. Column existence |
| 62 | 3. Nullability match |
| 63 | 4. JSONB vs JSON type drift |
| 64 | 5. VARCHAR length mismatch |
| 65 | 6. Index existence |
| 66 | 7. FK existence |
| 67 | |
| 68 | Returns a list of human-readable mismatch strings (empty = clean). |
| 69 | """ |
| 70 | from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB # noqa: PLC0415 |
| 71 | from sqlalchemy import String as SA_String # noqa: PLC0415 |
| 72 | |
| 73 | existing_tables: set[str] = set(inspector.get_table_names()) |
| 74 | mismatches: list[str] = [] |
| 75 | |
| 76 | for table_name, table in base.metadata.tables.items(): |
| 77 | if table_name not in existing_tables: |
| 78 | mismatches.append(f"missing table: {table_name!r}") |
| 79 | continue |
| 80 | |
| 81 | db_columns: dict[str, dict] = { |
| 82 | c["name"]: c for c in inspector.get_columns(table_name) |
| 83 | } |
| 84 | table_aliases: StrDict = alias_map.get(table_name, {}) |
| 85 | |
| 86 | # ── Per-column checks ──────────────────────────────────────────────── |
| 87 | for column in table.columns: |
| 88 | col_name: str = column.name |
| 89 | if col_name not in db_columns: |
| 90 | python_attr = table_aliases.get(col_name) |
| 91 | hint = f" (ORM attribute: {python_attr!r})" if python_attr else "" |
| 92 | mismatches.append( |
| 93 | f"{table_name}.{col_name!r}: column missing from DB{hint}" |
| 94 | ) |
| 95 | continue |
| 96 | |
| 97 | db_col = db_columns[col_name] |
| 98 | orm_type = column.type |
| 99 | db_type = db_col.get("type") |
| 100 | |
| 101 | # Nullability — most reliable cross-dialect check. |
| 102 | # Primary key columns report nullable=False in the ORM but the |
| 103 | # DB may omit the flag; skip when ORM nullable is None. |
| 104 | orm_nullable: bool | None = column.nullable |
| 105 | if orm_nullable is not None: |
| 106 | db_nullable: bool = bool(db_col.get("nullable", True)) |
| 107 | if orm_nullable != db_nullable: |
| 108 | mismatches.append( |
| 109 | f"{table_name}.{col_name!r}: nullable mismatch " |
| 110 | f"(ORM={orm_nullable}, DB={db_nullable})" |
| 111 | ) |
| 112 | |
| 113 | # JSONB vs JSON type drift — JSONB has different indexing semantics; |
| 114 | # silently falling back to JSON loses GIN index support. |
| 115 | if isinstance(orm_type, PG_JSONB): |
| 116 | db_type_name = type(db_type).__name__.upper() if db_type is not None else "" |
| 117 | if db_type_name != "JSONB": |
| 118 | mismatches.append( |
| 119 | f"{table_name}.{col_name!r}: type mismatch " |
| 120 | f"(ORM=JSONB, DB={db_type_name or 'unknown'})" |
| 121 | ) |
| 122 | |
| 123 | # VARCHAR length mismatch — catches truncation risk before data loss. |
| 124 | if ( |
| 125 | isinstance(orm_type, SA_String) |
| 126 | and orm_type.length is not None |
| 127 | and db_type is not None |
| 128 | and hasattr(db_type, "length") |
| 129 | and db_type.length is not None |
| 130 | and orm_type.length != db_type.length |
| 131 | ): |
| 132 | mismatches.append( |
| 133 | f"{table_name}.{col_name!r}: varchar length mismatch " |
| 134 | f"(ORM={orm_type.length}, DB={db_type.length})" |
| 135 | ) |
| 136 | |
| 137 | # ── Index checks ───────────────────────────────────────────────────── |
| 138 | db_index_names: set[str] = { |
| 139 | idx["name"] |
| 140 | for idx in inspector.get_indexes(table_name) |
| 141 | if idx.get("name") |
| 142 | } |
| 143 | for idx in table.indexes: |
| 144 | if idx.name and idx.name not in db_index_names: |
| 145 | mismatches.append( |
| 146 | f"{table_name}: index {idx.name!r} missing from DB" |
| 147 | ) |
| 148 | |
| 149 | # ── FK checks ──────────────────────────────────────────────────────── |
| 150 | db_fks = inspector.get_foreign_keys(table_name) |
| 151 | # Build (local_col, referred_table) pairs from DB |
| 152 | db_fk_pairs: set[tuple[str, str]] = { |
| 153 | (local_col, dbfk["referred_table"]) |
| 154 | for dbfk in db_fks |
| 155 | for local_col in dbfk["constrained_columns"] |
| 156 | } |
| 157 | for fk in table.foreign_keys: |
| 158 | local_col_name = fk.parent.name |
| 159 | try: |
| 160 | referred_table_name = fk.column.table.name |
| 161 | except Exception: |
| 162 | continue # unresolved FK ref — skip |
| 163 | if (local_col_name, referred_table_name) not in db_fk_pairs: |
| 164 | mismatches.append( |
| 165 | f"{table_name}.{local_col_name!r}: FK → {referred_table_name!r} missing from DB" |
| 166 | ) |
| 167 | |
| 168 | return mismatches |
| 169 | |
| 170 | |
| 171 | async def assert_schema_matches_orm( |
| 172 | engine: AsyncEngine, |
| 173 | base: type[DeclarativeBase], |
| 174 | ) -> None: |
| 175 | """Assert every ORM-declared table, column, index, and FK exists in the live DB. |
| 176 | |
| 177 | Compares ``base.metadata`` against the live schema via SQLAlchemy |
| 178 | inspection. Raises ``RuntimeError`` listing every discrepancy if any |
| 179 | are found. |
| 180 | |
| 181 | Args: |
| 182 | engine: The async engine to inspect. |
| 183 | base: The declarative base whose ``metadata`` describes the expected |
| 184 | schema. All model classes must be imported before calling this |
| 185 | function so they are registered on the metadata. |
| 186 | |
| 187 | Raises: |
| 188 | RuntimeError: If any table, column, index, or FK declared in the ORM |
| 189 | is absent from the live DB, or if a column's nullability, JSONB |
| 190 | type, or VARCHAR length differs. |
| 191 | """ |
| 192 | url_str = str(engine.url) |
| 193 | if url_str.startswith("sqlite"): |
| 194 | logger.debug("schema_check: skipping for SQLite engine") |
| 195 | return |
| 196 | |
| 197 | alias_map = _build_alias_map(base) |
| 198 | |
| 199 | def _inspect(conn: Connection) -> list[str]: |
| 200 | from sqlalchemy import inspect as sa_inspect # noqa: PLC0415 |
| 201 | inspector = sa_inspect(conn) |
| 202 | return _inspect_schema(inspector, base, alias_map) |
| 203 | |
| 204 | async with engine.connect() as conn: |
| 205 | mismatches: list[str] = await conn.run_sync(_inspect) |
| 206 | |
| 207 | if mismatches: |
| 208 | detail = "\n".join(f" • {m}" for m in mismatches) |
| 209 | raise RuntimeError( |
| 210 | f"Schema drift detected — {len(mismatches)} mismatch(es) between " |
| 211 | f"ORM metadata and live DB:\n{detail}\n\n" |
| 212 | "Run Alembic migrations to resolve: alembic upgrade head" |
| 213 | ) |
| 214 | |
| 215 | logger.info("✅ schema check passed — ORM and DB are in sync (%d tables)", len(base.metadata.tables)) |
| 216 | |
| 217 | |
| 218 | def assert_no_orm_column_aliases(base: type[DeclarativeBase]) -> None: |
| 219 | """Assert no ORM column has a Python key that differs from its DB name. |
| 220 | |
| 221 | A column alias — ``mapped_column("db_name", ...)`` on an attribute named |
| 222 | differently — hides naming debt: the Python code uses one name, the DB |
| 223 | uses another. This check runs statically against the ORM mapper |
| 224 | (no DB connection required) and is suitable for CI. |
| 225 | |
| 226 | Uses the mapper registry rather than ``Table.columns`` because SQLAlchemy |
| 227 | stores the Python attribute name only in the mapper's column_attrs, not in |
| 228 | the ``Table`` metadata (where both ``key`` and ``name`` reflect the DB |
| 229 | column name). |
| 230 | |
| 231 | Args: |
| 232 | base: The declarative base to inspect. |
| 233 | |
| 234 | Raises: |
| 235 | AssertionError: If any column has ``attr.key != col.name``. |
| 236 | """ |
| 237 | aliases: list[str] = [] |
| 238 | for mapper in base.registry.mappers: |
| 239 | if not isinstance(mapper.local_table, Table): |
| 240 | continue |
| 241 | table_name: str = mapper.local_table.name |
| 242 | for attr in mapper.column_attrs: |
| 243 | col = attr.columns[0] |
| 244 | if attr.key != col.name: |
| 245 | aliases.append( |
| 246 | f"{table_name}: ORM key={attr.key!r} != DB name={col.name!r}" |
| 247 | ) |
| 248 | assert not aliases, ( |
| 249 | "ORM column aliases detected — Python attribute name differs from DB column name.\n" |
| 250 | "Remove the alias and write a migration to rename the DB column instead:\n" |
| 251 | f"{chr(10).join(f' • {a}' for a in aliases)}" |
| 252 | ) |
File History
2 commits
sha256:969ddc5e88776e70af33c016b8777bb721356508ae5f304b3fb46c451494ece7
test(mwp3): Phase 4 — fix adjacent suite regressions, lock …
Sonnet 4.6
22 days ago
sha256:1749b9cc5cd2583c56d3261c4c00a342c27435f3946d4bc3e70f76aa2795458a
docs(mwp-3): TDD implementation plan for job-enqueue idempo…
Opus 4.8
22 days ago