#!/usr/bin/env python3
"""
03_partition_pruning.py
=======================
对比「裁剪 vs 不裁剪」的性能差异，并展示什么样的 WHERE 条件能 / 不能触发裁剪。

要求：先跑过 init.sql（已有 ch16_orders 分区表 + 6 万行数据）。
"""
import os
import sys
import time

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" + "=" * 70)
    print(f"  {s}")
    print("=" * 70)


def explain(cur, sql, label):
    print(f"\n[{label}]")
    print(f"SQL: {sql}")
    cur.execute(f"EXPLAIN (ANALYZE, BUFFERS) {sql}")
    rows = [r[0] for r in cur.fetchall()]
    # 统计扫描的分区数
    scanned = sum(1 for line in rows if "Scan on ch16_orders_2025" in line)
    pruned = next((line for line in rows if "Subplans Removed" in line), None)
    timing = next((line for line in rows if "Execution Time" in line), "")
    print("  --- EXPLAIN ---")
    for line in rows:
        print("  ", line)
    print(f"  >>> 扫描了 {scanned} 个分区"
          + (f"   {pruned.strip()}" if pruned else "")
          + f"   {timing.strip()}")


def main():
    with psycopg.connect(DSN, autocommit=True) as conn, conn.cursor() as cur:
        # 检查表是否存在
        cur.execute("""
            SELECT count(*) FROM pg_class WHERE relname = 'ch16_orders' AND relkind = 'p';
        """)
        if cur.fetchone()[0] == 0:
            print("[FATAL] 表 ch16_orders 不存在或不是分区表，请先跑 init.sql")
            return 1

        cur.execute("ANALYZE ch16_orders;")

        hr("场景 A：完美裁剪（WHERE 直接对分区键比较）")
        explain(cur,
                "SELECT count(*) FROM ch16_orders WHERE created_at >= '2025-03-01' AND created_at < '2025-04-01'",
                "A.1 范围常量")
        explain(cur,
                "SELECT count(*) FROM ch16_orders WHERE created_at = '2025-04-15 12:00:00'",
                "A.2 点查常量")
        explain(cur,
                "SELECT count(*) FROM ch16_orders WHERE created_at IN ('2025-02-15','2025-04-15','2025-05-15')",
                "A.3 IN 多值")

        hr("场景 B：失败裁剪（WHERE 套了表达式）")
        explain(cur,
                "SELECT count(*) FROM ch16_orders WHERE EXTRACT(MONTH FROM created_at) = 3",
                "B.1 函数包裹分区键")
        explain(cur,
                "SELECT count(*) FROM ch16_orders WHERE created_at + INTERVAL '1 day' >= '2025-03-01'",
                "B.2 加表达式")
        explain(cur,
                "SELECT count(*) FROM ch16_orders WHERE created_at::date = '2025-03-15'",
                "B.3 cast 改变类型")

        hr("场景 C：执行期裁剪（PREPARE 参数化查询）")
        cur.execute("DEALLOCATE ALL;")
        cur.execute("PREPARE q (timestamptz, timestamptz) AS "
                    "SELECT count(*) FROM ch16_orders WHERE created_at >= $1 AND created_at < $2;")
        # 跑几次让规划器生成 generic plan（PG 用启发式 5 次自定义后切 generic）
        for _ in range(7):
            cur.execute("EXECUTE q ('2025-04-01'::timestamptz, '2025-05-01'::timestamptz);")

        explain(cur,
                "EXECUTE q ('2025-04-01'::timestamptz, '2025-05-01'::timestamptz)",
                "C.1 PREPARE + EXECUTE（看 Subplans Removed: N）")
        cur.execute("DEALLOCATE q;")

        hr("场景 D：开 vs 关 enable_partition_pruning 性能对比")
        sql = ("SELECT count(*) FROM ch16_orders "
               "WHERE created_at >= '2025-03-01' AND created_at < '2025-04-01'")

        for setting in ("on", "off"):
            cur.execute(f"SET enable_partition_pruning = {setting};")
            # 跑 5 次取最小耗时
            times = []
            for _ in range(5):
                t0 = time.perf_counter()
                cur.execute(sql)
                cur.fetchall()
                times.append((time.perf_counter() - t0) * 1000)
            print(f"  enable_partition_pruning={setting:>3}  最快 {min(times):.2f}ms"
                  f"  平均 {sum(times)/len(times):.2f}ms")
        cur.execute("RESET enable_partition_pruning;")

        hr("场景 E：partitionwise_join 跨分区表 JOIN 优化（PG 11+）")
        # 简单演示：构造另一张同分区键的表然后 JOIN
        cur.execute("DROP TABLE IF EXISTS ch16_payments CASCADE;")
        cur.execute("""
            CREATE TABLE ch16_payments (
                id BIGSERIAL,
                order_id BIGINT,
                paid_at TIMESTAMPTZ NOT NULL,
                PRIMARY KEY (id, paid_at)
            ) PARTITION BY RANGE (paid_at);
        """)
        for ym, frm, to in [("2025_03", "2025-03-01", "2025-04-01"),
                            ("2025_04", "2025-04-01", "2025-05-01")]:
            cur.execute(f"CREATE TABLE ch16_payments_{ym} PARTITION OF ch16_payments "
                        f"FOR VALUES FROM ('{frm}') TO ('{to}');")
        cur.execute("""
            INSERT INTO ch16_payments (order_id, paid_at)
            SELECT id, created_at FROM ch16_orders
            WHERE created_at >= '2025-03-01' AND created_at < '2025-05-01'
            LIMIT 5000;
        """)
        cur.execute("ANALYZE ch16_payments;")

        sql_join = ("SELECT count(*) FROM ch16_orders o JOIN ch16_payments p "
                    "ON o.id = p.order_id AND o.created_at = p.paid_at "
                    "WHERE o.created_at >= '2025-03-01' AND o.created_at < '2025-05-01'")

        for setting in ("off", "on"):
            cur.execute(f"SET enable_partitionwise_join = {setting};")
            cur.execute(f"EXPLAIN (ANALYZE) {sql_join}")
            timing = next((r[0] for r in cur.fetchall() if "Execution Time" in r[0]), "")
            print(f"  enable_partitionwise_join={setting}  {timing.strip()}")
        cur.execute("RESET enable_partitionwise_join;")
        cur.execute("DROP TABLE ch16_payments CASCADE;")

        hr("总结")
        print("  ✓ WHERE 直接对分区键比较 → 完美裁剪")
        print("  ✓ PG 11+ PREPARE 后参数查询 → 执行期裁剪（Subplans Removed）")
        print("  ✗ 函数 / cast / 表达式包裹分区键 → 失败")
        print("  💡 partitionwise_join 默认 OFF，做大宽表 JOIN 记得开")

    return 0


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