"""
第 13 章 · 演示 1：ROLE 创建、组继承、GRANT/REVOKE 全流程
=========================================================

依赖：
    pip install "psycopg[binary]>=3.1"

前置：
    数据库 learn_pg 已存在，且当前用户是 superuser（postgres）。
    建议先 psql -f ../init.sql 跑过初始化，本脚本不依赖 init.sql 中的 ch13_tickets 表，
    会创建/销毁自己专属的演示对象。

运行：
    python 01_role_grant.py
"""

from __future__ import annotations

import psycopg

DSN = "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres"


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


def show_roles(conn: psycopg.Connection, names: list[str]) -> None:
    """打印若干 ROLE 的属性。"""
    sql = """
        SELECT rolname,
               rolsuper, rolcreatedb, rolcreaterole,
               rolcanlogin, rolinherit, rolbypassrls,
               rolconnlimit, rolvaliduntil
        FROM pg_roles
        WHERE rolname = ANY(%s)
        ORDER BY rolname;
    """
    with conn.cursor() as cur:
        cur.execute(sql, (names,))
        rows = cur.fetchall()
    for r in rows:
        print(f"  {r[0]:<10} super={r[1]} createdb={r[2]} createrole={r[3]} "
              f"login={r[4]} inherit={r[5]} bypassrls={r[6]} conn_limit={r[7]}")


def cleanup(conn: psycopg.Connection) -> None:
    """删掉演示用的对象（幂等）。"""
    with conn.cursor() as cur:
        cur.execute("DROP TABLE IF EXISTS demo_payroll CASCADE;")
        for role in ("hr_alice", "hr_bob", "engineer_carol", "hr_team"):
            cur.execute(f"DROP OWNED BY {role} CASCADE;") if role_exists(cur, role) else None
            cur.execute(f"DROP ROLE IF EXISTS {role};")
    conn.commit()


def role_exists(cur: psycopg.Cursor, name: str) -> bool:
    cur.execute("SELECT 1 FROM pg_roles WHERE rolname = %s", (name,))
    return cur.fetchone() is not None


def step1_create_roles(conn: psycopg.Connection) -> None:
    line("Step 1：创建组角色 + 用户角色")
    with conn.cursor() as cur:
        cur.execute("CREATE ROLE hr_team NOLOGIN;")                              # 组
        cur.execute("CREATE ROLE hr_alice LOGIN PASSWORD 'Hr_Alice@2025' INHERIT;")
        cur.execute("CREATE ROLE hr_bob   LOGIN PASSWORD 'Hr_Bob@2025'   NOINHERIT;")
        cur.execute("CREATE ROLE engineer_carol LOGIN PASSWORD 'Eng_Carol@2025';")
        cur.execute("GRANT hr_team TO hr_alice, hr_bob;")
    conn.commit()
    show_roles(conn, ["hr_alice", "hr_bob", "engineer_carol", "hr_team"])


def step2_create_table_and_grant(conn: psycopg.Connection) -> None:
    line("Step 2：创建一张 demo_payroll 表，把 SELECT/UPDATE 授权给 hr_team")
    with conn.cursor() as cur:
        cur.execute("""
            CREATE TABLE demo_payroll(
                id SERIAL PRIMARY KEY,
                emp_name TEXT,
                salary  NUMERIC(10,2)
            );
        """)
        cur.execute("INSERT INTO demo_payroll(emp_name, salary) VALUES "
                    "('Alice', 18000), ('Bob', 22000), ('Carol', 35000);")
        cur.execute("GRANT SELECT, UPDATE ON demo_payroll TO hr_team;")
        # SERIAL 列依赖 sequence
        cur.execute("GRANT USAGE, SELECT, UPDATE ON SEQUENCE demo_payroll_id_seq TO hr_team;")
    conn.commit()
    print("  表已建好，权限授给 hr_team 组。")


def step3_inherit_vs_noinherit() -> None:
    line("Step 3：演示 INHERIT vs NOINHERIT 的差别")
    # hr_alice INHERIT，连上来直接能 SELECT
    with psycopg.connect("host=127.0.0.1 port=5432 dbname=learn_pg "
                         "user=hr_alice password=Hr_Alice@2025") as alice:
        with alice.cursor() as cur:
            cur.execute("SELECT count(*) FROM demo_payroll;")
            print(f"  hr_alice (INHERIT) SELECT 成功，行数 = {cur.fetchone()[0]}")

    # hr_bob NOINHERIT，连上来默认不带 hr_team 权限，必须 SET ROLE
    with psycopg.connect("host=127.0.0.1 port=5432 dbname=learn_pg "
                         "user=hr_bob password=Hr_Bob@2025") as bob:
        with bob.cursor() as cur:
            try:
                cur.execute("SELECT count(*) FROM demo_payroll;")
                print(f"  hr_bob (NOINHERIT) 直接 SELECT = {cur.fetchone()[0]}（不应该看到这行）")
            except psycopg.errors.InsufficientPrivilege as e:
                print(f"  hr_bob (NOINHERIT) 直接 SELECT → 被拒绝：{type(e).__name__}")
            bob.rollback()
            cur.execute("SET ROLE hr_team;")
            cur.execute("SELECT count(*) FROM demo_payroll;")
            print(f"  hr_bob 在 SET ROLE hr_team 后 SELECT 成功，行数 = {cur.fetchone()[0]}")


def step4_engineer_denied() -> None:
    line("Step 4：未授权账号 engineer_carol 应该被拒绝")
    with psycopg.connect("host=127.0.0.1 port=5432 dbname=learn_pg "
                         "user=engineer_carol password=Eng_Carol@2025") as eng:
        with eng.cursor() as cur:
            try:
                cur.execute("SELECT count(*) FROM demo_payroll;")
            except psycopg.errors.InsufficientPrivilege:
                print("  engineer_carol SELECT demo_payroll → 被拒绝 ✅")


def step5_show_acl(conn: psycopg.Connection) -> None:
    line("Step 5：查看 ACL（access privileges）")
    with conn.cursor() as cur:
        cur.execute("""
            SELECT relname, relacl
            FROM pg_class
            WHERE relname = 'demo_payroll';
        """)
        row = cur.fetchone()
        print(f"  表 {row[0]} 的 ACL = {row[1]}")
        print("  解读：每条形如 'grantee=权限/grantor'，权限缩写见文档 13.2.2 节")


def main() -> None:
    with psycopg.connect(DSN) as conn:
        cleanup(conn)
        try:
            step1_create_roles(conn)
            step2_create_table_and_grant(conn)
            step3_inherit_vs_noinherit()
            step4_engineer_denied()
            step5_show_acl(conn)
        finally:
            line("收尾：清理演示对象")
            cleanup(conn)
            print("  已清理。")


if __name__ == "__main__":
    main()