document-management/backend/tests/dms/test_sql_dump.py

# -*- coding: utf-8 -*-
"""sql_dump 自有格式数据库备份单元测试（假连接，零真实数据库依赖）。

覆盖：语句切分、SQL 字面量渲染、外键拓扑排序、导出块结构（IDENTITY
包裹/主键行序/空表）、恢复执行序列（禁用外键→逆序清表→导入→恢复外键）、
空库与缺表守卫、外键启用失败与事务失败回补。
"""

from __future__ import annotations

import contextlib
import re
from datetime import date, datetime
from decimal import Decimal

import pytest

from dms.common.errors import BackupFailedError
from dms.services.sql_dump import (
    _build_dump_body,
    _foreign_key_edges,
    _render_value,
    _restore_via,
    _topological_order,
    split_statements,
)


# ------------------------------------------------------------- 假连接/假引擎


class FakeResult:
    """最小结果集：keys() + 行迭代 + scalar_one。"""

    def __init__(self, columns, rows):
        self._columns = columns
        self._rows = rows

    def keys(self):
        return self._columns

    def __iter__(self):
        return iter(self._rows)

    def scalar_one(self):
        return self._rows[0][0]


class FakeConnection:
    """按 SQL 前缀路由的假达梦连接，记录每一条下发的语句。"""

    def __init__(
        self,
        *,
        tables,
        rows_by_table=None,
        primary_keys=None,
        fk_edges=None,
        fk_constraints=None,
        identity_tables=(),
        fail_enable_constraints=(),
    ):
        self.tables = list(tables)
        self.rows_by_table = rows_by_table or {}
        self.primary_keys = primary_keys or {}
        self.fk_edges = list(fk_edges or [])
        self.fk_constraints = list(fk_constraints or [])
        self.identity_tables = set(identity_tables)
        self.fail_enable = set(fail_enable_constraints)
        self.executed = []

    def execute(self, sql, params=None):
        statement = getattr(sql, "text", str(sql))
        self.executed.append(statement)
        if statement.startswith("SELECT TABLE_NAME FROM ALL_TABLES"):
            return FakeResult(["TABLE_NAME"], [(t,) for t in self.tables])
        if statement.startswith("SELECT a.TABLE_NAME"):
            return FakeResult([], [(c, p) for c, p in self.fk_edges])
        if statement.startswith("SELECT TABLE_NAME, CONSTRAINT_NAME"):
            return FakeResult([], [(t, n) for t, n in self.fk_constraints])
        if statement.startswith("SELECT cols.COLUMN_NAME"):
            columns = self.primary_keys.get(params["t"], [])
            return FakeResult(["COLUMN_NAME"], [(c,) for c in columns])
        if match := re.match(r'SELECT COUNT\(\*\) FROM "[^"]+"\."([^"]+)"', statement):
            rows = self.rows_by_table.get(match.group(1), ([], []))[1]
            return FakeResult([], [(len(rows),)])
        if match := re.match(r'SELECT \* FROM "[^"]+"\."([^"]+)"', statement):
            columns, rows = self.rows_by_table.get(match.group(1), ([], []))
            return FakeResult(columns, rows)
        if match := re.match(
            r'SET IDENTITY_INSERT "[^"]+"\."([^"]+)" (ON|OFF)$', statement
        ):
            if match.group(1) not in self.identity_tables:
                raise RuntimeError("非 IDENTITY 表")
            return FakeResult([], [])
        if statement.startswith("DELETE FROM") or statement.startswith("INSERT INTO"):
            return FakeResult([], [])
        if match := re.match(
            r'ALTER TABLE "[^"]+"\."([^"]+)" (ENABLE|DISABLE) CONSTRAINT "([^"]+)"$',
            statement,
        ):
            _table, action, name = match.groups()
            if action == "ENABLE" and name in self.fail_enable:
                raise RuntimeError(f"约束 {name} 启用失败")
            return FakeResult([], [])
        raise AssertionError(f"未预期的 SQL：{statement}")


class FakeTransaction:
    def __init__(self, log, fail_on=None):
        self.log = log
        self.fail_on = fail_on

    def execute(self, stmt):
        text = stmt.text
        if self.fail_on and self.fail_on in text:
            raise RuntimeError(f"事务内语句失败：{text}")
        self.log.append(text)


