"""
05_procedure_with_commit.py
---------------------------
PG 11+ 存储过程演示：内部 COMMIT 实现"分批归档"。
  init.sql 已经创建好 ch12_archive_done_orders(batch_size INT) 过程。
  这里用 psycopg 调用 CALL，并观察日志输出（通过捕获 NOTICE）。

依赖: pip install psycopg[binary]>=3.1
运行: python 05_procedure_with_commit.py
"""
import psycopg

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


def prepare_data(conn):
    """先把一批订单状态改成 done，然后调用过程归档"""
    with conn.cursor() as cur:
        cur.execute("UPDATE ch12_orders SET status='done' "
                    "WHERE id IN (SELECT id FROM ch12_orders LIMIT 50)")
        cur.execute("SELECT count(*) FROM ch12_orders WHERE status='done'")
        print("待归档 done 订单数 =", cur.fetchone()[0])


def call_procedure():
    """注意：调用包含 COMMIT 的过程时，连接必须开启 autocommit。"""
    with psycopg.connect(DSN, autocommit=True) as conn:
        prepare_data(conn)

        notices = []
        conn.add_notice_handler(lambda diag: notices.append(diag.message_primary))

        with conn.cursor() as cur:
            print("\n>>> CALL ch12_archive_done_orders(20) ...")
            cur.execute("CALL ch12_archive_done_orders(20)")

        print("\n--- 过程内 RAISE NOTICE 输出 ---")
        for m in notices:
            print(" ", m)

        with conn.cursor() as cur:
            cur.execute("SELECT count(*) FROM ch12_orders WHERE status='done'")
            print("\n剩余 done 订单数 =", cur.fetchone()[0])
            cur.execute("SELECT count(*) FROM ch12_orders_archive")
            print("已归档总数     =", cur.fetchone()[0])


def main():
    call_procedure()


if __name__ == "__main__":
    main()