-- The contracts every custom node of pg_lion keeps with PostgreSQL, run
-- against each node from one list, so a new node gets a line here (as
-- reader_contracts.sql does for the direct readers).  A node stands in for
-- expressions the executor would have initialised - a clause's operator it
-- answers from a posting set, an aggregate it computes, a score it returns -
-- and so asks EXECUTE on each of them itself (DESIGN.md §9, "Privileges"):
--   - of the current user, when the plan starts, plain EXPLAIN included;
--   - whichever role made the plan: a cached plan answers a REVOKE and a SET
--     ROLE as the ordinary plan does;
--   - of the owner inside a SECURITY DEFINER function;
-- and with the privilege its answer is the ordinary plan's.  The hook a
-- security module sees for the same check (InvokeFunctionExecuteHook) needs a
-- module that installs one: test/modules/lion_hooktest/sql/exec_hook.sql runs
-- this list against it.  security_exec.sql covers each node's own functions
-- in more detail.
--
-- Revoking a function from PUBLIC affects the whole database, so nc_check()
-- grants it back before it returns, and the file ends by checking that
-- nothing is left revoked.
\set VERBOSITY terse
SET client_min_messages = warning;
CREATE EXTENSION IF NOT EXISTS pg_lion;
RESET client_min_messages;
SET synchronous_commit = on;
SET max_parallel_workers_per_gather = 0;
CREATE ROLE lion_nc_yes;
CREATE ROLE lion_nc_no;

-- o: a permutation of the ids, with a covering btree
CREATE TABLE nc (id int, k int, j int, o int, d tsvector);
INSERT INTO nc SELECT i, i % 10, i % 7, (i * 7919) % 3001,
	   to_tsvector('simple', 'w' || (i % 5) || ' w' || (i % 7) || ' w' || (i % 11))
  FROM generate_series(1, 3000) i;
CREATE INDEX nc_k ON nc USING lion (k);
CREATE INDEX nc_j ON nc USING lion (j);
CREATE INDEX nc_d ON nc USING lion (d) WITH (store_positions = true);
CREATE INDEX nc_o ON nc (o) INCLUDE (id);
-- an expression column (DESIGN.md §41) over a function nothing inlines
CREATE FUNCTION nc_bucket(int) RETURNS int IMMUTABLE LANGUAGE plpgsql
	AS 'BEGIN RETURN $1 % 4; END';
CREATE INDEX nc_b ON nc USING lion (nc_bucket(j));
-- a dimension on nc.k, three of whose keys have no fact row
CREATE TABLE ncd (pk int PRIMARY KEY, name text NOT NULL);
INSERT INTO ncd SELECT i, 'n' || (i % 2) FROM generate_series(0, 12) i;
-- a table on nc.j that only filters, its key of another type
CREATE TABLE ncx (pk int8 PRIMARY KEY, flag bool NOT NULL);
INSERT INTO ncx SELECT i, i % 3 = 0 FROM generate_series(0, 6) i;
-- the joins by key (DESIGN.md §46-§48): a dimension of lion columns under a
-- primary key, and a fact of few rows without a lion index, its key of
-- another type, so that the join's operator is a function of its own; a
-- table of keys that filters the fact's key, and one for a NOT EXISTS
CREATE TABLE ncpd (pk int PRIMARY KEY, r int NOT NULL, t text NOT NULL);
INSERT INTO ncpd SELECT i, i % 5, 't' || (i % 3) FROM generate_series(1, 2000) i;
CREATE INDEX ncpd_r ON ncpd USING lion (r);
CREATE INDEX ncpd_t ON ncpd USING lion (t);
CREATE TABLE ncpf (id int PRIMARY KEY, fk int8, ch int NOT NULL, rk int NOT NULL);
INSERT INTO ncpf SELECT i, 1 + (i * 37) % 2000, i % 3, (i * 7919) % 30001
  FROM generate_series(1, 30) i;
CREATE INDEX ncpf_rk ON ncpf (rk);
CREATE TABLE ncpk (fk int8 NOT NULL, prod int NOT NULL);
INSERT INTO ncpk SELECT 1 + (i * 37) % 2000, i % 2 FROM generate_series(1, 30) i;
CREATE TABLE ncpb (k int8);
INSERT INTO ncpb SELECT 1 + (i * 37 * 2) % 2000 FROM generate_series(1, 10) i;
VACUUM (FREEZE, ANALYZE) nc;
VACUUM (FREEZE, ANALYZE) ncd;
VACUUM (FREEZE, ANALYZE) ncx;
VACUUM (FREEZE, ANALYZE) ncpd;
VACUUM (FREEZE, ANALYZE) ncpf;
VACUUM (FREEZE, ANALYZE) ncpk;
VACUUM (FREEZE, ANALYZE) ncpb;
GRANT SELECT ON nc, ncd, ncx, ncpd, ncpf, ncpk, ncpb TO lion_nc_yes, lion_nc_no;

