Neo-ZQYY/scripts/ops/validate_p1_db_foundation.py

# -*- coding: utf-8 -*-
"""
P1 数据库基础设施层端到端验证脚本。

检查项：
  1. 业务库 auth / biz Schema 存在性
  2. ETL 库 app Schema 及 35 张 RLS 视图存在性
  3. 业务库 fdw_etl Schema 及外部表存在性 + 可查询性
  4. RLS 视图 site_id 过滤正确性
  5. app_user / app_reader 角色权限配置

用法：
  python scripts/ops/validate_p1_db_foundation.py
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

from dotenv import load_dotenv

# ── 环境变量加载 ──────────────────────────────────────────────
_ROOT = Path(__file__).resolve().parents[2]
load_dotenv(_ROOT / ".env", override=False)

PG_DSN = os.environ.get("PG_DSN")
APP_DB_DSN = os.environ.get("APP_DB_DSN")

_missing = []
if not PG_DSN:
    _missing.append("PG_DSN")
if not APP_DB_DSN:
    _missing.append("APP_DB_DSN")
if _missing:
    raise RuntimeError(
        f"必需环境变量缺失: {', '.join(_missing)}。"
        "请在根 .env 中配置 PG_DSN 和 APP_DB_DSN。"
    )

import psycopg2  # noqa: E402  — 延迟导入，确保环境变量校验先行

# ── 35 张 RLS 视图清单 ────────────────────────────────────────
EXPECTED_RLS_VIEWS: list[str] = [
    # DWD 层（11 张）
    "v_dim_member",
    "v_dim_assistant",
    "v_dim_member_card_account",
    "v_dim_table",
    "v_dwd_settlement_head",
    "v_dwd_table_fee_log",
    "v_dwd_assistant_service_log",
    "v_dwd_recharge_order",
    "v_dwd_store_goods_sale",
    "v_dim_staff",
    "v_dim_staff_ex",
    # DWS 层（24 张）
    "v_dws_member_consumption_summary",
    "v_dws_member_visit_detail",
    "v_dws_member_winback_index",
    "v_dws_member_newconv_index",
    "v_dws_member_recall_index",
    "v_dws_member_assistant_relation_index",
    "v_dws_member_assistant_intimacy",
    "v_dws_assistant_daily_detail",
    "v_dws_assistant_monthly_summary",
    "v_dws_assistant_salary_calc",
    "v_dws_assistant_customer_stats",
    "v_dws_assistant_finance_analysis",
    "v_dws_finance_daily_summary",
    "v_dws_finance_income_structure",
    "v_dws_finance_recharge_summary",
    "v_dws_finance_discount_detail",
    "v_dws_finance_expense_summary",
    "v_dws_platform_settlement",
    "v_dws_assistant_recharge_commission",
    "v_cfg_performance_tier",
    "v_cfg_assistant_level_price",
    "v_cfg_bonus_rules",
    "v_cfg_index_parameters",
    "v_dws_order_summary",
]


# ── 辅助函数 ──────────────────────────────────────────────────

def _connect(dsn: str, label: str):
    """建立数据库连接，失败时输出脱敏信息。"""
    try:
        conn = psycopg2.connect(dsn)
        return conn
    except psycopg2.OperationalError as exc:
        # 脱敏：只显示 host/dbname，不泄露密码
        safe = dsn.split("@")[-1] if "@" in dsn else "(unknown)"
        print(f"❌ 无法连接 {label}（{safe}）: {exc}", file=sys.stderr)
        raise


def _schema_exists(cur, schema_name: str) -> bool:
    cur.execute(
        "SELECT 1 FROM information_schema.schemata WHERE schema_name = %s",
        (schema_name,),
    )
    return cur.fetchone() is not None


def _view_exists(cur, schema_name: str, view_name: str) -> bool:
    cur.execute(
        "SELECT 1 FROM information_schema.views "
        "WHERE table_schema = %s AND table_name = %s",
        (schema_name, view_name),
    )
    return cur.fetchone() is not None


def _foreign_table_exists(cur, schema_name: str, table_name: str) -> bool:
    cur.execute(
        "SELECT 1 FROM information_schema.tables "
        "WHERE table_schema = %s AND table_name = %s AND table_type = 'FOREIGN'",
        (schema_name, table_name),
    )
    return cur.fetchone() is not None


# ── 核心验证逻辑 ──────────────────────────────────────────────

def validate_p1_db_foundation() -> dict:
    """
    返回验证结果字典：
    {
        "schemas": {"auth": bool, "biz": bool, "app": bool, "fdw_etl": bool},
        "rls_views": {"app.v_dim_member": bool, ...},
        "fdw_tables": {"fdw_etl.v_dim_member": bool, ...},
        "rls_filtering": bool | None,   # None = SKIP
        "permissions": {"app_user": bool, "app_reader": bool},
        "errors": [str, ...]
    }
    """
    result: dict = {
        "schemas": {},
        "rls_views": {},
        "fdw_tables": {},
        "rls_filtering": None,
        "permissions": {},
        "errors": [],
    }

    etl_conn = _connect(PG_DSN, "ETL 库")
    app_conn = _connect(APP_DB_DSN, "业务库")

    try:
        _check_schemas(etl_conn, app_conn, result)
        _check_rls_views(etl_conn, result)
        _check_fdw_tables(app_conn, result)
        _check_rls_filtering(etl_conn, result)
        _check_permissions(etl_conn, app_conn, result)
    finally:
        etl_conn.close()
        app_conn.close()

    return result


def _check_schemas(etl_conn, app_conn, result: dict):
    """检查 auth / biz（业务库）和 app / fdw_etl（ETL 库 + 业务库）Schema 存在性。"""
    with app_conn.cursor() as cur:
        for s in ("auth", "biz", "fdw_etl"):
            ok = _schema_exists(cur, s)
            result["schemas"][s] = ok
            if not ok:
                result["errors"].append(f"业务库缺少 Schema: {s}")

    with etl_conn.cursor() as cur:
        ok = _schema_exists(cur, "app")
        result["schemas"]["app"] = ok
        if not ok:
            result["errors"].append("ETL 库缺少 Schema: app")


def _check_rls_views(etl_conn, result: dict):
    """检查 ETL 库 app Schema 中 35 张 RLS 视图是否存在。"""
    with etl_conn.cursor() as cur:
        for vname in EXPECTED_RLS_VIEWS:
            ok = _view_exists(cur, "app", vname)
            result["rls_views"][f"app.{vname}"] = ok
            if not ok:
                result["errors"].append(f"ETL 库缺少 RLS 视图: app.{vname}")


def _check_fdw_tables(app_conn, result: dict):
    """检查业务库 fdw_etl Schema 中外部表存在性 + 可查询性。

    cfg_* 表无 RLS 过滤，可直接 SELECT count(*)。
    其余 RLS 表的远端视图需要 app.current_site_id，
    先获取一个有效 site_id 再统一查询。
    """
    # 无 RLS 过滤的配置表
    cfg_views = {
        "v_cfg_performance_tier",
        "v_cfg_assistant_level_price",
        "v_cfg_bonus_rules",
        "v_cfg_index_parameters",
    }

    # 先从一张 cfg 表确认 FDW 链路可用
    with app_conn.cursor() as cur:
        for vname in EXPECTED_RLS_VIEWS:
            exists = _foreign_table_exists(cur, "fdw_etl", vname)
            key = f"fdw_etl.{vname}"
            if not exists:
                result["fdw_tables"][key] = False
                result["errors"].append(f"业务库缺少外部表: {key}")
                continue

            if vname in cfg_views:
                # 无 RLS，直接查询
                try:
                    cur.execute(f"SELECT count(*) FROM fdw_etl.{vname}")
                    cur.fetchone()
                    result["fdw_tables"][key] = True
                except Exception as exc:
                    app_conn.rollback()
                    result["fdw_tables"][key] = False
                    result["errors"].append(f"外部表 {key} 查询失败: {exc}")
            else:
                # RLS 表：远端需要 app.current_site_id，此处仅验证存在性
                # 可查询性在 _check_fdw_rls_queryability 中统一验证
                result["fdw_tables"][key] = True  # 存在即通过

    # 对 RLS 外部表做可查询性抽查（通过 dblink 在远端设置 site_id）
    _check_fdw_rls_queryability(app_conn, result, cfg_views)


def _check_fdw_rls_queryability(app_conn, result: dict, cfg_views: set):
    """通过 ETL 库直连验证 RLS 外部表的可查询性。

    FDW 远端会话无法继承本地 SET 的 session 变量，
    因此改为：直连 ETL 库设置 site_id 后查询 app.v_* 视图，
    间接证明 FDW 映射的源视图可查询。
    """
    etl_conn = _connect(PG_DSN, "ETL 库（FDW 可查询性验证）")
    try:
        with etl_conn.cursor() as cur:
            # 获取一个有效 site_id
            try:
                cur.execute(
                    "SELECT DISTINCT site_id FROM dws.dws_member_consumption_summary LIMIT 1"
                )
                row = cur.fetchone()
            except Exception:
                etl_conn.rollback()
                return  # 无法获取 site_id，跳过

            if row is None:
                return

            site_id = row[0]
            cur.execute("SET app.current_site_id = %s", (str(site_id),))

            # 抽查一张 RLS 视图
            rls_sample = "v_dws_member_consumption_summary"
            try:
                cur.execute(f"SELECT count(*) FROM app.{rls_sample}")
                cnt = cur.fetchone()[0]
                if cnt == 0:
                    # 有 site_id 但无数据，不算失败
                    pass
            except Exception as exc:
                key = f"fdw_etl.{rls_sample}"
                result["fdw_tables"][key] = False
                result["errors"].append(
                    f"RLS 外部表源视图 app.{rls_sample} 查询失败: {exc}"
                )
                etl_conn.rollback()
    finally:
        etl_conn.close()


def _check_rls_filtering(etl_conn, result: dict):
    """设置 site_id 后验证 RLS 视图过滤正确性。"""
    # 从 DWS 表取一个实际存在的 site_id（DWS 表使用 site_id 列）
    with etl_conn.cursor() as cur:
        try:
            cur.execute(
                "SELECT DISTINCT site_id FROM dws.dws_member_consumption_summary LIMIT 1"
            )
            row = cur.fetchone()
        except Exception as exc:
            result["rls_filtering"] = None
            result["errors"].append(f"无法获取 site_id 样本: {exc}")
            etl_conn.rollback()
            return

    if row is None:
        result["rls_filtering"] = None
        return

    site_id = row[0]

    # 新连接，设置 site_id 后查询
    verify_conn = _connect(PG_DSN, "ETL 库（RLS 验证）")
    try:
        with verify_conn.cursor() as cur:
            cur.execute("SET app.current_site_id = %s", (str(site_id),))
            cur.execute(
                "SELECT site_id FROM app.v_dws_member_consumption_summary LIMIT 100"
            )
            rows = cur.fetchall()
            if not rows:
                result["rls_filtering"] = None  # SKIP — 该 site_id 无数据
                return
            all_match = all(r[0] == site_id for r in rows)
            result["rls_filtering"] = all_match
            if not all_match:
                result["errors"].append(
                    f"RLS 过滤失败: 设置 site_id={site_id} 后 "
                    f"v_dws_member_consumption_summary 返回了其他门店数据"
                )
    finally:
        verify_conn.close()


def _check_permissions(etl_conn, app_conn, result: dict):
    """验证 app_reader（ETL 库）和 app_user（业务库）角色权限。"""
    # ── app_reader：对 app Schema 的 USAGE + 视图 SELECT ──
    app_reader_ok = True
    with etl_conn.cursor() as cur:
        try:
            cur.execute(
                "SELECT has_schema_privilege('app_reader', 'app', 'USAGE')"
            )
            if not cur.fetchone()[0]:
                app_reader_ok = False
                result["errors"].append("app_reader 缺少 app Schema USAGE 权限")

            # 抽查一张视图的 SELECT 权限
            cur.execute(
                "SELECT has_table_privilege('app_reader', 'app.v_dim_member', 'SELECT')"
            )
            if not cur.fetchone()[0]:
                app_reader_ok = False
                result["errors"].append(
                    "app_reader 缺少 app.v_dim_member SELECT 权限"
                )
        except Exception as exc:
            etl_conn.rollback()
            app_reader_ok = False
            result["errors"].append(f"app_reader 权限检查异常: {exc}")

    result["permissions"]["app_reader"] = app_reader_ok

    # ── app_user：对 auth / biz 的 USAGE ──
    app_user_ok = True
    with app_conn.cursor() as cur:
        try:
            for schema in ("auth", "biz"):
                cur.execute(
                    "SELECT has_schema_privilege('app_user', %s, 'USAGE')",
                    (schema,),
                )
                if not cur.fetchone()[0]:
                    app_user_ok = False
                    result["errors"].append(
                        f"app_user 缺少 {schema} Schema USAGE 权限"
                    )
        except Exception as exc:
            app_conn.rollback()
            app_user_ok = False
            result["errors"].append(f"app_user 权限检查异常: {exc}")

    result["permissions"]["app_user"] = app_user_ok


# ── 输出格式化 ────────────────────────────────────────────────

def _icon(val) -> str:
    if val is None:
        return "⏭️"
    return "✅" if val else "❌"


def print_report(result: dict):
    """打印结构化验证报告。"""
    print("\n" + "=" * 60)
    print("  P1 数据库基础设施层验证报告")
    print("=" * 60)

    # Schema 检查
    print("\n📦 Schema 存在性")
    for name, ok in result["schemas"].items():
        print(f"  {_icon(ok)} {name}")

    # RLS 视图
    print(f"\n👁️  RLS 视图（共 {len(EXPECTED_RLS_VIEWS)} 张）")
    passed = sum(1 for v in result["rls_views"].values() if v)
    failed = sum(1 for v in result["rls_views"].values() if not v)
    print(f"  ✅ 通过: {passed}  ❌ 失败: {failed}")
    for name, ok in result["rls_views"].items():
        if not ok:
            print(f"  ❌ {name}")

    # FDW 外部表
    print(f"\n🔗 FDW 外部表（共 {len(EXPECTED_RLS_VIEWS)} 张）")
    passed_fdw = sum(1 for v in result["fdw_tables"].values() if v)
    failed_fdw = sum(1 for v in result["fdw_tables"].values() if not v)
    print(f"  ✅ 通过: {passed_fdw}  ❌ 失败: {failed_fdw}")
    for name, ok in result["fdw_tables"].items():
        if not ok:
            print(f"  ❌ {name}")

    # RLS 过滤
    print(f"\n🔒 RLS 过滤正确性: {_icon(result['rls_filtering'])}", end="")
    if result["rls_filtering"] is None:
        print("  (无数据，跳过验证)")
    else:
        print()

    # 权限
    print("\n🔑 角色权限")
    for role, ok in result["permissions"].items():
        print(f"  {_icon(ok)} {role}")

    # 汇总
    total_checks = (
        len(result["schemas"])
        + len(result["rls_views"])
        + len(result["fdw_tables"])
        + (1 if result["rls_filtering"] is not None else 0)
        + len(result["permissions"])
    )
    total_pass = (
        sum(1 for v in result["schemas"].values() if v)
        + sum(1 for v in result["rls_views"].values() if v)
        + sum(1 for v in result["fdw_tables"].values() if v)
        + (1 if result["rls_filtering"] is True else 0)
        + sum(1 for v in result["permissions"].values() if v)
    )
    total_skip = 1 if result["rls_filtering"] is None else 0

    print(f"\n{'=' * 60}")
    print(f"  汇总: {total_pass}/{total_checks} 通过", end="")
    if total_skip:
        print(f"  ({total_skip} 项跳过)", end="")
    print()

    if result["errors"]:
        print(f"\n⚠️  失败详情（共 {len(result['errors'])} 项）:")
        for err in result["errors"]:
            print(f"  • {err}")

    print("=" * 60 + "\n")


# ── 入口 ──────────────────────────────────────────────────────

if __name__ == "__main__":
    result = validate_p1_db_foundation()
    print_report(result)

    # 退出码：有失败项则非零
    has_failure = bool(result["errors"])
    sys.exit(1 if has_failure else 0)