190 lines
5.5 KiB
Python
190 lines
5.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
在测试库 test_etl_feiqiu 执行员工档案建表迁移脚本。
|
||
|
||
迁移脚本：db/etl_feiqiu/migrations/2026-02-22__add_staff_info_tables.sql
|
||
目标表：ods.staff_info_master, dwd.dim_staff, dwd.dim_staff_ex
|
||
|
||
使用方式：
|
||
python scripts/ops/run_migration_staff_info.py
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from dotenv import load_dotenv
|
||
import psycopg2
|
||
|
||
# 加载根 .env
|
||
_ROOT = Path(__file__).resolve().parents[2]
|
||
load_dotenv(_ROOT / ".env", override=False)
|
||
|
||
DSN = os.getenv("TEST_DB_DSN")
|
||
if not DSN:
|
||
print("ERROR: TEST_DB_DSN 未配置，请在根 .env 中设置")
|
||
sys.exit(1)
|
||
|
||
MIGRATION_FILE = (
|
||
_ROOT / "db" / "etl_feiqiu" / "migrations"
|
||
/ "2026-02-22__add_staff_info_tables.sql"
|
||
)
|
||
|
||
# 需要创建的三张表
|
||
TABLES = [
|
||
("ods", "staff_info_master"),
|
||
("dwd", "dim_staff"),
|
||
("dwd", "dim_staff_ex"),
|
||
]
|
||
|
||
|
||
def tables_exist(conn) -> dict[str, bool]:
|
||
"""检查目标表是否已存在，返回 {schema.table: bool}"""
|
||
cur = conn.cursor()
|
||
result = {}
|
||
for schema, table in TABLES:
|
||
cur.execute("""
|
||
SELECT 1 FROM information_schema.tables
|
||
WHERE table_schema = %s AND table_name = %s
|
||
""", (schema, table))
|
||
result[f"{schema}.{table}"] = cur.fetchone() is not None
|
||
cur.close()
|
||
return result
|
||
|
||
|
||
def execute_migration(conn) -> bool:
|
||
"""执行迁移脚本，返回是否成功"""
|
||
sql = MIGRATION_FILE.read_text(encoding="utf-8")
|
||
|
||
# 去掉注释中的回滚部分
|
||
main_lines = []
|
||
in_rollback = False
|
||
for line in sql.split("\n"):
|
||
stripped = line.strip()
|
||
if stripped.startswith("-- ====") and "回滚" in stripped:
|
||
in_rollback = True
|
||
if not in_rollback:
|
||
main_lines.append(line)
|
||
|
||
main_sql = "\n".join(main_lines).strip()
|
||
if not main_sql:
|
||
print("⚠️ 迁移脚本为空，跳过")
|
||
return False
|
||
|
||
try:
|
||
cur = conn.cursor()
|
||
cur.execute(main_sql)
|
||
cur.close()
|
||
print("✅ 迁移脚本执行成功")
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ 迁移脚本执行失败: {e}")
|
||
return False
|
||
|
||
|
||
def verify(conn) -> bool:
|
||
"""验证建表结果"""
|
||
cur = conn.cursor()
|
||
checks = []
|
||
|
||
# 1. 三张表都存在
|
||
for schema, table in TABLES:
|
||
cur.execute("""
|
||
SELECT 1 FROM information_schema.tables
|
||
WHERE table_schema = %s AND table_name = %s
|
||
""", (schema, table))
|
||
checks.append((f"表 {schema}.{table} 存在", cur.fetchone() is not None))
|
||
|
||
# 2. ods.staff_info_master 关键字段
|
||
cur.execute("""
|
||
SELECT column_name FROM information_schema.columns
|
||
WHERE table_schema = 'ods' AND table_name = 'staff_info_master'
|
||
""")
|
||
ods_cols = {r[0] for r in cur.fetchall()}
|
||
ods_required = {"id", "staff_name", "mobile", "content_hash", "payload", "tenant_id", "site_id"}
|
||
missing_ods = ods_required - ods_cols
|
||
checks.append((f"ODS 关键字段完整（{len(ods_cols)} 列）", len(missing_ods) == 0))
|
||
if missing_ods:
|
||
print(f" ODS 缺失字段: {missing_ods}")
|
||
|
||
# 3. dwd.dim_staff 主键包含 staff_id + scd2_start_time
|
||
cur.execute("""
|
||
SELECT column_name FROM information_schema.columns
|
||
WHERE table_schema = 'dwd' AND table_name = 'dim_staff'
|
||
""")
|
||
dwd_cols = {r[0] for r in cur.fetchall()}
|
||
dwd_required = {"staff_id", "staff_name", "scd2_start_time", "scd2_end_time", "scd2_is_current"}
|
||
missing_dwd = dwd_required - dwd_cols
|
||
checks.append((f"DWD 主表关键字段完整（{len(dwd_cols)} 列）", len(missing_dwd) == 0))
|
||
if missing_dwd:
|
||
print(f" DWD 主表缺失字段: {missing_dwd}")
|
||
|
||
# 4. dwd.dim_staff_ex 关键字段
|
||
cur.execute("""
|
||
SELECT column_name FROM information_schema.columns
|
||
WHERE table_schema = 'dwd' AND table_name = 'dim_staff_ex'
|
||
""")
|
||
ex_cols = {r[0] for r in cur.fetchall()}
|
||
ex_required = {"staff_id", "rank_name", "shop_name", "scd2_start_time"}
|
||
missing_ex = ex_required - ex_cols
|
||
checks.append((f"DWD 扩展表关键字段完整（{len(ex_cols)} 列）", len(missing_ex) == 0))
|
||
if missing_ex:
|
||
print(f" DWD 扩展表缺失字段: {missing_ex}")
|
||
|
||
cur.close()
|
||
|
||
print("\n" + "=" * 50)
|
||
print("建表验证结果")
|
||
print("=" * 50)
|
||
all_ok = True
|
||
for name, ok in checks:
|
||
status = "✅" if ok else "❌"
|
||
print(f" {status} {name}")
|
||
if not ok:
|
||
all_ok = False
|
||
|
||
return all_ok
|
||
|
||
|
||
def main():
|
||
dsn_display = DSN.split("@")[1] if "@" in DSN else DSN
|
||
print(f"连接测试库: {dsn_display}")
|
||
print(f"迁移脚本: {MIGRATION_FILE.name}\n")
|
||
|
||
if not MIGRATION_FILE.exists():
|
||
print(f"ERROR: 迁移脚本不存在: {MIGRATION_FILE}")
|
||
sys.exit(1)
|
||
|
||
conn = psycopg2.connect(DSN)
|
||
conn.autocommit = True
|
||
|
||
# 检查表是否已存在
|
||
existing = tables_exist(conn)
|
||
all_exist = all(existing.values())
|
||
|
||
if all_exist:
|
||
print("ℹ️ 所有目标表已存在，跳过建表")
|
||
else:
|
||
for name, exists in existing.items():
|
||
if exists:
|
||
print(f" ℹ️ {name} 已存在")
|
||
else:
|
||
print(f" 📋 {name} 待创建")
|
||
if not execute_migration(conn):
|
||
conn.close()
|
||
sys.exit(1)
|
||
|
||
# 验证
|
||
all_ok = verify(conn)
|
||
conn.close()
|
||
|
||
if all_ok:
|
||
print("\n✅ 员工档案建表迁移完成，所有验证通过")
|
||
else:
|
||
print("\n⚠️ 部分验证未通过，请检查")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|