document-management/backend/dms/services/sql_dump.py

# -*- coding: utf-8 -*-
"""自有格式数据库备份：数据 SQL 脚本导出/导入（替代 dexp/dimp，零外部工具）。

格式（database/dms-data.sql，UTF-8 文本）::

    -- DMS 数据库备份（自有格式 v1）
    -- 生成引擎: PYTHON / JAVA
    -- 生成时间: 2026-09-09 12:00:00
    -- 源模式: DMS
    -- 表数量: 14
    SET IDENTITY_INSERT "DMS"."DOC_CATEGORY" ON;
    INSERT INTO "DMS"."DOC_CATEGORY" ("ID", "NAME") VALUES (1, '规划');
    SET IDENTITY_INSERT "DMS"."DOC_CATEGORY" OFF;

设计要点：

- 表按外键依赖拓扑序（父表在前，字母序稳定），行按单列数值主键升序——
  同库同数据多次导出逐字节一致；自引用表父行 ID 必小于子行，天然满足；
- 业务表与迁移簿记表（alembic_version / flyway_schema_history）一并导出，
  版本戳随备份走，恢复后启动校验语义与 dexp 时代一致；
- IDENTITY 表以 SET IDENTITY_INSERT 包裹；包裹与否由导出时对目标库探测
  决定（SET ON 后立即 OFF，成功即含 IDENTITY 列），不猜测系统视图；
- 字面量：NULL / 数值原文（驱动字符串化保真）/ 二进制 0x hex / 日期时间
  'YYYY-MM-DD HH:MM:SS.ffffff' / 其余单引号转义；
- 恢复：禁用全部外键 → 清空全部业务表 → 事务内执行脚本 → 恢复（启用）
  全部外键——启用即完整性校验。外键维护是 DDL，在达梦隐式提交，置于
  业务事务之外；业务事务失败自动回滚后仍会重试恢复外键。

Java 版 com.dms.backup.SqlDataDumper 与本模块逐条对齐（同一 SQL、同一
格式、同一恢复语义），两侧产物可互相导入。
"""

from __future__ import annotations

import re
from datetime import date, datetime
from decimal import Decimal
from pathlib import Path
from typing import Any, Callable

from dms.common.errors import BackupFailedError

DUMP_FILE_NAME = "dms-data.sql"

_INSERT_TABLE_RE = re.compile(
    r'^INSERT\s+INTO\s+"([^"]+)"\."([^"]+)"', re.IGNORECASE
)


# ==================================================================== 导出


def export_dump(database_url: str, output_directory: Path, schema: str) -> Path:
    """全量导出 schema 业务数据到 output_directory/dms-data.sql，返回脚本路径。"""
    import sqlalchemy as sa

    output_directory.mkdir(parents=True, exist_ok=True)
    dump_file = output_directory / DUMP_FILE_NAME
    engine = sa.create_engine(database_url)
    try:
        with engine.connect() as conn:
            body = _build_dump_body(conn, schema, "PYTHON")
    finally:
        engine.dispose()
    dump_file.write_text(body, encoding="utf-8", newline="\n")
    return dump_file


def _build_dump_body(conn: Any, schema: str, engine_label: str) -> str:
    """导出主体（吃已连接的 conn，便于双端一致性测试）。"""
    tables = _list_tables(conn, schema)
    if not tables:
        raise BackupFailedError(f"数据库模式 {schema} 中没有任何表，无法导出")
    order = _topological_order(_foreign_key_edges(conn, schema), tables)
    blocks: list[str] = [
        "-- DMS 数据库备份（自有格式 v1）",
        f"-- 生成引擎: {engine_label}",
        f"-- 生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        f"-- 源模式: {schema}",
        f"-- 表数量: {len(order)}",
        "",
    ]
    for index, table in enumerate(order, start=1):
        blocks.append(
            f"-- 表 {index}/{len(order)}: {table}"
            f"（{_row_count(conn, schema, table)} 行）"
        )
        blocks.extend(_table_blocks(conn, schema, table))
    return "\n".join(blocks).rstrip("\n") + "\n"


def _list_tables(conn: Any, schema: str) -> list[str]:
    """schema 下全部表名字典序（ALL_TABLES 不含视图；DM Oracle 兼容字典）。"""
    from sqlalchemy import text

    rows = conn.execute(
        text("SELECT TABLE_NAME FROM ALL_TABLES WHERE OWNER = :o ORDER BY TABLE_NAME"),
        {"o": schema},
    )
    return [str(row[0]).upper() for row in rows]


def _foreign_key_edges(conn: Any, schema: str) -> list[tuple[str, str]]:
    """外键依赖边 (父表, 子表)；自引用与跨 schema 引用忽略。"""
    from sqlalchemy import text

    rows = conn.execute(
        text(
            "SELECT a.TABLE_NAME AS child_table, c.TABLE_NAME AS parent_table"
            " FROM ALL_CONSTRAINTS a"
            " JOIN ALL_CONSTRAINTS c ON a.R_CONSTRAINT_NAME = c.CONSTRAINT_NAME"
            " WHERE a.CONSTRAINT_TYPE = 'R' AND a.OWNER = :o"
        ),
        {"o": schema},
    )
    edges = []
    for child, parent in rows:
        child_table, parent_table = str(child).upper(), str(parent).upper()
        if child_table != parent_table:
            edges.append((parent_table, child_table))
    return edges