-- every custom node, with a query it answers, a function it stands in for,
-- and the settings that make the planner choose it; a new node goes here
CREATE TABLE nc_nodes (node text PRIMARY KEY, q text, fn regprocedure, force text);
INSERT INTO nc_nodes VALUES
	('LionCount',
	 'SELECT count(*) FROM nc WHERE k = 3', 'int4eq(int4,int4)', ''),
	('LionCount, its aggregate',
	 'SELECT count(*) FROM nc WHERE k = 3', 'count()', ''),
	('LionCount, an expression over its count',
	 'SELECT k, round(100.0 * count(*) / 3000, 2) FROM nc GROUP BY k', 'round(numeric,int4)', ''),
	('LionCount, a NOT IN list',
	 'SELECT count(*) FROM nc WHERE k NOT IN (3, 4)', 'int4ne(int4,int4)', ''),
	('LionCount, an IS DISTINCT FROM',
	 'SELECT count(*) FROM nc WHERE k IS DISTINCT FROM 3', 'int4eq(int4,int4)', ''),
	('LionCount, a complement in an OR',
	 'SELECT count(*) FROM nc WHERE j = 2 OR k <> 3', 'int4ne(int4,int4)', ''),
	('LionCount, a NOT of a multi-key query',
	 'SELECT count(*) FROM nc WHERE NOT (d @@ ''w1 & w3'')', 'ts_match_vq(tsvector,tsquery)', ''),
	('LionCount, a residual clause',
	 'SELECT count(*) FROM nc WHERE k = 3 AND o % 7 = 1', 'int4mod(int4,int4)', ''),
	('LionCount, a FILTER on its count',
	 'SELECT count(*) FILTER (WHERE j <> 2), count(*) FROM nc WHERE k = 3', 'int4ne(int4,int4)', ''),
	('LionCount, a conditional sum',
	 'SELECT k, sum(CASE WHEN j <> 2 THEN 1 ELSE 0 END) FROM nc GROUP BY k', 'int4ne(int4,int4)', ''),
	('LionCount, a residual clause beside a FILTER',
	 'SELECT count(*) FILTER (WHERE j = 2), count(*) FROM nc WHERE k = 3 AND o % 7 = 1', 'int4mod(int4,int4)',
	 'enable_seqscan enable_bitmapscan pg_lion.enable_plain_scan'),
	('LionCount, an expression column filtered',
	 'SELECT count(*) FROM nc WHERE nc_bucket(j) = 1', 'nc_bucket(int4)', ''),
	('LionCount, an expression column grouped',
	 'SELECT nc_bucket(j), count(*) FROM nc GROUP BY 1', 'nc_bucket(int4)', ''),
	('LionCount, a dimension folded into a fact filter',
	 'SELECT d.name, count(*) FROM nc JOIN ncd d ON d.pk = nc.k JOIN ncx x ON x.pk = nc.j WHERE x.flag GROUP BY d.name',
	 'int48eq(int4,int8)', 'enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionJoinAgg',
	 'SELECT count(DISTINCT d.name), count(*) FROM nc JOIN ncd d ON d.pk = nc.k WHERE nc.j = 3',
	 'int4eq(int4,int4)', ''),
	('LionJoinAgg, a LEFT JOIN',
	 'SELECT d.name, count(*) FROM ncd d LEFT JOIN nc ON nc.k = d.pk AND nc.j = 3 GROUP BY d.name',
	 'int4eq(int4,int4)', 'enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionJoinAgg, grouping sets',
	 'SELECT k, j, sum(j), count(*) FROM nc GROUP BY ROLLUP (k, j)', 'sum(int4)',
	 'enable_seqscan enable_bitmapscan pg_lion.enable_plain_scan'),
	('LionSemiJoin',
	 'SELECT d.pk FROM ncd d WHERE EXISTS (SELECT 1 FROM nc WHERE nc.k = d.pk AND nc.j = 3)',
	 'int4eq(int4,int4)', 'enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionAntiJoin',
	 'SELECT d.pk FROM ncd d WHERE NOT EXISTS (SELECT 1 FROM nc WHERE nc.k = d.pk AND nc.j = 3)',
	 'int4eq(int4,int4)', 'enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionOrdered',
	 'SELECT id FROM nc WHERE k = 3 ORDER BY o LIMIT 5', 'int4eq(int4,int4)',
	 'enable_sort enable_seqscan enable_bitmapscan pg_lion.enable_plain_scan'),
	('LionBtreeScan',
	 'SELECT id, o FROM nc WHERE o BETWEEN 100 AND 400 AND k = 3', 'int4eq(int4,int4)',
	 'enable_seqscan enable_bitmapscan pg_lion.enable_plain_scan'),
	('LionBm25',
	 'SELECT id, lion_bm25_score(d, ''w1 & w3'', ''nc_d'') FROM nc WHERE d @@ ''w1 & w3''
	   ORDER BY lion_bm25_score(d, ''w1 & w3'', ''nc_d'') DESC LIMIT 5',
	 'lion_bm25_score(tsvector,tsquery,regclass,float8,float8,name)',
	 'enable_seqscan enable_bitmapscan pg_lion.enable_plain_scan'),
	('LionBm25, a named column',
	 'SELECT id, lion_bm25_score(d, ''w1 & w3'', ''nc_d'', col => ''d'') FROM nc WHERE d @@ ''w1 & w3''
	   ORDER BY lion_bm25_score(d, ''w1 & w3'', ''nc_d'', col => ''d'') DESC LIMIT 5',
	 'lion_bm25_score(tsvector,tsquery,regclass,float8,float8,name)',
	 'enable_seqscan enable_bitmapscan pg_lion.enable_plain_scan'),
	('LionProbeJoin',
	 'SELECT f.ch, count(*) FROM ncpf f JOIN ncpd d ON f.fk = d.pk WHERE d.t = ''t1'' GROUP BY 1',
	 'int48eq(int4,int8)', 'enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionProbeJoin, its WHERE',
	 'SELECT f.ch, count(*) FROM ncpf f JOIN ncpd d ON f.fk = d.pk WHERE d.r = 3 GROUP BY 1',
	 'int4eq(int4,int4)', 'enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionProbeJoin, a predicate of its GROUP BY',
	 'SELECT d.r = 2, count(*) FROM ncpf f JOIN ncpd d ON f.fk = d.pk GROUP BY 1',
	 'int4eq(int4,int4)', 'enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionProbeJoin, a NOT EXISTS',
	 'SELECT f.ch, count(*) FROM ncpf f JOIN ncpd d ON f.fk = d.pk
	   WHERE d.t = ''t1'' AND NOT EXISTS (SELECT 1 FROM ncpb b WHERE b.k = f.fk) GROUP BY 1',
	 'int8eq(int8,int8)', 'enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionOrderedJoin',
	 'SELECT f.id, d.pk FROM ncpf f JOIN ncpd d ON f.fk = d.pk WHERE d.r = 3 ORDER BY f.rk, f.id LIMIT 5',
	 'int48eq(int4,int8)', 'enable_sort enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionOrderedJoin, a key set',
	 'SELECT f.id, d.pk FROM ncpf f JOIN ncpd d ON f.fk = d.pk
	   WHERE d.r = 3 AND f.fk IN (SELECT fk FROM ncpk WHERE prod = 1) ORDER BY f.rk, f.id LIMIT 5',
	 'int8eq(int8,int8)', 'enable_sort enable_hashjoin enable_mergejoin enable_nestloop'),
	('LionOrderedSplit',
	 'SELECT f.id, d.pk, d.r FROM ncpf f JOIN ncpd d ON f.fk = d.pk
	   WHERE d.r < 3 AND f.fk IN (SELECT fk FROM ncpk WHERE prod = 1) ORDER BY d.r, f.rk, f.id LIMIT 5',
	 'int4lt(int4,int4)', 'enable_sort enable_hashjoin enable_mergejoin enable_nestloop');

