"""Schema integrity check — compares ORM metadata against the live DB. Called from ``init_db()`` after the engine is created to catch Alembic migration drift before the application starts serving requests. A schema mismatch is a hard startup failure: the server raises ``RuntimeError`` and refuses to boot rather than silently serving incorrect data. Design ------ - Uses SQLAlchemy ``inspect()`` run synchronously inside an async connection via ``conn.run_sync()``. - Checks: every ORM-declared table exists, every column exists, nullability matches, JSONB/JSON type is correct, VARCHAR lengths match, all declared indexes exist, and all declared FK constraints exist. - Skipped automatically for SQLite engines (in-memory test databases are created via ``create_all`` and are by definition correct). """ import logging from typing import TYPE_CHECKING from sqlalchemy import Table from sqlalchemy.engine import Connection from musehub.types.json_types import StrDict type AliasMapDict = dict[str, StrDict] if TYPE_CHECKING: from sqlalchemy.engine import Inspector from sqlalchemy.ext.asyncio import AsyncEngine from sqlalchemy.orm import DeclarativeBase logger = logging.getLogger(__name__) def _build_alias_map(base: type[DeclarativeBase]) -> AliasMapDict: """Build {table_name: {db_col_name: attr_key}} for aliased columns.""" alias_map: AliasMapDict = {} for mapper in base.registry.mappers: if not isinstance(mapper.local_table, Table): continue tname: str = mapper.local_table.name alias_map[tname] = { attr.columns[0].name: attr.key for attr in mapper.column_attrs if attr.key != attr.columns[0].name } return alias_map def _inspect_schema( inspector: Inspector, base: type[DeclarativeBase], alias_map: AliasMapDict, ) -> list[str]: """Run all schema integrity checks against a live inspector. Checks (in order): 1. Table existence 2. Column existence 3. Nullability match 4. JSONB vs JSON type drift 5. VARCHAR length mismatch 6. Index existence 7. FK existence Returns a list of human-readable mismatch strings (empty = clean). """ from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB # noqa: PLC0415 from sqlalchemy import String as SA_String # noqa: PLC0415 existing_tables: set[str] = set(inspector.get_table_names()) mismatches: list[str] = [] for table_name, table in base.metadata.tables.items(): if table_name not in existing_tables: mismatches.append(f"missing table: {table_name!r}") continue db_columns: dict[str, dict] = { c["name"]: c for c in inspector.get_columns(table_name) } table_aliases: StrDict = alias_map.get(table_name, {}) # ── Per-column checks ──────────────────────────────────────────────── for column in table.columns: col_name: str = column.name if col_name not in db_columns: python_attr = table_aliases.get(col_name) hint = f" (ORM attribute: {python_attr!r})" if python_attr else "" mismatches.append( f"{table_name}.{col_name!r}: column missing from DB{hint}" ) continue db_col = db_columns[col_name] orm_type = column.type db_type = db_col.get("type") # Nullability — most reliable cross-dialect check. # Primary key columns report nullable=False in the ORM but the # DB may omit the flag; skip when ORM nullable is None. orm_nullable: bool | None = column.nullable if orm_nullable is not None: db_nullable: bool = bool(db_col.get("nullable", True)) if orm_nullable != db_nullable: mismatches.append( f"{table_name}.{col_name!r}: nullable mismatch " f"(ORM={orm_nullable}, DB={db_nullable})" ) # JSONB vs JSON type drift — JSONB has different indexing semantics; # silently falling back to JSON loses GIN index support. if isinstance(orm_type, PG_JSONB): db_type_name = type(db_type).__name__.upper() if db_type is not None else "" if db_type_name != "JSONB": mismatches.append( f"{table_name}.{col_name!r}: type mismatch " f"(ORM=JSONB, DB={db_type_name or 'unknown'})" ) # VARCHAR length mismatch — catches truncation risk before data loss. if ( isinstance(orm_type, SA_String) and orm_type.length is not None and db_type is not None and hasattr(db_type, "length") and db_type.length is not None and orm_type.length != db_type.length ): mismatches.append( f"{table_name}.{col_name!r}: varchar length mismatch " f"(ORM={orm_type.length}, DB={db_type.length})" ) # ── Index checks ───────────────────────────────────────────────────── db_index_names: set[str] = { idx["name"] for idx in inspector.get_indexes(table_name) if idx.get("name") } for idx in table.indexes: if idx.name and idx.name not in db_index_names: mismatches.append( f"{table_name}: index {idx.name!r} missing from DB" ) # ── FK checks ──────────────────────────────────────────────────────── db_fks = inspector.get_foreign_keys(table_name) # Build (local_col, referred_table) pairs from DB db_fk_pairs: set[tuple[str, str]] = { (local_col, dbfk["referred_table"]) for dbfk in db_fks for local_col in dbfk["constrained_columns"] } for fk in table.foreign_keys: local_col_name = fk.parent.name try: referred_table_name = fk.column.table.name except Exception: continue # unresolved FK ref — skip if (local_col_name, referred_table_name) not in db_fk_pairs: mismatches.append( f"{table_name}.{local_col_name!r}: FK → {referred_table_name!r} missing from DB" ) return mismatches async def assert_schema_matches_orm( engine: AsyncEngine, base: type[DeclarativeBase], ) -> None: """Assert every ORM-declared table, column, index, and FK exists in the live DB. Compares ``base.metadata`` against the live schema via SQLAlchemy inspection. Raises ``RuntimeError`` listing every discrepancy if any are found. Args: engine: The async engine to inspect. base: The declarative base whose ``metadata`` describes the expected schema. All model classes must be imported before calling this function so they are registered on the metadata. Raises: RuntimeError: If any table, column, index, or FK declared in the ORM is absent from the live DB, or if a column's nullability, JSONB type, or VARCHAR length differs. """ url_str = str(engine.url) if url_str.startswith("sqlite"): logger.debug("schema_check: skipping for SQLite engine") return alias_map = _build_alias_map(base) def _inspect(conn: Connection) -> list[str]: from sqlalchemy import inspect as sa_inspect # noqa: PLC0415 inspector = sa_inspect(conn) return _inspect_schema(inspector, base, alias_map) async with engine.connect() as conn: mismatches: list[str] = await conn.run_sync(_inspect) if mismatches: detail = "\n".join(f" • {m}" for m in mismatches) raise RuntimeError( f"Schema drift detected — {len(mismatches)} mismatch(es) between " f"ORM metadata and live DB:\n{detail}\n\n" "Run Alembic migrations to resolve: alembic upgrade head" ) logger.info("✅ schema check passed — ORM and DB are in sync (%d tables)", len(base.metadata.tables)) def assert_no_orm_column_aliases(base: type[DeclarativeBase]) -> None: """Assert no ORM column has a Python key that differs from its DB name. A column alias — ``mapped_column("db_name", ...)`` on an attribute named differently — hides naming debt: the Python code uses one name, the DB uses another. This check runs statically against the ORM mapper (no DB connection required) and is suitable for CI. Uses the mapper registry rather than ``Table.columns`` because SQLAlchemy stores the Python attribute name only in the mapper's column_attrs, not in the ``Table`` metadata (where both ``key`` and ``name`` reflect the DB column name). Args: base: The declarative base to inspect. Raises: AssertionError: If any column has ``attr.key != col.name``. """ aliases: list[str] = [] for mapper in base.registry.mappers: if not isinstance(mapper.local_table, Table): continue table_name: str = mapper.local_table.name for attr in mapper.column_attrs: col = attr.columns[0] if attr.key != col.name: aliases.append( f"{table_name}: ORM key={attr.key!r} != DB name={col.name!r}" ) assert not aliases, ( "ORM column aliases detected — Python attribute name differs from DB column name.\n" "Remove the alias and write a migration to rename the DB column instead:\n" f"{chr(10).join(f' • {a}' for a in aliases)}" )