Neo-ZQYY/scripts/ops/init_test_user.py

# -*- coding: utf-8 -*-
"""
一键初始化测试用户 — 打通认证→业务页面的完整链路。

用途：
  在没有租户管理后台的情况下，直接在 test_zqyy_app 中插入必要的
  映射数据，让测试用户能通过 dev-login 登录后正常访问业务页面。

操作内容：
  1. 确保 biz.sites 中有真实门店注册记录
  2. 创建/更新测试用户为 approved 状态
  3. 分配门店角色（user_site_roles）
  4. 创建助教绑定（user_assistant_binding）
  5. 输出可用于测试的 JWT token

使用方式：
  cd C:\\Project\\NeoZQYY
  python scripts/ops/init_test_user.py [--openid <openid>] [--role <role_code>] [--reset]

环境要求：
  - 根 .env 中配置 APP_DB_DSN（指向 test_zqyy_app）
  - WX_DEV_MODE=true（dev-login 端点需要）
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from pathlib import Path

import psycopg2
from dotenv import load_dotenv

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

APP_DB_DSN = os.environ.get("APP_DB_DSN")
if not APP_DB_DSN:
    print("❌ 环境变量 APP_DB_DSN 未定义，请检查根 .env", file=sys.stderr)
    sys.exit(1)

# 安全检查：禁止连正式库
if "test_" not in APP_DB_DSN and "test-" not in APP_DB_DSN:
    print("❌ APP_DB_DSN 不包含 'test_'，疑似正式库，拒绝执行", file=sys.stderr)
    sys.exit(1)

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger(__name__)

# ── ETL 真实数据（从 test_etl_feiqiu 查询所得） ──────────
# 朗朗桌球 — 唯一的测试门店
REAL_SITE_ID = 2790685415443269
REAL_TENANT_ID = 2790683160709957
REAL_SHOP_NAME = "朗朗桌球"
SITE_CODE = "LL001"  # 已存在的映射

# 默认绑定的助教（李楚欣，欣欣）
DEFAULT_ASSISTANT_ID = 2793361915547781
DEFAULT_ASSISTANT_NAME = "李楚欣"

# 默认绑定的员工（厉超）
DEFAULT_STAFF_ID = 3020235774380805
DEFAULT_STAFF_NAME = "厉超"


# ── 角色映射 ──────────────────────────────────────────────
ROLE_MAP = {
    "coach": 1,       # 助教
    "staff": 2,       # 员工
    "site_admin": 3,  # 店铺管理员
    "tenant_admin": 4,  # 租户管理员
}


def get_conn():
    """获取数据库连接。"""
    conn = psycopg2.connect(APP_DB_DSN)
    conn.autocommit = False
    return conn


def _safe_tenant_id():
    """返回真实 tenant_id（列已迁移为 bigint，无需范围检查）。"""
    return REAL_TENANT_ID


def ensure_site_registration(cur) -> int:
    """确保 biz.sites 中有朗朗桌球的注册记录，返回 site_id。"""
    cur.execute(
        "SELECT site_id FROM biz.sites WHERE site_code = %s",
        (SITE_CODE,),
    )
    row = cur.fetchone()
    if row:
        existing_site_id = row[0]
        if existing_site_id != REAL_SITE_ID:
            # site_id 不匹配（旧测试数据），更新为真实值
            cur.execute(
                """
                UPDATE biz.sites
                SET site_id = %s, site_name = %s, tenant_id = %s
                WHERE site_code = %s
                """,
                (REAL_SITE_ID, REAL_SHOP_NAME, _safe_tenant_id(), SITE_CODE),
            )
            log.info("✅ biz.sites 已更新: %s → %s (%s)（旧值 %s）",
                     SITE_CODE, REAL_SITE_ID, REAL_SHOP_NAME, existing_site_id)
        else:
            log.info("✅ biz.sites 已存在: %s → %s (%s)", SITE_CODE, REAL_SITE_ID, REAL_SHOP_NAME)
        return REAL_SITE_ID

    cur.execute(
        """
        INSERT INTO biz.sites (site_code, site_id, site_name, tenant_id)
        VALUES (%s, %s, %s, %s)
        RETURNING site_id
        """,
        (SITE_CODE, REAL_SITE_ID, REAL_SHOP_NAME, _safe_tenant_id()),
    )
    site_id = cur.fetchone()[0]
    log.info("✅ 已创建 biz.sites: %s → %s (%s)", SITE_CODE, site_id, REAL_SHOP_NAME)
    return site_id


def ensure_user(cur, openid: str, nickname: str, reset: bool) -> tuple[int, str]:
    """创建或更新测试用户，返回 (user_id, status)。"""
    cur.execute(
        "SELECT id, status FROM auth.users WHERE wx_openid = %s",
        (openid,),
    )
    row = cur.fetchone()

    if row and not reset:
        user_id, status = row
        if status != "approved":
            cur.execute(
                "UPDATE auth.users SET status = 'approved', updated_at = NOW() WHERE id = %s",
                (user_id,),
            )
            log.info("✅ 用户 %s (id=%d) 状态更新: %s → approved", openid, user_id, status)
            return user_id, "approved"
        log.info("✅ 用户已存在且已审核: %s (id=%d)", openid, user_id)
        return user_id, status

    if row and reset:
        user_id = row[0]
        # 清理旧数据
        cur.execute("DELETE FROM auth.user_assistant_binding WHERE user_id = %s", (user_id,))
        cur.execute("DELETE FROM auth.user_site_roles WHERE user_id = %s", (user_id,))
        cur.execute("DELETE FROM auth.user_applications WHERE user_id = %s", (user_id,))
        cur.execute(
            """
            UPDATE auth.users
            SET status = 'approved', nickname = %s, updated_at = NOW()
            WHERE id = %s
            """,
            (nickname, user_id),
        )
        log.info("✅ 用户 %s (id=%d) 已重置为 approved", openid, user_id)
        return user_id, "approved"

    # 新建用户
    cur.execute(
        """
        INSERT INTO auth.users (wx_openid, nickname, status)
        VALUES (%s, %s, 'approved')
        RETURNING id
        """,
        (openid, nickname),
    )
    user_id = cur.fetchone()[0]
    log.info("✅ 已创建用户: %s (id=%d, nickname=%s)", openid, user_id, nickname)
    return user_id, "approved"


def ensure_site_role(cur, user_id: int, site_id: int, role_code: str) -> None:
    """分配门店角色。"""
    role_id = ROLE_MAP.get(role_code)
    if not role_id:
        log.error("❌ 未知角色: %s（可选: %s）", role_code, ", ".join(ROLE_MAP))
        sys.exit(1)

    cur.execute(
        """
        SELECT id FROM auth.user_site_roles
        WHERE user_id = %s AND site_id = %s AND role_id = %s
        """,
        (user_id, site_id, role_id),
    )
    if cur.fetchone():
        log.info("✅ 角色已分配: user_id=%d, site_id=%d, role=%s", user_id, site_id, role_code)
        return

    cur.execute(
        """
        INSERT INTO auth.user_site_roles (user_id, site_id, role_id)
        VALUES (%s, %s, %s)
        """,
        (user_id, site_id, role_id),
    )
    log.info("✅ 已分配角色: user_id=%d, site_id=%d, role=%s (role_id=%d)", user_id, site_id, role_code, role_id)


def ensure_binding(cur, user_id: int, site_id: int, role_code: str) -> None:
    """创建助教/员工绑定。"""
    if role_code == "coach":
        binding_type = "assistant"
        assistant_id = DEFAULT_ASSISTANT_ID
        staff_id = None
        label = f"助教 {DEFAULT_ASSISTANT_NAME} (id={assistant_id})"
    elif role_code == "staff":
        binding_type = "staff"
        assistant_id = None
        staff_id = DEFAULT_STAFF_ID
        label = f"员工 {DEFAULT_STAFF_NAME} (id={staff_id})"
    elif role_code in ("site_admin", "tenant_admin"):
        binding_type = "manager"
        assistant_id = None
        staff_id = DEFAULT_STAFF_ID
        label = f"管理员绑定员工 {DEFAULT_STAFF_NAME} (id={staff_id})"
    else:
        log.warning("⚠️ 未知角色 %s，跳过绑定", role_code)
        return

    cur.execute(
        """
        SELECT id FROM auth.user_assistant_binding
        WHERE user_id = %s AND site_id = %s
        """,
        (user_id, site_id),
    )
    if cur.fetchone():
        log.info("✅ 绑定已存在: user_id=%d, site_id=%d", user_id, site_id)
        return

    cur.execute(
        """
        INSERT INTO auth.user_assistant_binding
            (user_id, site_id, assistant_id, staff_id, binding_type)
        VALUES (%s, %s, %s, %s, %s)
        """,
        (user_id, site_id, assistant_id, staff_id, binding_type),
    )
    log.info("✅ 已创建绑定: %s → %s", f"user_id={user_id}", label)


def print_summary(openid: str, user_id: int, site_id: int, role_code: str) -> None:
    """输出测试信息摘要。"""
    log.info("")
    log.info("=" * 60)
    log.info("  测试用户初始化完成")
    log.info("=" * 60)
    log.info("  openid:    %s", openid)
    log.info("  user_id:   %d", user_id)
    log.info("  site_id:   %d", site_id)
    log.info("  site_code: %s (%s)", SITE_CODE, REAL_SHOP_NAME)
    log.info("  角色:      %s (role_id=%d)", role_code, ROLE_MAP[role_code])
    log.info("  状态:      approved")
    log.info("")
    log.info("  下一步：")
    log.info("  1. 启动后端: cd apps/backend && uvicorn app.main:app --reload")
    log.info("  2. 调用 dev-login 获取 JWT:")
    log.info('     curl -X POST http://localhost:8000/api/xcx/dev-login \\')
    log.info('       -H "Content-Type: application/json" \\')
    log.info('       -d \'{"openid": "%s", "status": "approved"}\'', openid)
    log.info("")
    log.info("  3. 用返回的 access_token 测试业务接口:")
    log.info('     curl http://localhost:8000/api/xcx/me \\')
    log.info('       -H "Authorization: Bearer <access_token>"')
    log.info("=" * 60)


def main():
    parser = argparse.ArgumentParser(
        description="一键初始化测试用户，打通认证→业务页面完整链路",
    )
    parser.add_argument(
        "--openid",
        default="dev_test_openid",
        help="测试用户的 wx_openid（默认: dev_test_openid）",
    )
    parser.add_argument(
        "--nickname",
        default="测试管理员",
        help="测试用户昵称（默认: 测试管理员）",
    )
    parser.add_argument(
        "--role",
        default="site_admin",
        choices=list(ROLE_MAP.keys()),
        help="分配的角色（默认: site_admin）",
    )
    parser.add_argument(
        "--reset",
        action="store_true",
        help="重置用户：清除旧的角色/绑定/申请，重新初始化",
    )
    args = parser.parse_args()

    conn = get_conn()
    try:
        with conn.cursor() as cur:
            site_id = ensure_site_registration(cur)
            user_id, _ = ensure_user(cur, args.openid, args.nickname, args.reset)
            ensure_site_role(cur, user_id, site_id, args.role)
            ensure_binding(cur, user_id, site_id, args.role)
        conn.commit()
        print_summary(args.openid, user_id, site_id, args.role)
    except Exception:
        conn.rollback()
        log.exception("❌ 初始化失败，已回滚")
        sys.exit(1)
    finally:
        conn.close()


if __name__ == "__main__":
    main()