--
-- partprune_fuzz.sql
--
-- A randomized (fuzz) test suite for RANGE partition pruning correctness.
--
-- Motivation: commit ee58da4 ("Fix another issue with RANGE's DEFAULT
-- partition pruning") fixed a case where the DEFAULT partition was pruned by
-- mistake because an adjacent partition had a MINVALUE / MAXVALUE range bound
-- and the pruning code assumed the leading key columns were all matched with
-- equality when they were not.  e.g.:
--
--   p:  PARTITION BY RANGE (a, b);
--   p1: FOR VALUES FROM (10, 0) TO (10, MAXVALUE);
--   pd: DEFAULT
--   SELECT * FROM p WHERE a = 11 AND b > 10;   <-- pd must be scanned
--
-- Strategy: build random RANGE partitioned tables with 1..N partition key
-- columns, random bounds (including MINVALUE / MAXVALUE and gaps in the range
-- space), a DEFAULT partition and NULLs.  Populate with a fixed value domain,
-- then run random WHERE clauses twice: once with enable_partition_pruning=off
-- (the ground truth) and once with it on.  Any difference in COUNT(*) means a
-- partition holding matching rows was wrongly pruned (or, less likely, an
-- extra one was scanned and produced dup rows).
--
-- Usage (psql, with autocommit on - do NOT wrap this in BEGIN/COMMIT, the
-- driver is a procedure that commits per iteration):
--
--     \i partprune_fuzz.sql
--     SELECT * FROM pp_fuzz.regression_cases();  -- deterministic smoke test
--     CALL pp_fuzz.run(2000);                    -- 2000 random tables
--     CALL pp_fuzz.run(500, 4, 12345);           -- up to 4 key cols, fixed seed
--     SELECT * FROM pp_fuzz.report();            -- results of the last run
--
-- pp_fuzz.report() returns exactly one row on a clean run:
--     'OK: ... iterations, no pruning mismatches'
-- Any other row is a failure report containing the full DDL, the query, both
-- counts and the plan needed to reproduce it standalone.
--
-- Note: pp_fuzz.run() commits after every iteration.  Each iteration creates
-- and drops a partitioned table and its partitions, and those DDL locks are
-- only released at commit, so running the whole thing in one transaction
-- exhausts the backend lock table ("out of shared memory ... you might need to
-- increase max_locks_per_transaction").  Committing per iteration bounds the
-- number of locks held at any time to a single table's partition count.
--

DROP SCHEMA IF EXISTS pp_fuzz CASCADE;
CREATE SCHEMA pp_fuzz;

-- Sentinels used inside int[] bound representations.  Chosen well outside the
-- value domain so plain array comparison gives the same ordering semantics as
-- the partition bound ordering (MINVALUE < any value < MAXVALUE).
CREATE FUNCTION pp_fuzz.minvalue() RETURNS int LANGUAGE sql IMMUTABLE AS
$$ SELECT -1000000 $$;
CREATE FUNCTION pp_fuzz.maxvalue() RETURNS int LANGUAGE sql IMMUTABLE AS
$$ SELECT 1000000 $$;

-- Working table holding the candidate bound tuples for the current iteration.
CREATE TABLE pp_fuzz.bounds (b int[]);

--
-- Render an int[] bound tuple as partition bound SQL text, e.g. '10, MAXVALUE'
--
CREATE FUNCTION pp_fuzz.bound_text(b int[]) RETURNS text
LANGUAGE plpgsql IMMUTABLE AS
$$
DECLARE
	parts text[] := '{}';
	e     int;
BEGIN
	FOREACH e IN ARRAY b LOOP
		IF e = pp_fuzz.minvalue() THEN
			parts := parts || '{MINVALUE}';
		ELSIF e = pp_fuzz.maxvalue() THEN
			parts := parts || '{MAXVALUE}';
		ELSE
			parts := parts || e::text;
		END IF;
	END LOOP;
	RETURN array_to_string(parts, ', ');
END;
$$;

--
-- Build one random bound tuple of ncols elements.
--
-- PostgreSQL requires that once MINVALUE or MAXVALUE appears in a bound, every
-- following column uses the same sentinel.  We honour that here, which also
-- means the canonical int[] form sorts correctly with plain array comparison.
--
CREATE FUNCTION pp_fuzz.random_bound(ncols int, maxval int) RETURNS int[]
LANGUAGE plpgsql VOLATILE AS
$$
DECLARE
	b int[] := '{}';
	i int;
	r double precision;
BEGIN
	FOR i IN 1 .. ncols LOOP
		r := random();
		-- Bias sentinels towards trailing columns: that is exactly the shape
		-- that triggered the DEFAULT-pruning bugs.
		IF i > 1 AND r < 0.18 THEN
			RETURN b || array_fill(pp_fuzz.minvalue(), ARRAY[ncols - i + 1]);
		ELSIF i > 1 AND r < 0.36 THEN
			RETURN b || array_fill(pp_fuzz.maxvalue(), ARRAY[ncols - i + 1]);
		ELSIF i = 1 AND r < 0.06 THEN
			RETURN array_fill(pp_fuzz.minvalue(), ARRAY[ncols]);
		ELSIF i = 1 AND r < 0.12 THEN
			RETURN array_fill(pp_fuzz.maxvalue(), ARRAY[ncols]);
		ELSE
			b := b || (floor(random() * (maxval + 1))::int);
		END IF;
	END LOOP;
	RETURN b;
END;
$$;

--
-- Build a random WHERE clause over columns c1..cncols.
--
CREATE FUNCTION pp_fuzz.random_qual(ncols int, maxval int) RETURNS text
LANGUAGE plpgsql VOLATILE AS
$$
DECLARE
	ops   text[] := ARRAY['=', '<>', '<', '<=', '>', '>='];
	preds text[] := '{}';
	col   text;
	op    text;
	val   int;
	i     int;
	r     double precision;
	p     text;
BEGIN
	FOR i IN 1 .. ncols LOOP
		-- leading columns get a predicate more often; that is where the
		-- interesting equality-prefix logic lives
		IF random() > (0.75 - 0.1 * i) THEN
			col := format('c%s', i);
			r   := random();
			IF r < 0.05 THEN
				preds := preds || format('%s IS NULL', col);
				CONTINUE;
			ELSIF r < 0.10 THEN
				preds := preds || format('%s IS NOT NULL', col);
				CONTINUE;
			END IF;

			op  := ops[1 + floor(random() * array_length(ops, 1))::int];
			-- sample slightly outside the bound domain on purpose
			val := (floor(random() * (maxval + 3))::int) - 1;
			p   := format('%s %s %s', col, op, val);

			-- occasionally OR in a second predicate on the same column
			IF random() < 0.2 THEN
				p := format('(%s OR %s %s %s)', p, col,
							ops[1 + floor(random() * array_length(ops, 1))::int],
							(floor(random() * (maxval + 3))::int) - 1);
			END IF;
			preds := preds || p;
		END IF;
	END LOOP;

	IF array_length(preds, 1) IS NULL THEN
		RETURN 'true';
	END IF;
	RETURN array_to_string(preds, ' AND ');
END;
$$;

--
-- Count rows with pruning forced off, then on, and report any mismatch.
-- Returns NULL when the counts agree.
--
CREATE OR REPLACE FUNCTION pp_fuzz.check_query(tbl text, qual text, ddl text) RETURNS text
LANGUAGE plpgsql AS
$$
DECLARE
	sql     text;
	c_off   bigint;
	c_on    bigint;
	plan    text := '';
	r       record;
BEGIN

	sql := format('SELECT count(*) FROM %I WHERE %s', tbl, qual);

	PERFORM set_config('enable_partition_pruning', 'off', false);
	EXECUTE sql INTO c_off;

	PERFORM set_config('enable_partition_pruning', 'on', false);
	EXECUTE sql INTO c_on;

	INSERT INTO pp_fuzz.debug_output (line) VALUES(ddl);
	INSERT INTO pp_fuzz.debug_output (line) VALUES(sql);
	
	FOR r IN EXECUTE format('EXPLAIN (COSTS OFF) %s', sql) LOOP
		INSERT INTO pp_fuzz.debug_output (line) VALUES(r."QUERY PLAN");
	END LOOP;

	IF c_off IS DISTINCT FROM c_on THEN
		-- capture the pruned plan so the failure report is self-contained
		FOR r IN EXECUTE format('EXPLAIN (COSTS OFF) %s', sql) LOOP
			plan := plan || r."QUERY PLAN" || E'\n';
		END LOOP;

		RETURN format(
			E'FAIL: partition pruning changed the result\n%s\n\nQuery: %s\n  count with enable_partition_pruning = off : %s\n  count with enable_partition_pruning = on  : %s\n\nPlan (pruning on):\n%s',
			ddl, sql, c_off, c_on, plan);
	END IF;

	RETURN NULL;
END;
$$;

--
-- Failure reports land here rather than being returned as a result set: the
-- driver has to be a PROCEDURE so it can COMMIT (see below), and procedures
-- cannot return a set.
--
CREATE TABLE pp_fuzz.results (
	id        serial primary key,
	ts        timestamptz NOT NULL DEFAULT now(),
	iteration int,
	report    text
);

-- Debug output for recording plans and ddl.
CREATE TABLE pp_fuzz.debug_output (
	id SERIAL PRIMARY KEY,
	line TEXT NOT NULL
);

--
-- Main driver.
--
--   iterations - number of random tables to build (each gets several queries)
--   max_cols   - maximum number of RANGE partition key columns (default 3)
--   seed       - optional seed for reproducible runs
--
-- This is a PROCEDURE, not a function, so that it can COMMIT after every
-- iteration.  Each iteration creates and drops a partitioned table plus its
-- partitions, and every one of those takes an AccessExclusiveLock that is only
-- released at commit.  Running thousands of iterations inside a single
-- transaction exhausts the backend's lock table
-- ("out of shared memory / You might need to increase max_locks_per_transaction").
-- Committing per iteration keeps the number of held locks bounded by the
-- partition count of a single table.
--
-- Must be called with CALL and NOT inside an explicit transaction block:
--     CALL pp_fuzz.run(2000);
--     SELECT * FROM pp_fuzz.report();
--
CREATE PROCEDURE pp_fuzz.run(iterations int DEFAULT 500,
							 max_cols   int DEFAULT 3,
							 seed       double precision DEFAULT NULL)
LANGUAGE plpgsql AS
$$
DECLARE
	maxval    CONSTANT int  := 10;             -- bound value domain 0..10
	vals      CONSTANT int[] := ARRAY[0, 1, 2, 5, 9, 10];  -- inserted values
	tbl       CONSTANT text := 'pp_fuzz_t';
	ncols     int;
	nbounds   int;
	i         int;
	j         int;
	pno       int;
	coldefs   text;
	keycols   text;
	ddl       text;
	part_sql  text;
	rec       record;
	nparts    int;
	has_dflt  bool;
	qual      text;
	res       text;
	failures  int := 0;
BEGIN
	IF seed IS NOT NULL THEN
		PERFORM setseed(seed - floor(seed));
	END IF;

	DELETE FROM pp_fuzz.results;
	COMMIT;

	EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', tbl);

	FOR i IN 1 .. iterations LOOP
		ncols := 1 + floor(random() * max_cols)::int;

		SELECT string_agg(format('c%s int', k), ', ' ORDER BY k),
			   string_agg(format('c%s', k),     ', ' ORDER BY k)
		  INTO coldefs, keycols
		  FROM generate_series(1, ncols) k;

		ddl := format('CREATE TABLE %I (%s) PARTITION BY RANGE (%s);',
					  tbl, coldefs, keycols);
		EXECUTE ddl;

		----------------------------------------------------------------
		-- Random, strictly increasing bound tuples -> valid partitions
		----------------------------------------------------------------
		DELETE FROM pp_fuzz.bounds;   -- DELETE, not TRUNCATE: avoids an extra
									  -- AccessExclusiveLock + relfilenode churn
		nbounds := 2 + floor(random() * 6)::int;
		FOR j IN 1 .. nbounds LOOP
			INSERT INTO pp_fuzz.bounds VALUES (pp_fuzz.random_bound(ncols, maxval));
		END LOOP;

		pno    := 0;
		nparts := 0;
		FOR rec IN
			SELECT lo, hi FROM (
				SELECT lag(b) OVER (ORDER BY b) AS lo, b AS hi
				  FROM (SELECT DISTINCT b FROM pp_fuzz.bounds) d
			) s WHERE lo IS NOT NULL
		LOOP
			pno := pno + 1;
			-- deliberately leave holes in the range space so the DEFAULT
			-- partition has to be consulted for values in the gaps
			CONTINUE WHEN random() < 0.25;

			part_sql := format(
				'CREATE TABLE %I PARTITION OF %I FOR VALUES FROM (%s) TO (%s);',
				tbl || '_p' || pno, tbl,
				pp_fuzz.bound_text(rec.lo), pp_fuzz.bound_text(rec.hi));
			EXECUTE part_sql;
			ddl    := ddl || E'\n' || part_sql;
			nparts := nparts + 1;
		END LOOP;

		-- DEFAULT partition most of the time (it is the usual victim)
		has_dflt := random() < 0.85;
		IF has_dflt THEN
			part_sql := format('CREATE TABLE %I PARTITION OF %I DEFAULT;',
							   tbl || '_pd', tbl);
			EXECUTE part_sql;
			ddl    := ddl || E'\n' || part_sql;
			nparts := nparts + 1;
		END IF;

		IF nparts = 0 THEN
			-- nothing to prune; discard and try again (commit to free locks)
			EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', tbl);
			COMMIT;
			CONTINUE;
		END IF;

		----------------------------------------------------------------
		-- Populate: cross product of the value domain (+ NULLs, which land
		-- in the DEFAULT partition and must never be pruned away by an
		-- IS NULL / non-strict qual).
		----------------------------------------------------------------
		DECLARE
			sel  text := '';
			frm  text := '';
			dom  text;
		BEGIN
			dom := (SELECT string_agg(format('(%s)', v), ',')
					  FROM unnest(vals) v)
				   || ',(NULL)';
			FOR j IN 1 .. ncols LOOP
				sel := sel || CASE WHEN j > 1 THEN ', ' ELSE '' END
						   || format('v%s.x', j);
				frm := frm || CASE WHEN j > 1 THEN ' CROSS JOIN ' ELSE '' END
						   || format('(VALUES %s) AS v%s(x)', dom, j);
			END LOOP;

			BEGIN
				EXECUTE format('INSERT INTO %I SELECT %s FROM %s', tbl, sel, frm);
			EXCEPTION WHEN check_violation THEN
				-- No DEFAULT partition, so some values have no home.  Fall
				-- back to inserting row by row and skipping the ones that
				-- do not fit into any partition.
				DECLARE
					r    record;
					k    int;
					lits text;
				BEGIN
					FOR r IN EXECUTE format('SELECT ARRAY[%s] AS a FROM %s',
											sel, frm)
					LOOP
						lits := '';
						FOR k IN 1 .. ncols LOOP
							lits := lits || CASE WHEN k > 1 THEN ', ' ELSE '' END
										 || coalesce((r.a)[k]::text, 'NULL');
						END LOOP;
						BEGIN
							EXECUTE format('INSERT INTO %I VALUES (%s)',
										   tbl, lits);
						EXCEPTION WHEN check_violation THEN
							NULL;   -- expected: value falls in a gap
						END;
					END LOOP;
				END;
			END;
		END;

		EXECUTE format('ANALYZE %I', tbl);

		----------------------------------------------------------------
		-- Query the table several times, comparing pruning on vs off
		----------------------------------------------------------------
		FOR j IN 1 .. 8 LOOP
			qual := pp_fuzz.random_qual(ncols, maxval);
			res  := pp_fuzz.check_query(tbl, qual, ddl);
			IF res IS NOT NULL THEN
				failures := failures + 1;
				INSERT INTO pp_fuzz.results (iteration, report) VALUES (i, res);
			END IF;
		END LOOP;

		-- Drop the table and COMMIT before the next iteration so that the
		-- AccessExclusiveLocks taken by CREATE/DROP TABLE are released.
		-- Without this the backend runs out of lock slots after a few hundred
		-- iterations.
		EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', tbl);
		COMMIT;
	END LOOP;

	EXECUTE format('DROP TABLE IF EXISTS %I CASCADE', tbl);
	PERFORM set_config('enable_partition_pruning', 'on', false);

	IF failures = 0 THEN
		INSERT INTO pp_fuzz.results (iteration, report)
		VALUES (iterations,
				format('OK: %s iterations, no pruning mismatches', iterations));
	ELSE
		INSERT INTO pp_fuzz.results (iteration, report)
		VALUES (iterations,
				format('FAILED: %s mismatching queries over %s iterations',
					   failures, iterations));
	END IF;
	COMMIT;

	RAISE NOTICE 'pp_fuzz.run: % iterations, % mismatching queries. See SELECT * FROM pp_fuzz.report();',
				 iterations, failures;
END;
$$;

--
-- Convenience reader for the results collected by pp_fuzz.run().
--
CREATE FUNCTION pp_fuzz.report() RETURNS SETOF text
LANGUAGE sql AS
$$
	SELECT report FROM pp_fuzz.results ORDER BY id
$$;

--
-- Deterministic cases: the shapes that were actually broken before 709dfd27f
-- and ee58da4.  These are cheap and make a good smoke test.
--
CREATE FUNCTION pp_fuzz.regression_cases() RETURNS SETOF text
LANGUAGE plpgsql AS
$$
DECLARE
	ddl   text;
	quals text[];
	q     text;
	res   text;
	fails int := 0;
BEGIN
	-- Case 1: trailing MAXVALUE, lookup key does not match the leading column
	DROP TABLE IF EXISTS pp_fuzz_t CASCADE;
	ddl := 'CREATE TABLE pp_fuzz_t (c1 int, c2 int) PARTITION BY RANGE (c1, c2);'
		|| E'\nCREATE TABLE pp_fuzz_t_p1 PARTITION OF pp_fuzz_t FOR VALUES FROM (10, 0) TO (10, MAXVALUE);'
		|| E'\nCREATE TABLE pp_fuzz_t_pd PARTITION OF pp_fuzz_t DEFAULT;';
	EXECUTE 'CREATE TABLE pp_fuzz_t (c1 int, c2 int) PARTITION BY RANGE (c1, c2)';
	EXECUTE 'CREATE TABLE pp_fuzz_t_p1 PARTITION OF pp_fuzz_t FOR VALUES FROM (10, 0) TO (10, MAXVALUE)';
	EXECUTE 'CREATE TABLE pp_fuzz_t_pd PARTITION OF pp_fuzz_t DEFAULT';
	EXECUTE 'INSERT INTO pp_fuzz_t SELECT a.x, b.x FROM
			 (VALUES (9),(10),(11),(NULL)) a(x) CROSS JOIN
			 (VALUES (0),(10),(20),(NULL)) b(x)';
	ANALYZE pp_fuzz_t;

	quals := ARRAY[
		'c1 = 11 AND c2 > 10',
		'c1 = 11',
		'c1 > 10 AND c2 > 10',
		'c1 = 10 AND c2 < 0',
		'c1 = 9  AND c2 >= 0',
		'c1 IS NULL',
		'c2 IS NULL'
	];
	FOREACH q IN ARRAY quals LOOP
		res := pp_fuzz.check_query('pp_fuzz_t', q, ddl);
		IF res IS NOT NULL THEN fails := fails + 1; RETURN NEXT res; END IF;
	END LOOP;

	-- Case 2: trailing MINVALUE variant
	DROP TABLE pp_fuzz_t CASCADE;
	ddl := 'CREATE TABLE pp_fuzz_t (c1 int, c2 int, c3 int) PARTITION BY RANGE (c1, c2, c3);'
		|| E'\nCREATE TABLE pp_fuzz_t_p1 PARTITION OF pp_fuzz_t FOR VALUES FROM (10, MINVALUE, MINVALUE) TO (10, 5, 0);'
		|| E'\nCREATE TABLE pp_fuzz_t_pd PARTITION OF pp_fuzz_t DEFAULT;';
	EXECUTE 'CREATE TABLE pp_fuzz_t (c1 int, c2 int, c3 int) PARTITION BY RANGE (c1, c2, c3)';
	EXECUTE 'CREATE TABLE pp_fuzz_t_p1 PARTITION OF pp_fuzz_t FOR VALUES FROM (10, MINVALUE, MINVALUE) TO (10, 5, 0)';
	EXECUTE 'CREATE TABLE pp_fuzz_t_pd PARTITION OF pp_fuzz_t DEFAULT';
	EXECUTE 'INSERT INTO pp_fuzz_t SELECT a.x, b.x, c.x FROM
			 (VALUES (9),(10),(11),(NULL)) a(x) CROSS JOIN
			 (VALUES (0),(5),(9),(NULL)) b(x) CROSS JOIN
			 (VALUES (0),(5),(NULL)) c(x)';
	ANALYZE pp_fuzz_t;

	quals := ARRAY[
		'c1 = 9 AND c2 < 5',
		'c1 = 10 AND c2 < 5',
		'c1 = 10 AND c2 = 5 AND c3 >= 0',
		'c1 <= 10 AND c2 < 5',
		'c1 = 11 AND c2 < 5 AND c3 < 0'
	];
	FOREACH q IN ARRAY quals LOOP
		res := pp_fuzz.check_query('pp_fuzz_t', q, ddl);
		IF res IS NOT NULL THEN fails := fails + 1; RETURN NEXT res; END IF;
	END LOOP;

	DROP TABLE pp_fuzz_t CASCADE;
	PERFORM set_config('enable_partition_pruning', 'on', false);

	IF fails = 0 THEN
		RETURN NEXT 'OK: regression cases passed';
	ELSE
		RETURN NEXT format('FAILED: %s regression cases mismatched', fails);
	END IF;
END;
$$;
