"""
01_plpgsql_function.py
----------------------
演示 psycopg 调用三种 PL/pgSQL 函数：
  ① 标量函数 ch12_compound_interest(principal, rate, years)
  ② 表函数  ch12_list_orders(user_id, limit, offset)
  ③ 异常分支：故意传非法参数触发 RAISE EXCEPTION

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

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


def demo_scalar(conn):
    print("\n=== ① 标量函数：ch12_compound_interest ===")
    with conn.cursor() as cur:
        cur.execute(
            "SELECT ch12_compound_interest(10000, 0.05, %s) AS amount",
            (10,),
        )
        amt = cur.fetchone()[0]
        print(f"本金 10000、年化 5%、复利 10 年 = {amt}")


def demo_table(conn):
    print("\n=== ② 表函数：ch12_list_orders ===")
    with conn.cursor() as cur:
        cur.execute("SELECT * FROM ch12_list_orders(%s, %s, %s)", (1, 5, 0))
        for row in cur.fetchall():
            print(row)


def demo_exception(conn):
    print("\n=== ③ 异常分支：负参数 ===")
    try:
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT ch12_compound_interest(-1, 0.05, 10)"
                )
    except errors.RaiseException as e:
        print(f"PG 抛出业务异常: SQLSTATE={e.diag.sqlstate}, msg={e.diag.message_primary}")
    except psycopg.errors.Error as e:
        print(f"其他 PG 异常: {type(e).__name__}: {e}")


def main():
    with psycopg.connect(DSN, autocommit=True) as conn:
        demo_scalar(conn)
        demo_table(conn)
        demo_exception(conn)


if __name__ == "__main__":
    main()