-- lion's paths on, and the core settings in force turned off (for the
-- transaction); or every lion path off and core's settings back
CREATE FUNCTION nc_paths(force text, lion boolean) RETURNS void
LANGUAGE plpgsql AS $$
DECLARE
	g text;
BEGIN
	FOREACH g IN ARRAY ARRAY['pg_lion.enable_count_pushdown', 'pg_lion.enable_semijoin',
							 'pg_lion.enable_ordered_scan', 'pg_lion.enable_btree_scan',
							 'pg_lion.enable_bm25_scan', 'pg_lion.enable_probe_join',
							 'pg_lion.enable_ordered_join'] LOOP
		PERFORM set_config(g, CASE WHEN lion THEN 'on' ELSE 'off' END, true);
	END LOOP;
	FOREACH g IN ARRAY string_to_array(force, ' ') LOOP
		PERFORM set_config(g, CASE WHEN lion THEN 'off' ELSE 'on' END, true);
	END LOOP;
	-- core's index scans stay on, as they disable lion's scans too: lion's
	-- paths are marked down against them instead
	PERFORM set_config('pg_lion.pushdown_margin',
					   CASE WHEN lion AND force LIKE '%enable_plain_scan%' THEN '1000' ELSE '1' END, true);
END $$;

-- q's rows as one sorted string
CREATE FUNCTION nc_rows(q text) RETURNS text
LANGUAGE sql AS $$
	SELECT format('SELECT coalesce(string_agg(s::text, '' '' ORDER BY s::text), ''(none)'') FROM (%s) s', q)