class FakeEngine:
    def __init__(self, fail_on=None):
        self.executed = []
        self.fail_on = fail_on

    def begin(self):
        return contextlib.nullcontext(FakeTransaction(self.executed, self.fail_on))


# ------------------------------------------------------------- 标准测试数据

TABLES = ["DOC_CATEGORY", "SYS_ORGANIZATION", "SYS_USER"]
ROWS = {
    "DOC_CATEGORY": (["ID", "NAME"], [(1, "规划"), (2, "其它")]),
    "SYS_ORGANIZATION": (["ID", "NAME"], [(1, "总部")]),
    "SYS_USER": (["ID", "NAME", "ORG_ID"], [(1, "管理员", 1)]),
}
PKS = {"DOC_CATEGORY": ["ID"], "SYS_ORGANIZATION": ["ID"], "SYS_USER": ["ID"]}
FK_EDGES = [("SYS_USER", "SYS_ORGANIZATION"), ("SYS_USER", "DOC_CATEGORY")]
FK_CONSTRAINTS = [("SYS_USER", "FK_USER_ORG"), ("SYS_USER", "FK_USER_CATEGORY")]
IDENTITY = {"DOC_CATEGORY", "SYS_ORGANIZATION", "SYS_USER"}

RESTORE_STATEMENTS = [
    'SET IDENTITY_INSERT "DMS"."DOC_CATEGORY" ON',
    "INSERT INTO \"DMS\".\"DOC_CATEGORY\" (\"ID\", \"NAME\") VALUES (1, '规划')",
    'SET IDENTITY_INSERT "DMS"."DOC_CATEGORY" OFF',
    "INSERT INTO \"DMS\".\"SYS_USER\" (\"ID\", \"NAME\", \"ORG_ID\") VALUES (1, '管理员', 1)",
]


def _make_conn(**overrides):
    kwargs = dict(
        tables=TABLES,
        rows_by_table=ROWS,
        primary_keys=PKS,
        fk_edges=FK_EDGES,
        fk_constraints=FK_CONSTRAINTS,
        identity_tables=IDENTITY,
    )
    kwargs.update(overrides)
    return FakeConnection(**kwargs)


# ------------------------------------------------------------- 脚本切分


def test_split_statements_basic_and_empty():
    assert split_statements("A;\nB;") == ["A", "B"]
    assert split_statements("") == []
    assert split_statements("-- 仅注释; 无语句") == []
    assert split_statements("A;;;  ;") == ["A"]
    assert split_statements("无分号结尾") == ["无分号结尾"]


def test_split_statements_skips_comments():
    sql = "-- 头注释; 内含分号\nSET X ON; /* 块; 注释 */ DELETE FROM t; -- 尾注"
    assert split_statements(sql) == ["SET X ON", "DELETE FROM t"]


def test_split_statements_keeps_quoted_semicolons_and_comments():
    sql = "INSERT INTO t VALUES ('a;b--c'); INSERT INTO \"T;X\" VALUES ('it''s;');"
    assert split_statements(sql) == [
        "INSERT INTO t VALUES ('a;b--c')",
        'INSERT INTO "T;X" VALUES (\'it\'\'s;\')',
    ]


# ------------------------------------------------------------- 字面量渲染


def test_render_value_literals():
    assert _render_value(None) == "NULL"
    assert _render_value(True) == "1"
    assert _render_value(False) == "0"
    assert _render_value(42) == "42"
    assert _render_value(Decimal("1.25")) == "1.25"
    assert _render_value(1.5) == "1.5"
    assert _render_value(b"\x00\xff") == "0x00ff"
    assert _render_value(datetime(2026, 9, 9, 7, 8, 9, 123456)) == (
        "'2026-09-09 07:08:09.123456'"
    )
    assert _render_value(date(2026, 9, 9)) == "'2026-09-09'"
    assert _render_value("O'Neill") == "'O''Neill'"
    assert _render_value("规划") == "'规划'"


# ------------------------------------------------------------- 拓扑排序


def test_topological_order_parents_first_with_alpha_ties():
    # 边为 (父表, 子表)：SYS_USER 引用另两张表，故排最后
    edges = [("SYS_ORGANIZATION", "SYS_USER"), ("DOC_CATEGORY", "SYS_USER")]
    tables = ["SYS_USER", "SYS_ORGANIZATION", "DOC_CATEGORY"]
    assert _topological_order(edges, tables) == [
        "DOC_CATEGORY",
        "SYS_ORGANIZATION",
        "SYS_USER",
    ]


