264 lines
13 KiB
Python
264 lines
13 KiB
Python
"""
|
||
跨库数据泵（scripts/migrate_db.py）单元测试
|
||
==========================================
|
||
|
||
用 sqlite 文件同时充当「源库」和「目标库」，验证泵逻辑全链路：
|
||
读→分块→executemany→显式 id 保真→自增延续→Text/日期/BigInteger 保真→
|
||
软删/空时间原样→only-tables→dry-run→逐表 count 校验→建库→序列重置 SQL。
|
||
|
||
不依赖真实 MySQL/PostgreSQL（不挂 integration marker，pytest 默认即跑）。
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
from datetime import date
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
import pytest
|
||
from sqlalchemy import create_engine, select, text
|
||
|
||
# scripts/ 下无 __init__.py，把 scripts 目录加进 sys.path 以便直接 import
|
||
_SCRIPTS_DIR = Path(__file__).resolve().parents[3] / "scripts"
|
||
if str(_SCRIPTS_DIR) not in sys.path:
|
||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||
|
||
import migrate_db
|
||
|
||
|
||
def make_args(src_db: str, dst_db: str, **overrides) -> SimpleNamespace:
|
||
"""构造 sqlite→sqlite 迁移的 CLI 参数（等价 migrate_db.build_parser().parse_args）。"""
|
||
base = dict(
|
||
src_driver="sqlite", src_host="", src_port=None, src_db=src_db, src_user="", src_pass="",
|
||
dst_driver="sqlite", dst_host="", dst_port=None, dst_db=dst_db, dst_user="", dst_pass="",
|
||
drop_target_first=False, dry_run=False, chunk_size=1000, only_tables="",
|
||
skip_existing=False, skip_verify=False, single_transaction=False,
|
||
)
|
||
base.update(overrides)
|
||
return SimpleNamespace(**base)
|
||
|
||
|
||
@pytest.fixture(scope="module")
|
||
def loaded_metadata():
|
||
"""一次性加载全部模型填充 Base.metadata（与 migrate_db.run 同源）。"""
|
||
from core.app import create_app
|
||
from core.database import Base
|
||
create_app()
|
||
return Base
|
||
|
||
|
||
@pytest.fixture
|
||
def two_sqlite(tmp_path, loaded_metadata):
|
||
"""一对 sqlite 文件库：源库 + 目标库，均已按模型建好空结构。"""
|
||
src_path = tmp_path / "src.db"
|
||
dst_path = tmp_path / "dst.db"
|
||
src_engine = create_engine(f"sqlite:///{src_path}")
|
||
dst_engine = create_engine(f"sqlite:///{dst_path}")
|
||
loaded_metadata.metadata.create_all(src_engine)
|
||
loaded_metadata.metadata.create_all(dst_engine)
|
||
yield src_engine, dst_engine, src_path, dst_path
|
||
src_engine.dispose()
|
||
dst_engine.dispose()
|
||
|
||
|
||
def _table(loaded_metadata, name: str):
|
||
return loaded_metadata.metadata.tables["fastapi_" + name]
|
||
|
||
|
||
def test_make_url():
|
||
"""URL 形态与 core.config.database.build_database_url 逐字等价（特殊字符口令编码）。"""
|
||
assert migrate_db.make_url("mysql", "127.0.0.1", 3306, "db", "root", "p@ss") \
|
||
== "mysql+pymysql://root:p%40ss@127.0.0.1:3306/db?charset=utf8mb4"
|
||
assert migrate_db.make_url("postgresql", "127.0.0.1", 5432, "db", "root", "p@ss") \
|
||
== "postgresql+psycopg://root:p%40ss@127.0.0.1:5432/db"
|
||
assert migrate_db.make_url("mssql", "127.0.0.1", 1433, "db", "sa", "p@ss") \
|
||
== "mssql+pymssql://sa:p%40ss@127.0.0.1:1433/db?charset=utf8"
|
||
assert migrate_db.make_url("oracle", "127.0.0.1", 1521, "orcl", "hr", "p@ss") \
|
||
== "oracle+oracledb://hr:p%40ss@127.0.0.1:1521/?service_name=orcl"
|
||
assert migrate_db.make_url("sqlite", "x", None, "dev.db", "", "") == "sqlite:///./dev.db"
|
||
|
||
|
||
def test_pump_preserves_data_ids_and_autoincrement(two_sqlite, loaded_metadata):
|
||
"""泵后 id/父子引用/Text/BigInteger/软删/空时间全保真，且新插入自增延续。"""
|
||
src_engine, dst_engine, src_path, dst_path = two_sqlite
|
||
|
||
with src_engine.begin() as conn:
|
||
conn.execute(_table(loaded_metadata, "dept").insert(), [
|
||
{"id": 1, "name": "总部", "code": "HQ", "parent_id": 0, "sort": 1},
|
||
{"id": 2, "name": "研发部", "code": "RD", "parent_id": 1, "sort": 2},
|
||
])
|
||
conn.execute(_table(loaded_metadata, "user").insert(), [
|
||
{"id": 1, "username": "admin", "password": "$2b$12$abcdefghijklmnop", "salt": "salt1",
|
||
"realname": "管理员", "nickname": "管理员", "avatar": "", "mobile": "13800000000",
|
||
"email": "admin@test.com", "province_code": "110000", "city_code": "110100",
|
||
"district_code": "110101", "street_code": "", "address": "北京市", "is_delete": 0},
|
||
{"id": 2, "username": "softdel", "password": None, "salt": None,
|
||
"realname": "软删", "nickname": "软删", "avatar": "",
|
||
"mobile": "13900000000", "email": "s@test.com", "province_code": "", "city_code": "",
|
||
"district_code": "", "street_code": "", "address": "", "is_delete": 1},
|
||
])
|
||
# base_model 的 create_time 带 Python 默认值，Core insert 显式 None 也会被默认覆盖；
|
||
# 用 UPDATE 强制置 NULL，验证泵能原样搬空时间列
|
||
conn.execute(text("UPDATE fastapi_user SET create_time = NULL WHERE id = 2"))
|
||
conn.execute(_table(loaded_metadata, "city").insert(), [
|
||
{"id": 1, "name": "北京市", "short_name": "北京", "city_code": "010",
|
||
"area_code": "110000", "zip_code": "100000",
|
||
"parent_code": None, "parent_id": 0, "level": 1},
|
||
{"id": 2, "name": "东城区", "short_name": "东城", "city_code": "010",
|
||
"area_code": "110101", "zip_code": "100010",
|
||
"parent_code": "1101", "parent_id": 1, "level": 2},
|
||
])
|
||
conn.execute(_table(loaded_metadata, "article").insert(), [
|
||
{"id": 1, "title": "中文标题",
|
||
"content": "第一行\n第二行 with <b>html</b> 中文", "images": "a.jpg,b.jpg"},
|
||
])
|
||
conn.execute(_table(loaded_metadata, "job_log").insert(), [
|
||
{"id": 1, "job_name": "测试任务", "consume_time": 9223372036854775807}, # BIGINT 上限
|
||
])
|
||
|
||
assert migrate_db.run(make_args(str(src_path), str(dst_path))) == 0
|
||
|
||
# 逐表 count 相等（含空表）
|
||
for t in migrate_db.ordered_tables():
|
||
with src_engine.connect() as src, dst_engine.connect() as dst:
|
||
s = src.execute(text(f"SELECT COUNT(*) FROM {t.name}")).scalar()
|
||
d = dst.execute(text(f"SELECT COUNT(*) FROM {t.name}")).scalar()
|
||
assert s == d, t.name
|
||
|
||
with dst_engine.connect() as conn:
|
||
# 目标 id 与源逐行一致 + 子树 parent_id 有效
|
||
dept = [r for r in conn.execute(select(_table(loaded_metadata, "dept")).order_by(
|
||
_table(loaded_metadata, "dept").c.id))]
|
||
assert [(r.id, r.parent_id) for r in dept] == [(1, 0), (2, 1)]
|
||
|
||
# 软删行 + 空 create_time 原样；密码/salt 原样
|
||
users = [r for r in conn.execute(select(_table(loaded_metadata, "user")).order_by(
|
||
_table(loaded_metadata, "user").c.id))]
|
||
assert users[0].password == "$2b$12$abcdefghijklmnop"
|
||
assert users[1].is_delete == 1
|
||
assert users[1].create_time is None
|
||
|
||
# Text 逐字节相等（中文/换行/HTML）
|
||
content = conn.execute(select(_table(loaded_metadata, "article").c.content)).scalar_one()
|
||
assert content == "第一行\n第二行 with <b>html</b> 中文"
|
||
|
||
# BigInteger 保真
|
||
assert conn.execute(select(_table(loaded_metadata, "job_log").c.consume_time)).scalar_one() \
|
||
== 9223372036854775807
|
||
|
||
# 目标新插 dept 自增为 MAX(id)+1（sqlite rowid 自动跟随显式插 id）
|
||
with dst_engine.begin() as conn:
|
||
conn.execute(_table(loaded_metadata, "dept").insert(), [{"name": "新部门", "code": "NEW", "parent_id": 0}])
|
||
with dst_engine.connect() as conn:
|
||
new_id = conn.execute(
|
||
select(_table(loaded_metadata, "dept").c.id).order_by(
|
||
_table(loaded_metadata, "dept").c.id.desc())).scalar()
|
||
assert new_id == 3
|
||
|
||
|
||
def test_chunking_roundtrip(two_sqlite, loaded_metadata):
|
||
"""chunk 循环边界：源 2500 行 city、chunk=1000，目标 count/ids 一致。"""
|
||
src_engine, dst_engine, src_path, dst_path = two_sqlite
|
||
city = _table(loaded_metadata, "city")
|
||
rows = [{"id": i, "name": f"区域{i}", "short_name": f"区域{i}", "city_code": "010",
|
||
"area_code": f"{i:06d}", "zip_code": "100000",
|
||
"parent_id": 0, "parent_code": None, "level": 1}
|
||
for i in range(1, 2501)]
|
||
with src_engine.begin() as conn:
|
||
conn.execute(city.insert(), rows)
|
||
|
||
assert migrate_db.run(make_args(str(src_path), str(dst_path), chunk_size=1000)) == 0
|
||
|
||
with dst_engine.connect() as conn:
|
||
assert conn.execute(text("SELECT COUNT(*) FROM fastapi_city")).scalar() == 2500
|
||
ids = [r[0] for r in conn.execute(select(city.c.id).order_by(city.c.id))]
|
||
assert ids == list(range(1, 2501))
|
||
|
||
|
||
def test_only_tables_filter(two_sqlite, loaded_metadata):
|
||
"""--only-tables 只泵指定表（支持全名/短名），其余表保持空。"""
|
||
src_engine, dst_engine, src_path, dst_path = two_sqlite
|
||
with src_engine.begin() as conn:
|
||
conn.execute(_table(loaded_metadata, "dept").insert(), [{"id": 1, "name": "总部", "code": "HQ", "parent_id": 0}])
|
||
conn.execute(_table(loaded_metadata, "city").insert(), [
|
||
{"id": 1, "name": "北京市", "short_name": "北京", "city_code": "010",
|
||
"area_code": "110000", "zip_code": "100000", "parent_id": 0}])
|
||
conn.execute(_table(loaded_metadata, "article").insert(), [{"id": 1, "title": "x"}])
|
||
|
||
assert migrate_db.run(make_args(str(src_path), str(dst_path),
|
||
only_tables="fastapi_dept,city")) == 0
|
||
|
||
with dst_engine.connect() as conn:
|
||
assert conn.execute(text("SELECT COUNT(*) FROM fastapi_dept")).scalar() == 1
|
||
assert conn.execute(text("SELECT COUNT(*) FROM fastapi_city")).scalar() == 1
|
||
assert conn.execute(text("SELECT COUNT(*) FROM fastapi_article")).scalar() == 0
|
||
|
||
|
||
def test_dry_run_does_not_touch_target(two_sqlite, loaded_metadata):
|
||
"""--dry-run 只打印计划，目标库零写入。"""
|
||
src_engine, dst_engine, src_path, dst_path = two_sqlite
|
||
with src_engine.begin() as conn:
|
||
conn.execute(_table(loaded_metadata, "dept").insert(), [{"id": 1, "name": "总部", "code": "HQ", "parent_id": 0}])
|
||
|
||
assert migrate_db.run(make_args(str(src_path), str(dst_path), dry_run=True)) == 0
|
||
|
||
with dst_engine.connect() as conn:
|
||
assert conn.execute(text("SELECT COUNT(*) FROM fastapi_dept")).scalar() == 0
|
||
|
||
|
||
def test_pg_reset_sql():
|
||
sql = migrate_db.pg_reset_sql("fastapi_user")
|
||
assert "setval" in sql
|
||
assert "pg_get_serial_sequence('fastapi_user', 'id')" in sql
|
||
assert 'FROM "fastapi_user"' in sql
|
||
# 空表用 setval(seq, 1, false)（next=1），不能 setval(seq, 0)（PG 报 NumericValueOutOfRange）
|
||
assert "COALESCE((SELECT MAX(id) FROM \"fastapi_user\"), 1)" in sql
|
||
assert "(SELECT MAX(id) FROM \"fastapi_user\") IS NOT NULL" in sql
|
||
|
||
|
||
def test_oracle_autoincrement_sql():
|
||
"""Oracle 自增：sequence 从 MAX+1 起，触发器 id 为 NULL 时取 NEXTVAL。"""
|
||
seq = migrate_db.oracle_sequence_sql("fastapi_user", 27)
|
||
assert seq == "CREATE SEQUENCE fastapi_user_id_seq START WITH 27 INCREMENT BY 1 NOCACHE"
|
||
trg = migrate_db.oracle_trigger_sql("fastapi_user")
|
||
assert "BEFORE INSERT ON fastapi_user FOR EACH ROW" in trg
|
||
assert "WHEN (NEW.id IS NULL)" in trg
|
||
assert "SELECT fastapi_user_id_seq.NEXTVAL INTO :NEW.id FROM dual" in trg
|
||
|
||
|
||
def test_mssql_reset_sql():
|
||
assert migrate_db.mssql_reset_sql("fastapi_user", 5) == "DBCC CHECKIDENT ('fastapi_user', RESEED, 5)"
|
||
# 空表 max_id 为 None → RESEED 0，下一 id=1
|
||
assert migrate_db.mssql_reset_sql("fastapi_user", None) == "DBCC CHECKIDENT ('fastapi_user', RESEED, 0)"
|
||
|
||
|
||
def test_create_database_if_missing_sqlite(tmp_path):
|
||
"""sqlite 建库：父目录自动创建；drop_first 删除已存在文件。"""
|
||
db = tmp_path / "nested" / "new.db"
|
||
migrate_db.create_database_if_missing("sqlite", "", None, str(db), "", "", drop_first=False)
|
||
eng = create_engine(migrate_db.make_url("sqlite", "", None, str(db), "", ""))
|
||
with eng.connect():
|
||
pass # 触发 sqlite 建文件
|
||
eng.dispose()
|
||
assert db.exists()
|
||
|
||
migrate_db.create_database_if_missing("sqlite", "", None, str(db), "", "", drop_first=True)
|
||
assert not db.exists()
|
||
|
||
|
||
def test_verify_counts_reports_mismatch(two_sqlite, loaded_metadata):
|
||
"""count 校验能报出不一致的表。"""
|
||
src_engine, dst_engine, src_path, dst_path = two_sqlite
|
||
with src_engine.begin() as conn:
|
||
conn.execute(_table(loaded_metadata, "dept").insert(), [{"id": 1, "name": "总部", "code": "HQ", "parent_id": 0}])
|
||
|
||
assert migrate_db.run(make_args(str(src_path), str(dst_path))) == 0
|
||
|
||
# 泵完后删一行制造不一致
|
||
with dst_engine.begin() as conn:
|
||
conn.execute(text("DELETE FROM fastapi_dept WHERE id = 1"))
|
||
|
||
with src_engine.connect() as src, dst_engine.connect() as dst:
|
||
failed = migrate_db.verify_counts(src, dst, migrate_db.ordered_tables())
|
||
assert any(name == "fastapi_dept" for name, _, _ in failed)
|