Neo-ZQYY/scripts/ops/run_seed_spi_params.py

# -*- coding: utf-8 -*-
"""
在测试库 test_etl_feiqiu 执行 SPI 种子数据脚本。

种子脚本：db/etl_feiqiu/seeds/seed_index_parameters.sql
目标表：dws.cfg_index_parameters（index_type='SPI'）

使用方式：
    python scripts/ops/run_seed_spi_params.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)

SEED_FILE = _ROOT / "db" / "etl_feiqiu" / "seeds" / "seed_index_parameters.sql"


def execute_seed(conn) -> bool:
    """执行种子数据脚本，返回是否成功"""
    sql = SEED_FILE.read_text(encoding="utf-8")
    if not sql.strip():
        print("⚠️ 种子脚本为空，跳过")
        return False

    try:
        cur = conn.cursor()
        cur.execute(sql)
        cur.close()
        conn.commit()
        print("✅ 种子数据脚本执行成功")
        return True
    except Exception as e:
        conn.rollback()
        print(f"❌ 种子数据脚本执行失败: {e}")
        return False


def verify(conn) -> bool:
    """验证 SPI 参数插入结果"""
    cur = conn.cursor()
    checks = []

    # 1. SPI 参数总数（应为 28 个）
    cur.execute("""
        SELECT COUNT(*) FROM dws.cfg_index_parameters
        WHERE index_type = 'SPI'
    """)
    spi_count = cur.fetchone()[0]
    checks.append((f"SPI 参数数量 = {spi_count}（期望 28）", spi_count == 28))

    # 2. 关键参数存在且值正确
    key_params = [
        ("weight_level", 0.60),
        ("weight_speed", 0.30),
        ("weight_stability", 0.10),
        ("amount_base_spend_30", 500.0),
        ("compression_mode", 1.0),
    ]
    for pname, expected in key_params:
        cur.execute("""
            SELECT param_value FROM dws.cfg_index_parameters
            WHERE index_type = 'SPI' AND param_name = %s
            ORDER BY effective_from DESC LIMIT 1
        """, (pname,))
        row = cur.fetchone()
        if row:
            actual = float(row[0])
            ok = abs(actual - expected) < 1e-6
            checks.append((f"  {pname} = {actual}（期望 {expected}）", ok))
        else:
            checks.append((f"  {pname} 缺失", False))

    # 3. 权重归一化：Level 子分权重之和 = 1.0
    cur.execute("""
        SELECT SUM(param_value) FROM dws.cfg_index_parameters
        WHERE index_type = 'SPI'
          AND param_name IN ('w_level_spend_30', 'w_level_spend_90',
                             'w_level_ticket_90', 'w_level_recharge_90')
    """)
    level_sum = float(cur.fetchone()[0] or 0)
    checks.append((f"  Level 权重之和 = {level_sum:.2f}（期望 1.00）", abs(level_sum - 1.0) < 1e-6))

    # 4. 总分权重之和 = 1.0
    cur.execute("""
        SELECT SUM(param_value) FROM dws.cfg_index_parameters
        WHERE index_type = 'SPI'
          AND param_name IN ('weight_level', 'weight_speed', 'weight_stability')
    """)
    total_sum = float(cur.fetchone()[0] or 0)
    checks.append((f"  总分权重之和 = {total_sum:.2f}（期望 1.00）", abs(total_sum - 1.0) < 1e-6))

    cur.close()

    print("\n" + "=" * 50)
    print("SPI 种子数据验证结果")
    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"种子脚本: {SEED_FILE.name}\n")

    if not SEED_FILE.exists():
        print(f"ERROR: 种子脚本不存在: {SEED_FILE}")
        sys.exit(1)

    conn = psycopg2.connect(DSN)

    if not execute_seed(conn):
        conn.close()
        sys.exit(1)

    all_ok = verify(conn)
    conn.close()

    if all_ok:
        print("\n✅ SPI 种子数据执行完成，所有验证通过")
    else:
        print("\n⚠️ 部分验证未通过，请检查")
        sys.exit(1)


if __name__ == "__main__":
    main()