$$;

-- what statement stmt returns, or the privilege error it raised
CREATE FUNCTION nc_try(stmt text) RETURNS text
LANGUAGE plpgsql AS $$
DECLARE
	r text;
BEGIN
	EXECUTE stmt INTO r;
	RETURN r;
EXCEPTION WHEN insufficient_privilege THEN
	RETURN 'ERROR: ' || SQLERRM;
END $$;

-- plain EXPLAIN of q: the node it planned, or the privilege error it raised
CREATE FUNCTION nc_explain(q text, node text) RETURNS text
LANGUAGE plpgsql AS $$
DECLARE
	ln text;
	r text := 'NOT PLANNED';
BEGIN
	FOR ln IN EXECUTE 'EXPLAIN (COSTS OFF) ' || q LOOP
		IF ln LIKE '%Custom Scan (' || split_part(node, ',', 1) || ')%' THEN
			r := 'planned';
		END IF;
	END LOOP;
	RETURN r;
EXCEPTION WHEN insufficient_privilege THEN
	RETURN 'ERROR: ' || SQLERRM;
END $$;

-- q run as its owner
CREATE FUNCTION nc_def_yes(q text) RETURNS text SECURITY DEFINER
LANGUAGE plpgsql AS 'BEGIN RETURN nc_try(nc_rows(q)); END';
CREATE FUNCTION nc_def_no(q text) RETURNS text SECURITY DEFINER
LANGUAGE plpgsql AS 'BEGIN RETURN nc_try(nc_rows(q)); END';
ALTER FUNCTION nc_def_yes(text) OWNER TO lion_nc_yes;
ALTER FUNCTION nc_def_no(text) OWNER TO lion_nc_no;

-- what the node's plan says, against the ordinary plan's: "as the ordinary
-- plan" when they agree
CREATE FUNCTION nc_vs(lion text, ordinary text) RETURNS text
LANGUAGE sql AS $$
	SELECT CASE WHEN lion = ordinary THEN 'as the ordinary plan'
				ELSE lion || ' | the ordinary plan: ' || ordinary END
$$;

/*
 * nc_check() runs one node's query through the node and through the ordinary
 * plan: with every privilege, then with the function revoked from PUBLIC and
 * granted to lion_nc_yes alone - as lion_nc_no, as lion_nc_yes, through a
 * generic plan the superuser made, and through SECURITY DEFINER functions -
 * and grants it back.
 */
CREATE FUNCTION nc_check(name text) RETURNS SETOF text
LANGUAGE plpgsql AS $$
DECLARE
	n nc_nodes;
	answer text;
	a text;
	b text;
	g text;