def test_topological_order_cycle_falls_back_to_alpha():
    edges = [("A", "B"), ("B", "A")]
    assert _topological_order(edges, ["B", "A", "C"]) == ["C", "A", "B"]


def test_foreign_key_edges_parent_child_and_skip_self():
    conn = _make_conn(fk_edges=[("SYS_USER", "SYS_ORGANIZATION"), ("T", "T")])
    assert _foreign_key_edges(conn, "DMS") == [("SYS_ORGANIZATION", "SYS_USER")]


# ------------------------------------------------------------- 导出


def test_build_dump_body_structure_and_identity_wrap():
    body = _build_dump_body(_make_conn(), "DMS", "PYTHON")
    assert body.startswith("-- DMS 数据库备份（自有格式 v1）\n")
    assert "-- 生成引擎: PYTHON\n" in body
    assert "-- 源模式: DMS\n" in body
    assert "-- 表数量: 3\n" in body
    assert body.endswith("\n")

    lines = body.splitlines()
    inserts = [line for line in lines if line.startswith("INSERT INTO")]
    # 拓扑序（按各表首行出现位置）：子表 SYS_USER 在两个父表之后
    first_index = {
        table: next(
            index
            for index, line in enumerate(inserts)
            if line.startswith(f'INSERT INTO "DMS"."{table}"')
        )
        for table in ("DOC_CATEGORY", "SYS_ORGANIZATION", "SYS_USER")
    }
    assert first_index["DOC_CATEGORY"] < first_index["SYS_USER"]
    assert first_index["SYS_ORGANIZATION"] < first_index["SYS_USER"]
    # 带列清单、字符串中文原文
    assert (
        "INSERT INTO \"DMS\".\"DOC_CATEGORY\" (\"ID\", \"NAME\") VALUES (1, '规划');"
        in inserts
    )
    # IDENTITY 表以 SET ON/OFF 包裹且包裹行夹住 INSERT
    on_line = 'SET IDENTITY_INSERT "DMS"."DOC_CATEGORY" ON;'
    off_line = 'SET IDENTITY_INSERT "DMS"."DOC_CATEGORY" OFF;'
    assert lines.index(on_line) < lines.index(inserts[0]) < lines.index(off_line)
    # 行数注释
    assert "-- 表 1/3: DOC_CATEGORY（2 行）" in lines
    assert "-- 表 3/3: SYS_USER（1 行）" in lines


def test_build_dump_body_orders_rows_by_single_pk():
    conn = _make_conn()
    _build_dump_body(conn, "DMS", "PYTHON")
    assert 'SELECT * FROM "DMS"."DOC_CATEGORY" ORDER BY "ID"' in conn.executed
    assert 'SELECT * FROM "DMS"."SYS_USER" ORDER BY "ID"' in conn.executed


def test_build_dump_body_composite_pk_skips_order_by():
    conn = _make_conn(primary_keys={"DOC_CATEGORY": ["ID"], "SYS_ORGANIZATION": ["ID"], "SYS_USER": ["ID", "NAME"]})
    _build_dump_body(conn, "DMS", "PYTHON")
    assert 'SELECT * FROM "DMS"."SYS_USER"' in conn.executed


def test_build_dump_body_skips_identity_wrap_when_probe_fails():
    conn = _make_conn(identity_tables=set())
    body = _build_dump_body(conn, "DMS", "PYTHON")
    assert "SET IDENTITY_INSERT" not in body
    # 探测语句确实已下发：非 IDENTITY 表 ON 即报错，每表仅 1 条
    probes = [s for s in conn.executed if s.startswith("SET IDENTITY_INSERT")]
    assert len(probes) == 3


def test_build_dump_body_empty_table_only_comment():
    rows = {
        "DOC_CATEGORY": (["ID", "NAME"], []),
        "SYS_ORGANIZATION": ROWS["SYS_ORGANIZATION"],
        "SYS_USER": ROWS["SYS_USER"],
    }
    body = _build_dump_body(_make_conn(rows_by_table=rows), "DMS", "JAVA")
    assert "-- 表 1/3: DOC_CATEGORY（0 行）" in body.splitlines()
    assert 'INSERT INTO "DMS"."DOC_CATEGORY"' not in body
    # 空表整块省略，IDENTITY 包裹也不出现
    assert 'SET IDENTITY_INSERT "DMS"."DOC_CATEGORY"' not in body


