document-management/tools/verify_full_demo_import.py

"""Verify a DMS data export by importing it into an isolated scratch database."""

from __future__ import annotations

import argparse
import os
from pathlib import Path

from alembic import command
from alembic.config import Config
from dotenv import dotenv_values
from pymysql.constants import CLIENT
from sqlalchemy import create_engine, text
from sqlalchemy.engine import make_url


TABLES = {
    "sys_organization": 3,
    "sys_user": 3,
    "doc_category": 21,
    "doc_document": 31,
    "doc_attachment_binding": 19,
    "doc_permission": 7,
    "sys_audit_log": 291,
}


def execute_script(database_url: str, path: Path) -> None:
    engine = create_engine(
        database_url,
        connect_args={"client_flag": CLIENT.MULTI_STATEMENTS},
    )
    connection = engine.raw_connection()
    try:
        driver_connection = connection.driver_connection
        driver_connection.autocommit(True)
        with driver_connection.cursor() as cursor:
            cursor.execute(path.read_text(encoding="utf-8"))
            while cursor.nextset():
                pass
    finally:
        connection.close()
        engine.dispose()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--env-file", type=Path, required=True)
    parser.add_argument("--schema", type=Path, required=True)
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--backend", type=Path, required=True)
    parser.add_argument("--scratch-database", required=True)
    args = parser.parse_args()

    if not args.scratch_database.startswith("dms_full_demo_verify_"):
        raise RuntimeError("Scratch database name must start with dms_full_demo_verify_")

    source_url = make_url(dotenv_values(args.env_file).get("DMS_DATABASE_URL") or "")
    server_url = source_url.set(database=None)
    scratch_url = source_url.set(database=args.scratch_database)
    server_engine = create_engine(server_url)
    scratch_created = False
    try:
        with server_engine.connect().execution_options(isolation_level="AUTOCOMMIT") as connection:
            exists = connection.execute(
                text("SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME=:name"),
                {"name": args.scratch_database},
            ).first()
            if exists:
                raise RuntimeError("Scratch database already exists; refusing to overwrite it")
            connection.execute(text(f"CREATE DATABASE `{args.scratch_database}` CHARACTER SET utf8mb4"))
            scratch_created = True

        scratch_url_text = scratch_url.render_as_string(hide_password=False)
        execute_script(scratch_url_text, args.schema.resolve())
        previous_directory = Path.cwd()
        try:
            os.chdir(args.backend.resolve())
            os.environ["DMS_DATABASE_URL"] = scratch_url_text
            os.environ["DMS_JWT_SECRET"] = "verification-only-secret-not-for-deployment"
            command.upgrade(Config("migrations/alembic.ini"), "head")
        finally:
            os.chdir(previous_directory)
        execute_script(scratch_url_text, args.data.resolve())

        scratch_engine = create_engine(scratch_url)
        try:
            with scratch_engine.connect() as connection:
                for table_name, expected in TABLES.items():
                    actual = connection.execute(text(f"SELECT COUNT(*) FROM `{table_name}`")).scalar_one()
                    if actual != expected:
                        raise RuntimeError(f"{table_name}: expected {expected}, got {actual}")
                    print(f"{table_name}={actual}")
        finally:
            scratch_engine.dispose()
        print("scratch import verification passed")
    finally:
        if scratch_created:
            with server_engine.connect().execution_options(isolation_level="AUTOCOMMIT") as connection:
                connection.execute(text(f"DROP DATABASE `{args.scratch_database}`"))
            print("scratch database removed")
        server_engine.dispose()


if __name__ == "__main__":
    main()