JhHardwareWRS_BackPoint/scripts/migrate_report_over_limit_snapshot.py

from pathlib import Path
import re
import sys

from sqlalchemy import text

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from app.database import engine  # noqa: E402


IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")

REQUIRED_COLUMNS = {
    "production_reports": {
        "has_over_limit": "BOOLEAN NOT NULL DEFAULT FALSE AFTER workload_rate",
        "over_limit_summary": "VARCHAR(500) NULL AFTER has_over_limit",
    },
    "production_report_items": {
        "over_limit_status": "VARCHAR(32) NOT NULL DEFAULT 'not_checked' AFTER scrap_qty",
        "over_limit_type": "VARCHAR(64) NULL AFTER over_limit_status",
        "over_limit_reason": "VARCHAR(500) NULL AFTER over_limit_type",
        "over_limit_checked_at": "DATETIME NULL AFTER over_limit_reason",
        "over_limit_check_source": "VARCHAR(32) NULL AFTER over_limit_checked_at",
        "over_limit_batch_no": "VARCHAR(128) NULL AFTER over_limit_check_source",
        "over_limit_product_gross_weight_kg": "DECIMAL(12, 4) NULL AFTER over_limit_batch_no",
        "over_limit_issued_weight_kg": "DECIMAL(18, 6) NULL AFTER over_limit_product_gross_weight_kg",
        "over_limit_max_reportable_qty": "DECIMAL(12, 2) NULL AFTER over_limit_issued_weight_kg",
        "over_limit_previous_good_qty": "DECIMAL(12, 2) NULL AFTER over_limit_max_reportable_qty",
        "over_limit_current_cumulative_qty": "DECIMAL(12, 2) NULL AFTER over_limit_previous_good_qty",
        "over_limit_current_report_qty": "DECIMAL(12, 2) NULL AFTER over_limit_current_cumulative_qty",
    },
}

REQUIRED_INDEXES = {
    "production_reports": {
        "idx_reports_over_limit": ("has_over_limit", "submitted_at"),
    },
    "production_report_items": {
        "idx_report_items_over_limit_status": ("over_limit_status",),
        "idx_report_items_process_batch": (
            "attendance_point_name",
            "project_no",
            "product_name",
            "raw_material_batch_no",
            "process_name",
        ),
    },
}


def quote_identifier(identifier: str) -> str:
    if not IDENTIFIER_RE.fullmatch(identifier):
        raise ValueError(f"unexpected SQL identifier: {identifier}")
    return f"`{identifier}`"


def table_exists(conn, table_name: str) -> bool:
    return bool(
        conn.execute(
            text(
                """
                SELECT COUNT(1)
                FROM information_schema.tables
                WHERE table_schema = DATABASE()
                  AND table_name = :table_name
                """
            ),
            {"table_name": table_name},
        ).scalar_one()
    )


def existing_columns(conn, table_name: str) -> set[str]:
    rows = conn.execute(
        text(
            """
            SELECT column_name
            FROM information_schema.columns
            WHERE table_schema = DATABASE()
              AND table_name = :table_name
            """
        ),
        {"table_name": table_name},
    ).all()
    return {row[0] for row in rows}


def existing_indexes(conn, table_name: str) -> dict[str, tuple[str, ...]]:
    rows = conn.execute(
        text(
            """
            SELECT index_name, column_name
            FROM information_schema.statistics
            WHERE table_schema = DATABASE()
              AND table_name = :table_name
            ORDER BY index_name, seq_in_index
            """
        ),
        {"table_name": table_name},
    ).all()
    indexes: dict[str, list[str]] = {}
    for index_name, column_name in rows:
        indexes.setdefault(index_name, []).append(column_name)
    return {index_name: tuple(columns) for index_name, columns in indexes.items()}


def add_missing_columns(conn, table_name: str) -> None:
    columns = existing_columns(conn, table_name)
    for column_name, column_definition in REQUIRED_COLUMNS[table_name].items():
        if column_name in columns:
            continue
        conn.execute(
            text(
                f"ALTER TABLE {quote_identifier(table_name)} "
                f"ADD COLUMN {quote_identifier(column_name)} {column_definition}"
            )
        )
        columns.add(column_name)


def create_missing_indexes(conn, table_name: str) -> None:
    indexes = existing_indexes(conn, table_name)
    for index_name, columns in REQUIRED_INDEXES[table_name].items():
        existing_columns_for_index = indexes.get(index_name)
        if existing_columns_for_index == columns:
            continue
        if existing_columns_for_index is not None:
            raise RuntimeError(
                f"{table_name}.{index_name} exists with columns "
                f"{existing_columns_for_index}, expected {columns}; "
                "please drop or rename the conflicting index before rerunning"
            )
        column_sql = ", ".join(quote_identifier(column) for column in columns)
        conn.execute(
            text(
                f"CREATE INDEX {quote_identifier(index_name)} "
                f"ON {quote_identifier(table_name)} ({column_sql})"
            )
        )


def verify_schema(conn) -> None:
    missing: list[str] = []
    for table_name, columns in REQUIRED_COLUMNS.items():
        if not table_exists(conn, table_name):
            missing.append(f"{table_name} table")
            continue

        existing_column_names = existing_columns(conn, table_name)
        for column_name in columns:
            if column_name not in existing_column_names:
                missing.append(f"{table_name}.{column_name}")

        indexes = existing_indexes(conn, table_name)
        for index_name, index_columns in REQUIRED_INDEXES[table_name].items():
            if indexes.get(index_name) != index_columns:
                missing.append(f"{table_name}.{index_name}")

    if missing:
        raise RuntimeError(f"missing report over limit snapshot objects: {', '.join(missing)}")


def main() -> None:
    with engine.begin() as conn:
        for table_name in REQUIRED_COLUMNS:
            if not table_exists(conn, table_name):
                raise RuntimeError(f"missing required table: {table_name}")
            add_missing_columns(conn, table_name)
            create_missing_indexes(conn, table_name)
        verify_schema(conn)
    print("report over limit snapshot migrated")


if __name__ == "__main__":
    main()