xsk-education-management/app/scripts/migrate_text_to_sqlite.py

from __future__ import annotations

import argparse
from datetime import datetime
import hashlib
import json
import os
from pathlib import Path
import shutil
import sqlite3
import sys
import tarfile


APP_ROOT = Path(__file__).resolve().parents[1]
REPO_ROOT = APP_ROOT.parent
if str(APP_ROOT) not in sys.path:
    sys.path.insert(0, str(APP_ROOT))

from app.db import SCHEMA_VERSION, connect, initialize_schema  # noqa: E402
from app.repository import replace_database_from_paths, source_file_hashes  # noqa: E402


EXPECTED_COUNTS = {
    "records": 1548,
    "accounts": 45,
    "teachers": 13,
}


def sha256_path(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def data_paths(data_root: Path) -> dict[str, Path]:
    return {
        "classnotes_path": data_root / "classnotes.txt",
        "accounts_path": data_root / "学生课时账户.md",
        "teachers_path": data_root / "教师档案.md",
        "tasks_path": data_root / "admin_tasks.json",
        "summaries_root": data_root / "course_summaries",
        "state_path": data_root / "course_summary_state.json",
        "operation_logs_path": data_root / "operation_logs.jsonl",
    }


def validate_counts(summary: dict, *, strict_expected: bool) -> list[str]:
    errors: list[str] = []
    for key in ["records", "accounts", "teachers", "admin_tasks", "operation_logs", "course_summaries"]:
        if int(summary.get(key) or 0) < 0:
            errors.append(f"{key} 数量异常")
    if strict_expected:
        for key, expected in EXPECTED_COUNTS.items():
            actual = int(summary.get(key) or 0)
            if actual != expected:
                errors.append(f"{key} 数量应为 {expected}，实际为 {actual}")
        mismatch_count = len(summary.get("balance_mismatches") or [])
        if mismatch_count != 4:
            errors.append(f"余额差异应为 4 个学生，实际为 {mismatch_count}")
    return errors


def validate_database(db_path: Path, summary: dict, *, strict_expected: bool) -> dict:
    errors = validate_counts(summary, strict_expected=strict_expected)
    with connect(db_path) as conn:
        table_counts = {
            "records": conn.execute("SELECT COUNT(*) FROM class_records").fetchone()[0],
            "accounts": conn.execute("SELECT COUNT(*) FROM students").fetchone()[0],
            "teachers": conn.execute("SELECT COUNT(*) FROM teachers").fetchone()[0],
            "admin_tasks": conn.execute("SELECT COUNT(*) FROM admin_tasks").fetchone()[0],
            "operation_logs": conn.execute("SELECT COUNT(*) FROM operation_logs").fetchone()[0],
            "course_summaries": conn.execute("SELECT COUNT(*) FROM course_summaries").fetchone()[0],
        }
        for key, value in table_counts.items():
            if int(summary.get(key) or 0) != int(value):
                errors.append(f"SQLite {key} 数量不一致: source={summary.get(key)} sqlite={value}")
        duplicate_records = conn.execute(
            """
            SELECT record_key, COUNT(*) AS c
            FROM class_records
            GROUP BY record_key
            HAVING c > 1
            LIMIT 1
            """
        ).fetchone()
        if duplicate_records:
            errors.append(f"SQLite 中存在重复课程记录: {duplicate_records['record_key']}")
        missing_account = conn.execute(
            """
            SELECT student
            FROM class_records
            WHERE student_id IS NULL
            LIMIT 1
            """
        ).fetchone()
        if missing_account:
            errors.append(f"SQLite 中存在没有账户的上课学生: {missing_account['student']}")
    return {"ok": not errors, "errors": errors, "table_counts": table_counts}


def write_audit(conn: sqlite3.Connection, payload: dict) -> None:
    conn.execute(
        "INSERT INTO migration_audit(created_at, kind, payload_json) VALUES(?, ?, ?)",
        (
            datetime.now().isoformat(timespec="seconds"),
            "text_to_sqlite",
            json.dumps(payload, ensure_ascii=False, sort_keys=True),
        ),
    )


def create_archive(data_root: Path, archives_root: Path, report_path: Path, timestamp: str) -> tuple[Path, Path]:
    archives_root.mkdir(parents=True, exist_ok=True)
    archive_path = archives_root / f"text-source-before-sqlite-{timestamp}.tar.gz"
    with tarfile.open(archive_path, "w:gz") as archive:
        for name in [
            "classnotes.txt",
            "学生课时账户.md",
            "教师档案.md",
            "admin_tasks.json",
            "operation_logs.jsonl",
            "course_summary_state.json",
        ]:
            path = data_root / name
            if path.exists():
                archive.add(path, arcname=name)
        summaries_root = data_root / "course_summaries"
        if summaries_root.exists():
            archive.add(summaries_root, arcname="course_summaries")
        archive.add(report_path, arcname=report_path.name)
    sha_path = archive_path.with_suffix(archive_path.suffix + ".sha256")
    sha_path.write_text(f"{sha256_path(archive_path)}  {archive_path.name}\n", encoding="utf-8")
    return archive_path, sha_path


def migrate(args: argparse.Namespace) -> dict:
    data_root = args.data_root.resolve()
    db_path = args.db_path.resolve()
    tmp_path = Path(str(db_path) + ".tmp")
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    report_path = data_root / f"sqlite_migration_report_{timestamp}.json"

    if tmp_path.exists():
        tmp_path.unlink()
    if tmp_path.with_suffix(tmp_path.suffix + "-wal").exists():
        tmp_path.with_suffix(tmp_path.suffix + "-wal").unlink()
    if tmp_path.with_suffix(tmp_path.suffix + "-shm").exists():
        tmp_path.with_suffix(tmp_path.suffix + "-shm").unlink()

    with connect(tmp_path) as conn:
        initialize_schema(conn)
        conn.commit()
        conn.execute("BEGIN IMMEDIATE")
        try:
            summary = replace_database_from_paths(
                conn,
                **data_paths(data_root),
                allow_balance_adjustments=False,
            )
            conn.commit()
        except Exception:
            conn.rollback()
            raise

    validation = validate_database(tmp_path, summary, strict_expected=not args.no_strict_expected)
    source_hashes = source_file_hashes(
        [
            data_root / "classnotes.txt",
            data_root / "学生课时账户.md",
            data_root / "教师档案.md",
            data_root / "admin_tasks.json",
            data_root / "operation_logs.jsonl",
            data_root / "course_summary_state.json",
        ]
    )
    report = {
        "created_at": datetime.now().isoformat(timespec="seconds"),
        "schema_version": SCHEMA_VERSION,
        "data_root": str(data_root),
        "db_path": str(db_path),
        "summary": summary,
        "validation": validation,
        "source_hashes": source_hashes,
        "dry_run": bool(args.dry_run),
    }
    report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    if not validation["ok"]:
        raise SystemExit("迁移校验失败：\n" + "\n".join(validation["errors"]))

    with connect(tmp_path) as conn:
        write_audit(conn, {**report, "report_path": str(report_path)})

    if args.dry_run:
        print(json.dumps({**report, "tmp_db_path": str(tmp_path), "report_path": str(report_path)}, ensure_ascii=False, indent=2))
        return report

    db_path.parent.mkdir(parents=True, exist_ok=True)
    if db_path.exists():
        backup_db = db_path.with_name(f"{db_path.name}.before-sqlite-migration-{timestamp}")
        shutil.copy2(db_path, backup_db)
        report["previous_db_backup"] = str(backup_db)
    os.replace(tmp_path, db_path)
    for suffix in ["-wal", "-shm"]:
        sidecar = Path(str(tmp_path) + suffix)
        if sidecar.exists():
            os.replace(sidecar, Path(str(db_path) + suffix))

    archive_path, sha_path = create_archive(data_root, args.archives_root.resolve(), report_path, timestamp)
    report["archive_path"] = str(archive_path)
    report["archive_sha256_path"] = str(sha_path)
    report["dry_run"] = False
    report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({**report, "report_path": str(report_path)}, ensure_ascii=False, indent=2))
    return report


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="将新时空纯文本教务数据迁移到 SQLite")
    parser.add_argument("--data-root", type=Path, default=REPO_ROOT / "data")
    parser.add_argument("--db-path", type=Path, default=Path("/data/xsk_education.db"))
    parser.add_argument("--archives-root", type=Path, default=REPO_ROOT / "archives")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--no-strict-expected", action="store_true", help="不校验当前生产数据的固定计数")
    return parser.parse_args()


if __name__ == "__main__":
    migrate(parse_args())