#!/bin/bash
# 盈尺（Surinch）PostgreSQL 分区表 / 分区裁剪实验 v1.0
# 只能在测试实例上跑：会 DROP/CREATE 库 ${DB}（名字必须以 _lab 结尾）。
# 测试：PostgreSQL 17.11（Debian 包），shared_buffers=1GB，自造数据 400 万行（setseed 固定），非客户数据。
# 用法：PSQL="psql -h 127.0.0.1 -U postgres" DB=part_lab ./pg-partition-pruning-lab-v1.0.sh | tee result.txt
set -uo pipefail
PSQL=${PSQL:-"psql"}
DB=${DB:-part_lab}
case "$DB" in *_lab) ;; *) echo "DB 名必须以 _lab 结尾（防止误删业务库）"; exit 1;; esac
P="$PSQL -X -q -v ON_ERROR_STOP=0 -d $DB"
ts(){ date '+%F %T.%3N'; }

init(){
  echo "[$(ts)] 建库造数"
  $PSQL -X -q -d postgres -c "DROP DATABASE IF EXISTS $DB" -c "CREATE DATABASE $DB"
  $P <<'SQL'
\timing on
-- 1) 按月分区（2025-01 .. 2026-12，共 24 个）
CREATE TABLE ev_m (id bigint NOT NULL, user_id int NOT NULL, event_type smallint NOT NULL,
                   created_at timestamp NOT NULL, payload text NOT NULL) PARTITION BY RANGE (created_at);
DO $$ DECLARE d date := '2025-01-01'; BEGIN
  WHILE d < '2027-01-01' LOOP
    EXECUTE format('CREATE TABLE ev_m_%s PARTITION OF ev_m FOR VALUES FROM (%L) TO (%L)', to_char(d,'YYYYMM'), d, d + interval '1 month');
    d := d + interval '1 month';
  END LOOP; END $$;
-- 2) 按天分区（2025-01-01 .. 2026-12-31，共 730 个）
CREATE TABLE ev_d (LIKE ev_m) PARTITION BY RANGE (created_at);
DO $$ DECLARE d date := '2025-01-01'; BEGIN
  WHILE d < '2027-01-01' LOOP
    EXECUTE format('CREATE TABLE ev_d_%s PARTITION OF ev_d FOR VALUES FROM (%L) TO (%L)', to_char(d,'YYYYMMDD'), d, d + 1);
    d := d + 1;
  END LOOP; END $$;
-- 3) 不分区对照表
CREATE TABLE ev_flat (LIKE ev_m);
-- 自造 400 万行：2025-01-01 到 2026-10-24 均匀分布，10 万用户
SELECT setseed(0.42);
INSERT INTO ev_flat
SELECT n, 1 + floor(random()*100000)::int, (n % 8)::smallint,
       timestamp '2025-01-01' + random() * (timestamp '2026-10-24' - timestamp '2025-01-01'),
       md5(n::text)
FROM generate_series(1, 4000000) n;
INSERT INTO ev_m SELECT * FROM ev_flat;
INSERT INTO ev_d SELECT * FROM ev_flat;
CREATE INDEX ON ev_flat (created_at); CREATE INDEX ON ev_flat (user_id, created_at);
CREATE INDEX ON ev_m (created_at);    CREATE INDEX ON ev_m (user_id, created_at);
CREATE INDEX ON ev_d (created_at);    CREATE INDEX ON ev_d (user_id, created_at);
VACUUM ANALYZE ev_flat; VACUUM ANALYZE ev_m; VACUUM ANALYZE ev_d;
\timing off
SELECT relname, (SELECT count(*) FROM pg_inherits WHERE inhparent = c.oid) AS partitions,
       pg_size_pretty(pg_total_relation_size(c.oid)) AS flat_size
FROM pg_class c WHERE relname IN ('ev_flat','ev_m','ev_d') ORDER BY relname;
SELECT count(*), min(created_at), max(created_at) FROM ev_flat;
SQL
}