def _topological_order(
    edges: list[tuple[str, str]], tables: list[str]
) -> list[str]:
    """Kahn 拓扑排序（父表在前），同层字母序保证确定性；环内剩余表字母序追加。"""
    remaining = {
        table: {parent for parent, child in edges if child == table}
        for table in tables
    }
    ordered: list[str] = []
    while remaining:
        ready = sorted(t for t, parents in remaining.items() if not parents)
        if not ready:  # 依赖环：DISABLE 外键后导入顺序无关，字母序收尾
            ordered.extend(sorted(remaining))
            break
        for table in ready:
            ordered.append(table)
            remaining.pop(table)
        for table in remaining:
            remaining[table] -= set(ordered)
    return ordered


def _primary_key(conn: Any, schema: str, table: str) -> list[str]:
    """主键列名（按键序）；无主键返回空。"""
    from sqlalchemy import text

    rows = conn.execute(
        text(
            "SELECT cols.COLUMN_NAME"
            " FROM ALL_CONSTRAINTS con"
            " JOIN ALL_CONS_COLUMNS cols"
            "   ON con.CONSTRAINT_NAME = cols.CONSTRAINT_NAME"
            "  AND con.OWNER = cols.OWNER"
            " WHERE con.CONSTRAINT_TYPE = 'P' AND con.OWNER = :o"
            "   AND con.TABLE_NAME = :t"
            " ORDER BY cols.POSITION"
        ),
        {"o": schema, "t": table},
    )
    return [str(row[0]) for row in rows]


def _row_count(conn: Any, schema: str, table: str) -> int:
    from sqlalchemy import text

    return int(
        conn.execute(
            text(f'SELECT COUNT(*) FROM "{schema}"."{table}"')
        ).scalar_one()
    )


def _table_blocks(conn: Any, schema: str, table: str) -> list[str]:
    """单表导出块：IDENTITY 包裹（探测）+ 逐行 INSERT（带列清单）。"""
    from sqlalchemy import text

    target = f'"{schema}"."{table}"'
    identity = _table_has_identity(conn, target)
    key_columns = _primary_key(conn, schema, table)
    if len(key_columns) == 1:
        select = f"SELECT * FROM {target} ORDER BY \"{key_columns[0]}\""
    else:
        select = f"SELECT * FROM {target}"
    result = conn.execute(text(select))
    columns = [str(name).upper() for name in result.keys()]
    column_list = ", ".join(f'"{name}"' for name in columns)
    rows = [
        "INSERT INTO {target} ({columns}) VALUES ({values});".format(
            target=target,
            columns=column_list,
            values=", ".join(_render_value(value) for value in row),
        )
        for row in result
    ]
    if not rows:
        return []
    if not identity:
        return rows
    return [f"SET IDENTITY_INSERT {target} ON;", *rows, f"SET IDENTITY_INSERT {target} OFF;"]


def _table_has_identity(conn: Any, target: str) -> bool:
    """探测表是否含 IDENTITY 列：SET ON 后立即 OFF 成功即为真。"""
    from sqlalchemy import text

    try:
        conn.execute(text(f"SET IDENTITY_INSERT {target} ON"))
        conn.execute(text(f"SET IDENTITY_INSERT {target} OFF"))
        return True
    except Exception:  # noqa: BLE001 - 非 IDENTITY 表达梦直接报错，按无处理
        return False


def _render_value(value: Any) -> str:
    """SQL 字面量渲染（与 Java SqlDataDumper.renderValue 语义对齐）。"""
    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, (bytes, bytearray, memoryview)):
        return "0x" + bytes(value).hex()
    if isinstance(value, datetime):
        return "'" + value.strftime("%Y-%m-%d %H:%M:%S.%f") + "'"
    if isinstance(value, date):
        return "'" + value.strftime("%Y-%m-%d") + "'"
    return "'" + str(value).replace("'", "''") + "'"


# ==================================================================== 恢复


def restore_dump(
    database_url: str,
    dump_file: Path,
    schema: str,
    ensure_schema_ready: Callable[[], None] | None = None,
) -> None:
    """导入自有格式数据脚本：禁用外键 → 清表 → 事务导入 → 恢复外键。

    ensure_schema_ready：目标库为空表时的建库回调（Python CLI 传 Alembic
    upgrade head；为 None 且空库则报错引导）。
    """
    import sqlalchemy as sa

    statements = split_statements(dump_file.read_text(encoding="utf-8"))
    if not statements:
        raise BackupFailedError(f"数据库脚本为空或不含任何语句：{dump_file}")
    engine = sa.create_engine(database_url)
    try:
        with engine.connect() as conn:
            _restore_via(conn, engine, schema, statements, ensure_schema_ready)
    finally:
        engine.dispose()


