#!/usr/bin/env python3
"""
01_range_partition_demo.py
==========================
演示 RANGE 分区表：
  1. 从零建表 + 6 个月分区 + default 分区
  2. 灌入跨月数据
  3. 验证数据自动路由到对应分区
  4. EXPLAIN 看分区裁剪
  5. 演示「插入 default 分区」和「添加新分区」

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

import psycopg

DSN = os.environ.get(
    "PG_DSN",
    "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres password=postgres",
)


def hr(s):
    print("\n" + "=" * 68)
    print(f"  {s}")
    print("=" * 68)


def run(cur, sql, fetch=True):
    cur.execute(sql)
    if fetch and cur.description:
        rows = cur.fetchall()
        for r in rows:
            print(" ", r)
        return rows


def main():
    with psycopg.connect(DSN, autocommit=True) as conn, conn.cursor() as cur:
        hr("Step 1: 清理 + 建分区父表")
        run(cur, "DROP TABLE IF EXISTS ch16_demo_orders CASCADE;", fetch=False)
        run(cur, """
            CREATE TABLE ch16_demo_orders (
                id BIGSERIAL,
                user_id BIGINT NOT NULL,
                amount NUMERIC(12,2),
                created_at TIMESTAMPTZ NOT NULL,
                PRIMARY KEY (id, created_at)
            ) PARTITION BY RANGE (created_at);
        """, fetch=False)
        print("  ✓ 父表 ch16_demo_orders 已建（PARTITION BY RANGE created_at）")

        hr("Step 2: 建 3 个月分区 + default 分区")
        for ym, frm, to in [("2025_01", "2025-01-01", "2025-02-01"),
                            ("2025_02", "2025-02-01", "2025-03-01"),
                            ("2025_03", "2025-03-01", "2025-04-01")]:
            run(cur, f"""
                CREATE TABLE ch16_demo_orders_{ym} PARTITION OF ch16_demo_orders
                    FOR VALUES FROM ('{frm}') TO ('{to}');
            """, fetch=False)
        run(cur, "CREATE TABLE ch16_demo_orders_default PARTITION OF ch16_demo_orders DEFAULT;",
            fetch=False)
        print("  ✓ 3 个月分区 + default 分区建好")

        hr("Step 3: 父表索引（PG 11+ 自动级联到所有子分区）")
        run(cur, "CREATE INDEX ON ch16_demo_orders (user_id);", fetch=False)
        print("  ✓ 在父表建索引，自动下发到所有子分区")
        run(cur, """
            SELECT indexrelid::regclass AS index_name
            FROM pg_index
            WHERE indrelid IN (
                SELECT inhrelid FROM pg_inherits WHERE inhparent='ch16_demo_orders'::regclass
            )
            ORDER BY index_name;
        """)

        hr("Step 4: 插入跨月数据（PG 自动路由到对应分区）")
        run(cur, """
            INSERT INTO ch16_demo_orders (user_id, amount, created_at) VALUES
                (1, 100.00, '2025-01-15 10:00:00'),
                (2, 200.00, '2025-01-20 11:00:00'),
                (3, 300.00, '2025-02-05 12:00:00'),
                (4, 400.00, '2025-02-25 13:00:00'),
                (5, 500.00, '2025-03-10 14:00:00'),
                (6, 600.00, '2025-09-01 15:00:00')   -- 落入 default！
            RETURNING id, created_at;
        """)

        hr("Step 5: 各分区行数")
        run(cur, """
            SELECT
                relname AS partition,
                pg_stat_get_live_tuples(c.oid)::INT AS rows
            FROM pg_class c
            JOIN pg_inherits i ON i.inhrelid = c.oid
            WHERE i.inhparent = 'ch16_demo_orders'::regclass
            ORDER BY relname;
        """)
        # 强制 ANALYZE 更新统计信息
        run(cur, "ANALYZE ch16_demo_orders;", fetch=False)

        hr("Step 6: EXPLAIN 演示分区裁剪")
        print("\n--- 6.1 WHERE 用字面量常量（计划期裁剪） ---")
        run(cur, """
            EXPLAIN SELECT * FROM ch16_demo_orders
            WHERE created_at >= '2025-02-01' AND created_at < '2025-03-01';
        """)

        print("\n--- 6.2 WHERE 不能用：表达式套在分区键上 ---")
        run(cur, """
            EXPLAIN SELECT * FROM ch16_demo_orders
            WHERE EXTRACT(MONTH FROM created_at) = 2;
        """)

        print("\n--- 6.3 关闭分区裁剪做对比 ---")
        run(cur, "SET enable_partition_pruning = off;", fetch=False)
        run(cur, """
            EXPLAIN SELECT * FROM ch16_demo_orders
            WHERE created_at >= '2025-02-01' AND created_at < '2025-03-01';
        """)
        run(cur, "RESET enable_partition_pruning;", fetch=False)

        hr("Step 7: 添加 4 月分区，把 default 里的数据搬走")
        # 演示常见运维操作：default 里有了未匹配数据，需要先建对应的分区
        run(cur, """
            CREATE TABLE ch16_demo_orders_2025_04 PARTITION OF ch16_demo_orders
                FOR VALUES FROM ('2025-04-01') TO ('2025-05-01');
        """, fetch=False)
        print("  ✓ 4 月分区已建，但 default 里 2025-09-01 的行还需要单独处理")
        # 实际上 default 里的 09 月行不在 04 月范围，只能新建 09 月分区或保留
        run(cur, """
            SELECT created_at FROM ch16_demo_orders_default;
        """)

        hr("Step 8: DETACH CONCURRENTLY（PG 14+ 不锁表下线分区）")
        # 注意：DETACH CONCURRENTLY 不能在事务里执行
        try:
            run(cur, "ALTER TABLE ch16_demo_orders DETACH PARTITION ch16_demo_orders_2025_01 CONCURRENTLY;",
                fetch=False)
            print("  ✓ 1 月分区已下线（变成普通独立表 ch16_demo_orders_2025_01）")
        except psycopg.Error as e:
            print(f"  ⚠️  DETACH 失败（可能 PG 版本 < 14）：{e}")

        hr("最终各分区状态")
        run(cur, """
            SELECT
                relname,
                pg_stat_get_live_tuples(c.oid)::INT AS rows,
                pg_size_pretty(pg_relation_size(c.oid)) AS size
            FROM pg_class c
            JOIN pg_inherits i ON i.inhrelid = c.oid
            WHERE i.inhparent = 'ch16_demo_orders'::regclass
            ORDER BY relname;
        """)

        print("\n[DONE] 完整演示结束。可以打开 demo.html 看可视化效果。")
        return 0


if __name__ == "__main__":
    sys.exit(main() or 0)