# 完整 EXPLAIN ANALYZE 一次，再跑 2 次只看 Planning / Execution Time
q(){
  local tag=$1 sql=$2
  echo; echo "### [$tag] $sql"
  # 计划很长（730 个分区）时只打印前 24 行，另外统计各分区扫描节点执行/未执行的个数（不含 Bitmap Index Scan，避免同一分区数两次）
  $P -c "EXPLAIN (ANALYZE, BUFFERS) $sql" | awk '
    / Scan( .*)? on ev_/ && !/Bitmap Index Scan/ { if ($0 ~ /never executed/) ne++; else ex++ }
    NR<=24 {print; next}
    /Planning Time|Execution Time|Subplans Removed/ {print}
    END { if (NR>24) print "  ...（计划共 " NR " 行，已截断）"; if (ex+ne>0) printf "  分区扫描节点：执行 %d 个，never executed %d 个\n", ex, ne }' | cut -c1-200
  for i in 2 3; do
    echo "  [$tag 第 $i 次] $($P -At -c "EXPLAIN (ANALYZE, BUFFERS) $sql" | grep -E 'Planning Time|Execution Time' | tr '\n' ' ')"
  done
}
# 只看计划（不执行）
plan(){ echo; echo "### [$1 计划] $2"; $P -c "EXPLAIN (COSTS OFF) $2" | cut -c1-200; }

[ "${SKIP_INIT:-0}" = 1 ] || init
echo "[$(ts)] 预热"; $P -At -c "SELECT count(*) FROM ev_flat" -c "SELECT count(*) FROM ev_m" -c "SELECT count(*) FROM ev_d" >/dev/null
$P -At -c "SELECT 'today=' || CURRENT_DATE"

echo; echo "===== P1 常量条件：计划期裁剪"
W="created_at >= '2026-09-01' AND created_at < '2026-09-08'"
q P1a "SELECT count(*), sum(event_type) FROM ev_flat WHERE $W"
q P1b "SELECT count(*), sum(event_type) FROM ev_m WHERE $W"
q P1c "SELECT count(*), sum(event_type) FROM ev_d WHERE $W"

echo; echo "===== P2 CURRENT_DATE / now()：执行期裁剪（Subplans Removed）"
plan P2a "SELECT count(*) FROM ev_m WHERE created_at >= CURRENT_DATE - 7"
q    P2a "SELECT count(*) FROM ev_m WHERE created_at >= CURRENT_DATE - 7"
plan P2b "SELECT count(*) FROM ev_m WHERE created_at >= now() - interval '7 days'"
q    P2b "SELECT count(*) FROM ev_m WHERE created_at >= now() - interval '7 days'"

echo; echo "===== P3 分区键上套函数 / 类型转换：裁剪失效"
q P3a "SELECT count(*) FROM ev_m WHERE date_trunc('month', created_at) = '2026-09-01'"
q P3b "SELECT count(*) FROM ev_m WHERE created_at::date = '2026-09-03'"
q P3c "SELECT count(*) FROM ev_m WHERE created_at >= '2026-09-03' AND created_at < '2026-09-04'"
q P3d "SELECT count(*) FROM ev_m WHERE to_char(created_at,'YYYY-MM') = '2026-09'"

echo; echo "===== P4 查询不带分区键：每个分区都要查一遍"
q P4a "SELECT id, created_at FROM ev_flat WHERE user_id = 4242 ORDER BY created_at DESC LIMIT 20"
q P4b "SELECT id, created_at FROM ev_m    WHERE user_id = 4242 ORDER BY created_at DESC LIMIT 20"
q P4c "SELECT id, created_at FROM ev_d    WHERE user_id = 4242 ORDER BY created_at DESC LIMIT 20"
q P4d "SELECT id, created_at FROM ev_d    WHERE user_id = 4242 AND created_at >= '2026-10-01' ORDER BY created_at DESC LIMIT 20"

echo; echo "===== P5 预编译语句：custom plan vs generic plan"
for mode in force_custom_plan force_generic_plan; do
  echo; echo "### [P5 $mode]"
  $P <<SQL | cut -c1-200
SET plan_cache_mode = $mode;
PREPARE s(timestamp, timestamp) AS SELECT count(*) FROM ev_d WHERE created_at >= \$1 AND created_at < \$2;
EXPLAIN (ANALYZE, BUFFERS) EXECUTE s('2026-09-01','2026-09-08');
EXPLAIN (ANALYZE) EXECUTE s('2026-09-01','2026-09-08');
SQL
done

echo; echo "===== P6 主键 / 唯一约束必须带分区键"
$P -c "ALTER TABLE ev_m ADD PRIMARY KEY (id)"
$P -c "\\timing on" -c "ALTER TABLE ev_m ADD PRIMARY KEY (id, created_at)"

