fix: 避免报工归属索引重复

scripts/migrate_report_allocations.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+import re
 import sys
 
 from sqlalchemy import text
@@ -9,30 +10,151 @@ sys.path.insert(0, str(ROOT))
 from app.database import engine  # noqa: E402
 
 
-def index_exists(conn, table_name: str, index_name: str) -> bool:
+TABLE_NAME = "production_report_allocations"
+IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+REQUIRED_COLUMNS = {
+    "id",
+    "report_id",
+    "report_item_id",
+    "attendance_point_name",
+    "employee_phone",
+    "allocation_date",
+    "day_minutes",
+    "overtime_minutes",
+    "night_minutes",
+    "effective_minutes",
+    "good_qty",
+    "defect_qty",
+    "scrap_qty",
+    "changeover_count",
+    "reference_wage",
+    "created_at",
+    "updated_at",
+}
+REQUIRED_INDEXES = {
+    "idx_report_allocations_report": ("report_id",),
+    "idx_report_allocations_item": ("report_item_id",),
+    "idx_report_allocations_date": ("allocation_date",),
+    "idx_report_allocations_point_date": ("attendance_point_name", "allocation_date"),
+    "idx_report_allocations_employee_date": ("employee_phone", "allocation_date"),
+}
+REQUIRED_FOREIGN_KEYS = {
+    "report_id": ("production_reports", "id", "CASCADE"),
+    "report_item_id": ("production_report_items", "id", "CASCADE"),
+}
+
+
+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) -> bool:
     return bool(
         conn.execute(
             text(
                 """
                 SELECT COUNT(1)
-                FROM information_schema.statistics
+                FROM information_schema.tables
                 WHERE table_schema = DATABASE()
                   AND table_name = :table_name
-                  AND index_name = :index_name
                 """
             ),
-            {"table_name": table_name, "index_name": index_name},
+            {"table_name": TABLE_NAME},
-        ).scalar()
+        ).scalar_one()
     )
 
 
