"""
第 13 章 · 演示 3：行级安全 RLS · 多租户隔离
==============================================

场景：ch13_tickets 表有 tenant_id 字段，开启 RLS。
      用 tenant_a / tenant_b 两个账号分别连进去，验证：
      1) 各自只看到自己 tenant 的数据；
      2) 不能跨 tenant 写入；
      3) RESTRICTIVE 策略屏蔽 archived；
      4) dba_user (BYPASSRLS) 看到全部。

依赖：
    pip install "psycopg[binary]>=3.1"
    必须先 psql -f ../init.sql 跑过初始化。

运行：
    python 03_rls_multi_tenant.py
"""

from __future__ import annotations

import psycopg

DSN_BASE = "host=127.0.0.1 port=5432 dbname=learn_pg "

USERS = {
    "tenant_a": ("TenantA@2025", 1),    # (密码, 租户号)
    "tenant_b": ("TenantB@2025", 2),
    "dba_user": ("Dba@2025",     None), # BYPASSRLS，无需 setting
}


def line(t: str) -> None:
    print("\n" + "=" * 60 + "\n" + t + "\n" + "=" * 60)


def query_as(user: str, password: str, tenant: int | None) -> None:
    """以指定用户身份查询 ch13_tickets，可选注入 app.tenant_id。"""
    dsn = f"{DSN_BASE}user={user} password={password}"
    print(f"\n--- 以 {user} 身份连接 (tenant={tenant}) ---")
    with psycopg.connect(dsn) as conn, conn.cursor() as cur:
        cur.execute("BEGIN")
        if tenant is not None:
            cur.execute("SET LOCAL app.tenant_id = %s", (str(tenant),))
        cur.execute("SELECT id, tenant_id, title, status FROM ch13_tickets ORDER BY id;")
        rows = cur.fetchall()
        for r in rows:
            print(f"    id={r[0]:<3} tenant={r[1]} status={r[2]:<8} title={r[3]}")
        if not rows:
            print("    (空)")
        conn.rollback()


def try_cross_tenant_insert() -> None:
    """tenant_a 试图写入 tenant_id=2 的行，应被 WITH CHECK 拦截。"""
    line("Step 2：tenant_a 试图越权插入 tenant_id=2 的行")
    dsn = f"{DSN_BASE}user=tenant_a password=TenantA@2025"
    with psycopg.connect(dsn) as conn, conn.cursor() as cur:
        cur.execute("BEGIN")
        cur.execute("SET LOCAL app.tenant_id = '1'")   # 自报身份是租户1
        try:
            cur.execute(
                "INSERT INTO ch13_tickets(tenant_id, title) VALUES (2, '我是 A 试图伪装成 B')"
            )
            print("  ❌ 写入成功了，RLS 失效！请检查 init.sql。")
        except psycopg.errors.InsufficientPrivilege as e:
            # 注：触发 WITH CHECK 失败时报的是 InsufficientPrivilege 或 CheckViolation
            print(f"  ✅ 被 WITH CHECK 拦截：{type(e).__name__}: {str(e).splitlines()[0]}")
        except Exception as e:
            print(f"  ✅ 被拦截：{type(e).__name__}: {str(e).splitlines()[0]}")
        conn.rollback()


def show_archived_filter() -> None:
    """RESTRICTIVE 策略：ch13_no_archived 应该让 archived 行隐身（即使 owner 也不行）。"""
    line("Step 3：验证 RESTRICTIVE 策略 ch13_no_archived（archived 行不可见）")
    dsn = f"{DSN_BASE}user=tenant_a password=TenantA@2025"
    with psycopg.connect(dsn) as conn, conn.cursor() as cur:
        cur.execute("BEGIN")
        cur.execute("SET LOCAL app.tenant_id = '1'")
        cur.execute("SELECT status, count(*) FROM ch13_tickets GROUP BY status ORDER BY status;")
        for r in cur.fetchall():
            print(f"    status={r[0]:<8} count={r[1]}")
        print("  → 应该看不到 archived 状态（init.sql 写入时确实有一行 archived）")
        conn.rollback()


def superuser_view() -> None:
    line("Step 4：用 superuser (postgres) 连接，确认数据库里实际数据全貌")
    dsn = f"{DSN_BASE}user=postgres"
    with psycopg.connect(dsn) as conn, conn.cursor() as cur:
        # 由于 init.sql 用了 FORCE ROW LEVEL SECURITY，但 superuser 默认 BYPASSRLS
        cur.execute("SELECT tenant_id, status, count(*) FROM ch13_tickets "
                    "GROUP BY tenant_id, status ORDER BY tenant_id, status;")
        for r in cur.fetchall():
            print(f"    tenant={r[0]} status={r[1]:<8} count={r[2]}")


def main() -> None:
    line("Step 1：以两个租户身份分别 SELECT，应该只看到自己的行")
    query_as("tenant_a", "TenantA@2025", 1)
    query_as("tenant_b", "TenantB@2025", 2)

    try_cross_tenant_insert()
    show_archived_filter()
    superuser_view()

    line("Step 5：dba_user (BYPASSRLS) 不需要设置 app.tenant_id 也能看全表")
    query_as("dba_user", "Dba@2025", None)


if __name__ == "__main__":
    main()