echo; echo "===== P7 默认分区的坑：先落进 DEFAULT 的数据挡住新分区"
$P <<'SQL'
\timing on
DROP TABLE IF EXISTS ev_def;
CREATE TABLE ev_def (LIKE ev_m) PARTITION BY RANGE (created_at);
CREATE TABLE ev_def_202610 PARTITION OF ev_def FOR VALUES FROM ('2026-10-01') TO ('2026-11-01');
CREATE TABLE ev_def_default PARTITION OF ev_def DEFAULT;
-- 11 月分区忘了建，11 月数据先进了 DEFAULT（另造 200,000 行无关的老数据也在 DEFAULT）
INSERT INTO ev_def SELECT id, user_id, event_type, timestamp '2026-11-01' + (id % 86400) * interval '1 second', payload FROM ev_flat WHERE id <= 3;
INSERT INTO ev_def SELECT id, user_id, event_type, created_at, payload FROM ev_flat WHERE created_at < '2025-03-01' LIMIT 200000;
SELECT tableoid::regclass, count(*) FROM ev_def GROUP BY 1 ORDER BY 1;
CREATE TABLE ev_def_202611 PARTITION OF ev_def FOR VALUES FROM ('2026-11-01') TO ('2026-12-01');
-- 把 3 行挪走后再建：要扫 DEFAULT 验证
DELETE FROM ev_def_default WHERE created_at >= '2026-11-01';
CREATE TABLE ev_def_202611 PARTITION OF ev_def FOR VALUES FROM ('2026-11-01') TO ('2026-12-01');
SQL

echo; echo "===== P8 ATTACH PARTITION：先建 CHECK 约束可以跳过全表验证"
$P <<'SQL'
\timing on
DROP TABLE IF EXISTS ev_att;
CREATE TABLE ev_att (LIKE ev_m) PARTITION BY RANGE (created_at);
CREATE TABLE ev_att_a (LIKE ev_m); INSERT INTO ev_att_a SELECT * FROM ev_m_202609;
CREATE TABLE ev_att_b (LIKE ev_m); INSERT INTO ev_att_b SELECT * FROM ev_m_202609;
SELECT count(*) FROM ev_att_a;
-- a：直接 ATTACH（要全表扫描验证分区约束）
ALTER TABLE ev_att ATTACH PARTITION ev_att_a FOR VALUES FROM ('2026-09-01') TO ('2026-10-01');
ALTER TABLE ev_att DETACH PARTITION ev_att_a;
-- b：先加 CHECK（验证扫描在这一步做，此时表还没挂到分区树上），ATTACH 时跳过扫描
ALTER TABLE ev_att_b ADD CONSTRAINT ck CHECK (created_at >= '2026-09-01' AND created_at < '2026-10-01');
ALTER TABLE ev_att ATTACH PARTITION ev_att_b FOR VALUES FROM ('2026-09-01') TO ('2026-10-01');
SQL

echo; echo "===== P9 删一个月的老数据：DELETE vs DETACH + DROP"
$P <<'SQL'
\timing on
SELECT count(*) AS rows_2025_01 FROM ev_flat WHERE created_at < '2025-02-01';
SELECT pg_size_pretty(pg_relation_size('ev_flat')) AS flat_before;
SELECT pg_current_wal_lsn() AS lsn0 \gset
DELETE FROM ev_flat WHERE created_at < '2025-02-01';
SELECT pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_lsn(), :'lsn0')) AS delete_wal;
SELECT pg_size_pretty(pg_relation_size('ev_flat')) AS flat_after_delete;
VACUUM ev_flat;
SELECT pg_size_pretty(pg_relation_size('ev_flat')) AS flat_after_vacuum;
SELECT count(*) AS rows_ev_m_202501 FROM ev_m_202501;
SELECT pg_current_wal_lsn() AS lsn1 \gset
ALTER TABLE ev_m DETACH PARTITION ev_m_202501 CONCURRENTLY;
DROP TABLE ev_m_202501;
SELECT pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_lsn(), :'lsn1')) AS detach_drop_wal;
ALTER TABLE ev_d DETACH PARTITION ev_d_20250101;
DROP TABLE ev_d_20250101;
SQL
echo "[$(ts)] 全部完成"