-def create_index(conn, table_name: str, index_name: str, columns: str) -> None:
+def existing_columns(conn) -> set[str]:
-    if not index_exists(conn, table_name, index_name):
+    rows = conn.execute(
-        conn.execute(text(f"CREATE INDEX {index_name} ON {table_name} ({columns})"))
+        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 main() -> None:
+def existing_indexes(conn) -> dict[str, tuple[str, ...]]:
-    with engine.begin() as conn:
+    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 existing_foreign_keys(conn) -> dict[str, tuple[str, str, str]]:
+    rows = conn.execute(
+        text(
+            """
+            SELECT k.column_name, k.referenced_table_name, k.referenced_column_name, r.delete_rule
+            FROM information_schema.key_column_usage k
+            JOIN information_schema.referential_constraints r
+              ON r.constraint_schema = k.constraint_schema
+             AND r.constraint_name = k.constraint_name
+            WHERE k.table_schema = DATABASE()
+              AND k.table_name = :table_name
+              AND k.referenced_table_name IS NOT NULL
+            ORDER BY k.constraint_name, k.ordinal_position
+            """
+        ),
+        {"table_name": TABLE_NAME},
+    ).all()
+    return {row[0]: (row[1], row[2], row[3]) for row in rows}
+
+
+def find_equivalent_index(indexes: dict[str, tuple[str, ...]], columns: tuple[str, ...]) -> str | None:
+    for index_name, index_columns in indexes.items():
+        if index_name != "PRIMARY" and index_columns == columns:
+            return index_name
+    return None
+
+
+def create_index(conn, index_name: str, columns: tuple[str, ...]) -> None:
+    if index_name not in REQUIRED_INDEXES or REQUIRED_INDEXES[index_name] != columns:
+        raise ValueError(f"unexpected index definition: {index_name} {columns}")
+
+    indexes = existing_indexes(conn)
+    if index_name in indexes:
+        return
+
+    equivalent_index = find_equivalent_index(indexes, columns)
+    if equivalent_index:
+        conn.execute(
+            text(
+                f"ALTER TABLE {quote_identifier(TABLE_NAME)} "
+                f"RENAME INDEX {quote_identifier(equivalent_index)} TO {quote_identifier(index_name)}"
+            )
+        )
+        return
+
+    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 ensure_schema(conn) -> None:
     conn.execute(
         text(
             """
@@ -55,6 +177,11 @@ def main() -> None:
                 created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
                 updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                 PRIMARY KEY (id),
+                KEY idx_report_allocations_report (report_id),
+                KEY idx_report_allocations_item (report_item_id),
+                KEY idx_report_allocations_date (allocation_date),
+                KEY idx_report_allocations_point_date (attendance_point_name, allocation_date),
+                KEY idx_report_allocations_employee_date (employee_phone, allocation_date),
                 CONSTRAINT fk_report_allocations_report
                     FOREIGN KEY (report_id) REFERENCES production_reports (id)
                     ON DELETE CASCADE,
@@ -66,22 +193,59 @@ def main() -> None:
         )
     )
 
-        create_index(conn, "production_report_allocations", "idx_report_allocations_report", "report_id")
+    for index_name, columns in REQUIRED_INDEXES.items():
-        create_index(conn, "production_report_allocations", "idx_report_allocations_item", "report_item_id")
+        create_index(conn, index_name, columns)
-        create_index(conn, "production_report_allocations", "idx_report_allocations_date", "allocation_date")
+
-        create_index(
+
-            conn,
+def verify_schema(conn) -> None:
-            "production_report_allocations",
+    if not table_exists(conn):
-            "idx_report_allocations_point_date",
+        raise RuntimeError(f"{TABLE_NAME} table is missing")
-            "attendance_point_name, allocation_date",
+
+    missing_columns = sorted(REQUIRED_COLUMNS - existing_columns(conn))
+    if missing_columns:
+        raise RuntimeError(f"{TABLE_NAME} missing columns: {', '.join(missing_columns)}")
+
+    indexes = existing_indexes(conn)
+    missing_indexes = [index_name for index_name in REQUIRED_INDEXES if index_name not in indexes]
+    if missing_indexes:
+        raise RuntimeError(f"{TABLE_NAME} missing indexes: {', '.join(missing_indexes)}")
+
+    wrong_indexes = [
+        f"{index_name}({', '.join(indexes[index_name])})"
+        for index_name, columns in REQUIRED_INDEXES.items()
+        if indexes.get(index_name) != columns
+    ]
+    if wrong_indexes:
+        raise RuntimeError(f"{TABLE_NAME} incompatible indexes: {', '.join(wrong_indexes)}")
+
+    duplicate_single_column_indexes = [
+        index_name
+        for index_name, columns in indexes.items()
+        if (
+            (columns == ("report_id",) and index_name != "idx_report_allocations_report")
+            or (columns == ("report_item_id",) and index_name != "idx_report_allocations_item")
         )
-        create_index(
+    ]
-            conn,
+    if duplicate_single_column_indexes:
-            "production_report_allocations",
+        raise RuntimeError(
-            "idx_report_allocations_employee_date",
+            f"{TABLE_NAME} duplicate FK indexes: {', '.join(sorted(duplicate_single_column_indexes))}"
-            "employee_phone, allocation_date",
         )
 
+    foreign_keys = existing_foreign_keys(conn)
+    missing_foreign_keys = [
+        column
+        for column, definition in REQUIRED_FOREIGN_KEYS.items()
+        if foreign_keys.get(column) != definition
+    ]
+    if missing_foreign_keys:
+        raise RuntimeError(f"{TABLE_NAME} missing foreign keys: {', '.join(missing_foreign_keys)}")
+
+
+def main() -> None:
+    with engine.begin() as conn:
+        ensure_schema(conn)
+        verify_schema(conn)
+
     print("report allocations schema migrated")