def _restore_via(
    conn: Any,
    engine: Any,
    schema: str,
    statements: list[str],
    ensure_schema_ready: Callable[[], None] | None,
) -> None:
    """恢复主体（吃已连接的 conn 与 engine，便于双端一致性测试）。"""
    from sqlalchemy import text

    tables = _list_tables(conn, schema)
    if not tables:
        if ensure_schema_ready is None:
            raise BackupFailedError(
                "目标库为空（无任何表）。请先完成数据库初始化"
                "（upgrade-db 或参照部署文档建库）后再执行恢复。"
            )
        ensure_schema_ready()
        tables = _list_tables(conn, schema)

    dump_tables = sorted(
        {
            f"{match.group(1)}.{match.group(2)}".upper()
            for statement in statements
            if (match := _INSERT_TABLE_RE.match(statement))
        }
    )
    known = set(tables)
    missing = [name for name in dump_tables if name.split(".", 1)[1] not in known]
    if missing:
        raise BackupFailedError(
            f"目标库缺少备份中的表：{', '.join(missing)}。"
            "请将程序升级到与备份匹配的版本后重试。"
        )

    # 外键禁用（DDL 隐式提交；个别失败不阻断，导入顺序已按拓扑序兜底）
    disabled = _set_foreign_keys(conn, schema, False)
    try:
        with engine.begin() as transaction:
            for table in reversed(tables):
                transaction.execute(text(f'DELETE FROM "{schema}"."{table}"'))
            for statement in statements:
                transaction.execute(text(statement))
    finally:
        # 成功：启用即完整性校验；失败：回滚后恢复原状，启用应通过
        failed = _set_foreign_keys(conn, schema, True, only=disabled)
    if failed:
        detail = "；".join(f'"{schema}"."{table}" 约束 "{name}"' for table, name in failed)
        raise BackupFailedError(f"外键恢复失败（数据完整性校验未通过）：{detail}")


def _set_foreign_keys(
    conn: Any,
    schema: str,
    enable: bool,
    only: list[tuple[str, str]] | None = None,
) -> list[tuple[str, str]]:
    """批量启用/禁用外键，返回成功操作的约束清单（enable 失败返回失败项语义不同）。

    disable 返回成功禁用的清单；enable（only 指定目标）返回失败清单。
    """
    from sqlalchemy import text

    if only is None:
        rows = conn.execute(
            text(
                "SELECT TABLE_NAME, CONSTRAINT_NAME FROM ALL_CONSTRAINTS"
                " WHERE CONSTRAINT_TYPE = 'R' AND OWNER = :o"
                " ORDER BY TABLE_NAME, CONSTRAINT_NAME"
            ),
            {"o": schema},
        )
        constraints = [(str(row[0]).upper(), str(row[1])) for row in rows]
    else:
        constraints = only
    action = "ENABLE" if enable else "DISABLE"
    succeeded: list[tuple[str, str]] = []
    failed: list[tuple[str, str]] = []
    for table, name in constraints:
        try:
            conn.execute(
                text(f'ALTER TABLE "{schema}"."{table}" {action} CONSTRAINT "{name}"')
            )
            succeeded.append((table, name))
        except Exception:  # noqa: BLE001 - 单条失败继续，禁用失败靠拓扑序兜底
            failed.append((table, name))
    return failed if enable else succeeded


# ==================================================================== 脚本切分


def split_statements(sql: str) -> list[str]:
    """按分号切分；跳过行/块注释；保留单引号字符串与双引号标识符（含转义）。"""
    out: list[str] = []
    buf: list[str] = []
    i, n = 0, len(sql)
    while i < n:
        char = sql[i]
        nxt = sql[i + 1] if i + 1 < n else ""
        if char == "-" and nxt == "-":
            while i < n and sql[i] != "\n":
                i += 1
            continue
        if char == "/" and nxt == "*":
            i += 2
            while i + 1 < n and not (sql[i] == "*" and sql[i + 1] == "/"):
                i += 1
            i += 2
            continue
        if char == "'":
            buf.append(char)
            i += 1
            while i < n:
                buf.append(sql[i])
                if sql[i] == "'" and i + 1 < n and sql[i + 1] == "'":
                    buf.append("'")
                    i += 2
                    continue
                if sql[i] == "'":
                    break
                i += 1
            i += 1
            continue
        if char == '"':
            buf.append(char)
            i += 1
            while i < n:
                buf.append(sql[i])
                if sql[i] == '"' and i + 1 < n and sql[i + 1] == '"':
                    buf.append('"')
                    i += 2
                    continue
                if sql[i] == '"':
                    break
                i += 1
            i += 1  # 跳过结束引号（与单引号分支对齐，缺失会重入字符串状态）
            continue
        if char == ";":
            statement = "".join(buf).strip()
            if statement:
                out.append(statement)
            buf = []
            i += 1
            continue
        buf.append(char)
        i += 1
    tail = "".join(buf).strip()
    if tail:
        out.append(tail)
    return out