"""
03_event_trigger.py
-------------------
事件触发器演示：
  ① 安装 "禁止 DROP TABLE" 事件触发器
  ② 安装 "记录所有 DDL" 事件触发器
  ③ 故意 CREATE / DROP TABLE，观察拦截与日志
  ④ 清理触发器

依赖:  pip install psycopg[binary]>=3.1
运行:  python 03_event_trigger.py
"""
import psycopg
from psycopg import errors

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


SQL_CREATE_NO_DROP_FUNC = """
CREATE OR REPLACE FUNCTION ch12_trg_no_drop_table() RETURNS event_trigger
LANGUAGE plpgsql AS $$
DECLARE obj RECORD;
BEGIN
    FOR obj IN SELECT * FROM pg_event_trigger_dropped_objects()
                WHERE object_type = 'table'
    LOOP
        RAISE EXCEPTION '禁止 DROP TABLE %', obj.object_identity
              USING ERRCODE = 'P0001', HINT = '请联系 DBA';
    END LOOP;
END;
$$;
"""

SQL_CREATE_LOG_DDL_FUNC = """
CREATE OR REPLACE FUNCTION ch12_trg_log_ddl() RETURNS event_trigger
LANGUAGE plpgsql AS $$
DECLARE r RECORD;
BEGIN
    FOR r IN SELECT * FROM pg_event_trigger_ddl_commands() LOOP
        INSERT INTO ch12_ddl_log (usr, db, command_tag, object_type, object_identity)
        VALUES (current_user, current_database(),
                r.command_tag, r.object_type, r.object_identity);
    END LOOP;
END;
$$;
"""


def setup(conn):
    with conn.cursor() as cur:
        cur.execute(SQL_CREATE_NO_DROP_FUNC)
        cur.execute(SQL_CREATE_LOG_DDL_FUNC)
        cur.execute("DROP EVENT TRIGGER IF EXISTS ch12_no_drop_table")
        cur.execute("DROP EVENT TRIGGER IF EXISTS ch12_log_ddl")
        cur.execute(
            "CREATE EVENT TRIGGER ch12_no_drop_table "
            "ON sql_drop EXECUTE FUNCTION ch12_trg_no_drop_table()"
        )
        cur.execute(
            "CREATE EVENT TRIGGER ch12_log_ddl "
            "ON ddl_command_end EXECUTE FUNCTION ch12_trg_log_ddl()"
        )
    print("事件触发器已安装")


def teardown(conn):
    with conn.cursor() as cur:
        cur.execute("DROP EVENT TRIGGER IF EXISTS ch12_no_drop_table")
        cur.execute("DROP EVENT TRIGGER IF EXISTS ch12_log_ddl")
    print("事件触发器已清理")


def show_ddl_log(conn, n=5):
    with conn.cursor() as cur:
        cur.execute("SELECT id, ts, usr, command_tag, object_type, object_identity "
                    "  FROM ch12_ddl_log ORDER BY id DESC LIMIT %s", (n,))
        print("\n--- ch12_ddl_log 最近 {} 行 ---".format(n))
        for r in cur.fetchall():
            print(r)


def main():
    with psycopg.connect(DSN, autocommit=True) as conn:
        setup(conn)
        try:
            with conn.cursor() as cur:
                print("\n>>> 创建临时表 ch12_demo_evt_test ...")
                cur.execute("DROP TABLE IF EXISTS ch12_demo_evt_test")
                cur.execute("CREATE TABLE ch12_demo_evt_test (id INT)")

                print(">>> 尝试 DROP TABLE，应被禁止 ...")
                try:
                    cur.execute("DROP TABLE ch12_demo_evt_test")
                    print("意外！没有被拦截")
                except errors.RaiseException as e:
                    print(f"被事件触发器拦截 ✓: {e.diag.message_primary}")

            show_ddl_log(conn, 5)

        finally:
            teardown(conn)
            with conn.cursor() as cur:
                cur.execute("DROP TABLE IF EXISTS ch12_demo_evt_test")


if __name__ == "__main__":
    main()