secondary-school-grade-archive/backend/app/services/migrate.py

from sqlalchemy import inspect, text

from app.core.config import settings
from app.core.database import engine


def run_migrations() -> None:
    """Apply lightweight schema updates for existing databases."""
    inspector = inspect(engine)
    tables = set(inspector.get_table_names())
    if "students" not in tables:
        return

    columns = {col["name"] for col in inspector.get_columns("students")}
    if "school_level" not in columns:
        with engine.begin() as conn:
            conn.execute(
                text(
                    "ALTER TABLE students ADD COLUMN school_level VARCHAR(32) "
                    "NOT NULL DEFAULT 'junior_high'"
                )
            )

    if "users" in tables:
        user_columns = {col["name"] for col in inspector.get_columns("users")}
        if "is_superuser" not in user_columns:
            with engine.begin() as conn:
                conn.execute(
                    text("ALTER TABLE users ADD COLUMN is_superuser BOOLEAN NOT NULL DEFAULT FALSE")
                )
                conn.execute(
                    text(
                        f"UPDATE users SET is_superuser = TRUE "
                        f"WHERE username = '{settings.ADMIN_DEFAULT_USERNAME}'"
                    )
                )

    if "wrong_questions" in tables:
        wq_columns = {col["name"] for col in inspector.get_columns("wrong_questions")}
        if "category" not in wq_columns:
            with engine.begin() as conn:
                conn.execute(
                    text(
                        "ALTER TABLE wrong_questions ADD COLUMN category VARCHAR(32) "
                        "NOT NULL DEFAULT 'regular'"
                    )
                )

    if "system_settings" in tables:
        ss_columns = {col["name"] for col in inspector.get_columns("system_settings")}
        alters: list[str] = []
        if "ai_provider" not in ss_columns:
            alters.append("ADD COLUMN ai_provider VARCHAR(16) NOT NULL DEFAULT 'ollama'")
        if "ollama_base_url" not in ss_columns:
            alters.append("ADD COLUMN ollama_base_url VARCHAR(256)")
        if "ollama_model" not in ss_columns:
            alters.append("ADD COLUMN ollama_model VARCHAR(128)")
        if "openai_base_url" not in ss_columns:
            alters.append("ADD COLUMN openai_base_url VARCHAR(256)")
        if "openai_model" not in ss_columns:
            alters.append("ADD COLUMN openai_model VARCHAR(128)")
        if "openai_api_key" not in ss_columns:
            alters.append("ADD COLUMN openai_api_key VARCHAR(512)")
        if "ocr_service_url" not in ss_columns:
            alters.append("ADD COLUMN ocr_service_url VARCHAR(256)")
        if alters:
            with engine.begin() as conn:
                for clause in alters:
                    conn.execute(text(f"ALTER TABLE system_settings {clause}"))

    if "wrong_questions" in tables:
        wq_columns = {col["name"] for col in inspector.get_columns("wrong_questions")}
        wq_alters: list[str] = []
        if "solution_approach" not in wq_columns:
            wq_alters.append("ADD COLUMN solution_approach TEXT")
        if "mark_regions_json" not in wq_columns:
            wq_alters.append("ADD COLUMN mark_regions_json TEXT")
        if "annotated_image_path" not in wq_columns:
            wq_alters.append("ADD COLUMN annotated_image_path VARCHAR(512)")
        if "cropped_image_path" not in wq_columns:
            wq_alters.append("ADD COLUMN cropped_image_path VARCHAR(512)")
        if "error_message" not in wq_columns:
            wq_alters.append("ADD COLUMN error_message TEXT")
        if wq_alters:
            with engine.begin() as conn:
                for clause in wq_alters:
                    conn.execute(text(f"ALTER TABLE wrong_questions {clause}"))

    if "subject_scores" in tables:
        ss_columns = {col["name"] for col in inspector.get_columns("subject_scores")}
        if "review_statuses_json" not in ss_columns:
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE subject_scores ADD COLUMN review_statuses_json TEXT"))

    if "system_settings" in tables:
        ss_columns = {col["name"] for col in inspector.get_columns("system_settings")}
        if "ai_review_enabled" not in ss_columns:
            with engine.begin() as conn:
                conn.execute(
                    text(
                        "ALTER TABLE system_settings ADD COLUMN ai_review_enabled BOOLEAN NOT NULL DEFAULT TRUE"
                    )
                )

    student_columns = {col["name"] for col in inspector.get_columns("students")}
    student_alters: list[str] = []
    if "school_name" not in student_columns:
        student_alters.append("ADD COLUMN school_name VARCHAR(128)")
    if "avatar_path" not in student_columns:
        student_alters.append("ADD COLUMN avatar_path VARCHAR(512)")
    if student_alters:
        with engine.begin() as conn:
            for clause in student_alters:
                conn.execute(text(f"ALTER TABLE students {clause}"))