document-management/tools/export_full_demo_data.py

"""Export the local DMS demo database rows and referenced storage files.

The connection URL is read from an existing dotenv file and is never written
to the export. The resulting SQL is data-only and targets an already migrated
empty DMS database.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from datetime import date, datetime, time
from decimal import Decimal
from pathlib import Path
from typing import Any

from dotenv import dotenv_values
from sqlalchemy import MetaData, Table, create_engine, select


TABLES = (
    "sys_organization",
    "sys_user",
    "doc_category",
    "doc_document",
    "doc_attachment_binding",
    "doc_permission",
    "sys_audit_log",
)


def sql_literal(value: Any) -> str:
    if value is None:
        return "NULL"
    if isinstance(value, bool):
        return "1" if value else "0"
    if isinstance(value, (int, Decimal)):
        return str(value)
    if isinstance(value, float):
        return repr(value)
    if isinstance(value, (datetime, date, time)):
        value = value.isoformat(sep=" ") if isinstance(value, datetime) else value.isoformat()
    if isinstance(value, (dict, list)):
        value = json.dumps(value, ensure_ascii=False, separators=(",", ":"))
    if isinstance(value, bytes):
        return "X'" + value.hex() + "'"
    escaped = str(value).replace("\\", "\\\\").replace("'", "''")
    return "'" + escaped + "'"


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


def export(env_file: Path, source_storage: Path, output: Path) -> None:
    config = dotenv_values(env_file)
    database_url = (config.get("DMS_DATABASE_URL") or "").strip()
    if not database_url:
        raise RuntimeError("DMS_DATABASE_URL is missing from the source dotenv file")

    output.mkdir(parents=True, exist_ok=True)
    storage_output = output / "storage"
    storage_output.mkdir(exist_ok=True)

    engine = create_engine(database_url, pool_pre_ping=True)
    metadata = MetaData()
    reflected = {
        name: Table(name, metadata, autoload_with=engine)
        for name in TABLES
    }

    rows_by_table: dict[str, list[dict[str, Any]]] = {}
    with engine.connect() as connection:
        for name in TABLES:
            rows_by_table[name] = [dict(row) for row in connection.execute(select(reflected[name])).mappings()]
    engine.dispose()

    document_paths: dict[str, dict[str, Any]] = {}
    for row in rows_by_table["doc_document"]:
        relative = str(row.get("file_relative_path") or "").strip().replace("\\", "/")
        if not relative:
            continue
        if relative.startswith("/") or ".." in Path(relative).parts:
            raise RuntimeError(f"Unsafe storage path in database: {relative}")
        source = source_storage / Path(relative)
        if not source.is_file():
            raise RuntimeError(f"Referenced storage file is missing: {relative}")
        expected_size = row.get("file_size")
        if expected_size is not None and source.stat().st_size != int(expected_size):
            raise RuntimeError(f"Referenced storage file has unexpected size: {relative}")
        actual_hash = sha256(source)
        expected_hash = str(row.get("file_hash") or "").lower()
        if expected_hash and actual_hash != expected_hash:
            raise RuntimeError(f"Referenced storage file has unexpected SHA-256: {relative}")
        destination = storage_output / Path(relative)
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(source, destination)
        document_paths[relative] = {
            "size": source.stat().st_size,
            "sha256": actual_hash,
        }

    sql_path = output / "full_demo_data.sql"
    with sql_path.open("w", encoding="utf-8", newline="\n") as sql:
        sql.write("-- DMS full demo data export (data only; no schema or secrets)\n")
        sql.write("SET NAMES utf8mb4;\nSET FOREIGN_KEY_CHECKS=0;\nSTART TRANSACTION;\n")
        for name in reversed(TABLES):
            sql.write(f"DELETE FROM `{name}`;\n")
        for name in TABLES:
            table = reflected[name]
            # MySQL generated columns (for example active_marker) must not be
            # supplied explicitly in INSERT statements.
            columns = [
                column.name
                for column in table.columns
                if column.computed is None
            ]
            quoted_columns = ", ".join(f"`{column}`" for column in columns)
            for row in rows_by_table[name]:
                values = ", ".join(sql_literal(row.get(column)) for column in columns)
                sql.write(f"INSERT INTO `{name}` ({quoted_columns}) VALUES ({values});\n")
        sql.write("COMMIT;\nSET FOREIGN_KEY_CHECKS=1;\n")

    manifest = {
        "format": "DMS full demo data v1",
        "database_tables": {name: len(rows_by_table[name]) for name in TABLES},
        "storage_file_count": len(document_paths),
        "storage_total_bytes": sum(item["size"] for item in document_paths.values()),
        "storage_files": document_paths,
        "sql_sha256": sha256(sql_path),
    }
    (output / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(manifest, ensure_ascii=False, indent=2))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--env-file", type=Path, required=True)
    parser.add_argument("--storage", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    arguments = parser.parse_args()
    export(arguments.env_file.resolve(), arguments.storage.resolve(), arguments.output.resolve())


if __name__ == "__main__":
    main()