leocrm/scripts/schema_drift_check.py

"""Schema drift checker — compares SQLAlchemy models against production DB.

Usage: python scripts/schema_drift_check.py
Outputs a list of all columns where the model definition differs from the DB.
"""
from __future__ import annotations

import asyncio
import sys
from typing import Any

from sqlalchemy import inspect, text
from sqlalchemy.ext.asyncio import create_async_engine

# Import all models so they register with Base.metadata
import app.core.db  # noqa
from app.core.db import Base
from app.plugins.registry import get_registry

# Discover builtins to load plugin models
r = get_registry()
r.discover_builtins()
for name in r.list_discovered():
    try:
        import importlib
        importlib.import_module(f"app.plugins.builtins.{name}.models")
    except Exception:
        pass

# Also import known model modules
import app.models  # noqa
from app.models.notification import NotificationType  # noqa
from app.models.compliance import ComplianceIncident  # noqa
from app.plugins.builtins.knowledge.models import KnowledgeExtraction  # noqa
from app.plugins.builtins.self_improvement.models import (  # noqa
    ImprovementSignal, ImprovementPattern, ImprovementProposal, ImpactMeasurement,
)


async def check_drift(db_url: str):
    engine = create_async_engine(db_url)
    drifts = []

    async with engine.connect() as conn:
        # Get all DB tables and columns
        result = await conn.execute(text("""
            SELECT table_name, column_name, data_type, character_maximum_length,
                   is_nullable, column_default
            FROM information_schema.columns
            WHERE table_schema = 'public'
              AND table_name NOT LIKE 'pg_%'
              AND table_name NOT LIKE 'alembic_%'
            ORDER BY table_name, ordinal_position
        """))
        db_columns = {}
        for row in result:
            table = row[0]
            if table not in db_columns:
                db_columns[table] = {}
            db_columns[table][row[1]] = {
                "data_type": row[2],
                "max_length": row[3],
                "nullable": row[4],
                "default": row[5],
            }

    # Compare with models
    for table_name, table_obj in Base.metadata.tables.items():
        if table_name not in db_columns:
            drifts.append({
                "table": table_name,
                "issue": "TABLE MISSING IN DB",
                "column": "*",
                "model": "exists",
                "db": "missing",
            })
            continue

        for col_name, col_obj in table_obj.columns.items():
            if col_name not in db_columns[table_name]:
                drifts.append({
                    "table": table_name,
                    "issue": "COLUMN MISSING IN DB",
                    "column": col_name,
                    "model": str(col_obj.type),
                    "db": "missing",
                })
                continue

            db_col = db_columns[table_name][col_name]

            # Compare VARCHAR lengths
            model_type = str(col_obj.type)
            db_type = db_col["data_type"]
            db_max = db_col["max_length"]

            if "VARCHAR" in model_type or "character varying" in db_type:
                # Extract model length
                model_len = None
                if "(" in model_type:
                    try:
                        model_len = int(model_type.split("(")[1].split(")")[0])
                    except (ValueError, IndexError):
                        pass

                if model_len is not None and db_max is not None and model_len != db_max:
                    drifts.append({
                        "table": table_name,
                        "issue": "VARCHAR LENGTH MISMATCH",
                        "column": col_name,
                        "model": f"VARCHAR({model_len})",
                        "db": f"VARCHAR({db_max})",
                    })

            # Check for missing columns in model (extra in DB)
        for db_col_name in db_columns[table_name]:
            if db_col_name not in [c.name for c in table_obj.columns]:
                drifts.append({
                    "table": table_name,
                    "issue": "COLUMN IN DB NOT IN MODEL",
                    "column": db_col_name,
                    "model": "missing",
                    "db": db_columns[table_name][db_col_name]["data_type"],
                })

    # Check for tables in DB not in model
    for db_table in db_columns:
        if db_table not in Base.metadata.tables:
            drifts.append({
                "table": db_table,
                "issue": "TABLE IN DB NOT IN MODEL",
                "column": "*",
                "model": "missing",
                "db": "exists",
            })

    await engine.dispose()
    return drifts


async def main():
    db_url = "postgresql+asyncpg://crm_user:86FkF5vJ_qKYgO6Myj0eQ4Dtm3Dyb1ge@localhost:5432/crm_db"
    if len(sys.argv) > 1:
        db_url = sys.argv[1]

    print("Checking schema drift between models and DB...")
    print(f"DB URL: {db_url.split('@')[1]}")
    print()

    drifts = await check_drift(db_url)

    if not drifts:
        print("No drifts found! Schema is in sync.")
        return

    print(f"Found {len(drifts)} drift(s):")
    print()
    for d in drifts:
        print(f"  [{d['issue']}] {d['table']}.{d['column']}")
        print(f"    Model: {d['model']}")
        print(f"    DB:    {d['db']}")
        print()

    # Group by issue type
    by_issue = {}
    for d in drifts:
        by_issue.setdefault(d["issue"], []).append(d)
    print("Summary:")
    for issue, items in by_issue.items():
        print(f"  {issue}: {len(items)}")


if __name__ == "__main__":
    asyncio.run(main())