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()