def test_build_dump_body_empty_schema_rejected():
    with pytest.raises(BackupFailedError, match="没有任何表"):
        _build_dump_body(_make_conn(tables=[]), "DMS", "PYTHON")


# ------------------------------------------------------------- 恢复


def test_restore_via_disable_clear_import_enable_sequence():
    conn = _make_conn()
    engine = FakeEngine()

    _restore_via(conn, engine, "DMS", RESTORE_STATEMENTS, None)

    disables = [s for s in conn.executed if " DISABLE CONSTRAINT" in s]
    enables = [s for s in conn.executed if " ENABLE CONSTRAINT" in s]
    assert len(disables) == len(enables) == 2
    # 先全部禁用、业务执行后（finally）再启用
    assert conn.executed.index(disables[-1]) < conn.executed.index(enables[0])
    # 事务内：拓扑逆序清空全部业务表 → 按脚本顺序导入
    expected_clears = [
        'DELETE FROM "DMS"."SYS_USER"',
        'DELETE FROM "DMS"."SYS_ORGANIZATION"',
        'DELETE FROM "DMS"."DOC_CATEGORY"',
    ]
    assert engine.executed == expected_clears + RESTORE_STATEMENTS


def test_restore_via_empty_database_without_callback_rejected():
    conn = _make_conn(tables=[])
    with pytest.raises(BackupFailedError, match="目标库为空"):
        _restore_via(conn, FakeEngine(), "DMS", RESTORE_STATEMENTS, None)
    assert not [s for s in conn.executed if "CONSTRAINT" in s]


def test_restore_via_empty_database_runs_schema_callback():
    conn = _make_conn(tables=[])
    engine = FakeEngine()
    calls = []

    def ensure_schema_ready():
        calls.append("upgrade")
        conn.tables = list(TABLES)

    _restore_via(conn, engine, "DMS", RESTORE_STATEMENTS, ensure_schema_ready)

    assert calls == ["upgrade"]
    assert engine.executed[-len(RESTORE_STATEMENTS):] == RESTORE_STATEMENTS


def test_restore_via_missing_table_rejected_before_any_change():
    conn = _make_conn()
    engine = FakeEngine()
    statements = ['INSERT INTO "DMS"."NEW_TABLE" ("ID") VALUES (1)']

    with pytest.raises(BackupFailedError, match="目标库缺少备份中的表"):
        _restore_via(conn, engine, "DMS", statements, None)

    assert not [s for s in conn.executed if "CONSTRAINT" in s]
    assert engine.executed == []


def test_restore_via_reports_failed_fk_enable():
    conn = _make_conn(fail_enable_constraints={"FK_USER_ORG"})
    engine = FakeEngine()
    statements = ["INSERT INTO \"DMS\".\"DOC_CATEGORY\" (\"ID\") VALUES (1)"]

    with pytest.raises(BackupFailedError, match="外键恢复失败"):
        _restore_via(conn, engine, "DMS", statements, None)

    enables = [s for s in conn.executed if " ENABLE CONSTRAINT" in s]
    assert len(enables) == 2
    assert any('CONSTRAINT "FK_USER_CATEGORY"' in s for s in enables)


def test_restore_via_dml_failure_still_reenables_and_raises():
    conn = _make_conn()
    engine = FakeEngine(fail_on="BOOM")
    statements = [
        'INSERT INTO "DMS"."DOC_CATEGORY" ("ID") VALUES (1)',
        'INSERT INTO "DMS"."DOC_CATEGORY" ("ID") VALUES (2) /* BOOM */',
    ]

    with pytest.raises(RuntimeError, match="事务内语句失败"):
        _restore_via(conn, engine, "DMS", statements, None)

    # 失败前的语句已执行，失败语句未入库
    assert statements[0] in engine.executed
    assert statements[1] not in engine.executed
    # finally 仍完成外键恢复
    assert len([s for s in conn.executed if " DISABLE CONSTRAINT" in s]) == 2
    assert len([s for s in conn.executed if " ENABLE CONSTRAINT" in s]) == 2