BEGIN
	SELECT * INTO STRICT n FROM nc_nodes WHERE node = name;
	PERFORM nc_paths(n.force, true);
	answer := nc_try(nc_rows(n.q));
	RETURN NEXT 'plan: ' || nc_explain(n.q, n.node);
	PERFORM nc_paths(n.force, false);
	RETURN NEXT 'every privilege: ' ||
		CASE WHEN answer = nc_try(nc_rows(n.q)) THEN 'the ordinary plan''s answer'
			 ELSE 'MISMATCH: ' || answer || ' | ' || nc_try(nc_rows(n.q)) END;

	EXECUTE format('REVOKE EXECUTE ON FUNCTION %s FROM PUBLIC', n.fn);
	EXECUTE format('GRANT EXECUTE ON FUNCTION %s TO lion_nc_yes', n.fn);

	-- a role without the privilege, and plain EXPLAIN
	PERFORM set_config('role', 'lion_nc_no', true);
	PERFORM nc_paths(n.force, true);
	a := nc_try(nc_rows(n.q));
	PERFORM nc_paths(n.force, false);
	RETURN NEXT 'without EXECUTE on ' || n.fn || ': ' || a || ', ' || nc_vs(a, nc_try(nc_rows(n.q)));
	PERFORM nc_paths(n.force, true);
	a := nc_explain(n.q, n.node);
	PERFORM nc_paths(n.force, false);
	RETURN NEXT 'plain EXPLAIN: ' || nc_vs(a, nc_explain(n.q, n.node));

	-- a role with it
	PERFORM set_config('role', 'lion_nc_yes', true);
	PERFORM nc_paths(n.force, true);
	a := nc_try(nc_rows(n.q));
	RETURN NEXT 'granted to the role: ' ||
		CASE WHEN a = answer THEN 'the answer' ELSE a END;

	-- a generic plan made by the superuser, run by each role
	PERFORM set_config('role', 'none', true);
	PERFORM set_config('plan_cache_mode', 'force_generic_plan', true);
	EXECUTE 'PREPARE nc_p AS ' || nc_rows(n.q);
	a := nc_try('EXECUTE nc_p');
	b := nc_explain('EXECUTE nc_p', n.node);
	PERFORM set_config('role', 'lion_nc_no', true);
	a := nc_try('EXECUTE nc_p');
	PERFORM set_config('role', 'lion_nc_yes', true);
	b := b || ', ' || CASE WHEN nc_try('EXECUTE nc_p') = answer THEN 'the answer'
						   ELSE nc_try('EXECUTE nc_p') END;
	RETURN NEXT 'a cached plan (' || b || ') without it: ' || a;
	PERFORM set_config('role', 'none', true);
	DEALLOCATE nc_p;

	-- SECURITY DEFINER: the owner's privileges, whoever calls
	PERFORM set_config('role', 'lion_nc_no', true);
	a := nc_def_yes(n.q);
	PERFORM set_config('role', 'lion_nc_yes', true);
	b := nc_def_no(n.q);
	RETURN NEXT 'SECURITY DEFINER, owned by the role with it: ' ||
		CASE WHEN a = answer THEN 'the answer' ELSE a END ||
		'; by the role without it: ' || b;

	PERFORM set_config('role', 'none', true);
	EXECUTE format('REVOKE EXECUTE ON FUNCTION %s FROM lion_nc_yes', n.fn);
	EXECUTE format('GRANT EXECUTE ON FUNCTION %s TO PUBLIC', n.fn);

	-- every node's calls share one transaction: the next node's ordinary plan
	-- starts from the settings this one found
	FOREACH g IN ARRAY string_to_array(n.force, ' ') ||
			ARRAY['pg_lion.enable_count_pushdown', 'pg_lion.enable_semijoin',
				  'pg_lion.enable_ordered_scan', 'pg_lion.enable_btree_scan',
				  'pg_lion.enable_bm25_scan', 'pg_lion.enable_probe_join',
				  'pg_lion.enable_ordered_join', 'pg_lion.pushdown_margin',
				  'plan_cache_mode'] LOOP
		EXECUTE 'RESET ' || g;
	END LOOP;
END $$;

SELECT node, nc_check(node) FROM nc_nodes ORDER BY node;

-- nothing is left revoked from PUBLIC by this file
SELECT fn FROM nc_nodes WHERE NOT has_function_privilege('public', fn, 'EXECUTE');

DROP FUNCTION nc_check(text);
DROP FUNCTION nc_vs(text, text);
DROP FUNCTION nc_def_yes(text);
DROP FUNCTION nc_def_no(text);
DROP FUNCTION nc_explain(text, text);
DROP FUNCTION nc_try(text);
DROP FUNCTION nc_rows(text);
DROP FUNCTION nc_paths(text, boolean);
DROP TABLE nc_nodes, nc, ncd, ncx, ncpd, ncpf, ncpk, ncpb;
DROP FUNCTION nc_bucket(int);
DROP ROLE lion_nc_yes;
DROP ROLE lion_nc_no;