--
-- MULTI_PREPARE_PLSQL
--
-- Many of the queries are taken from other regression test files
-- and converted into both plain SQL and PL/pgsql functions, which
-- use prepared statements internally.
-- many of the tests in this file is intended for testing non-fast-path
-- router planner, so we're explicitly disabling it in this file.
-- We've bunch of other tests that triggers fast-path-router
SET citus.enable_fast_path_router_planner TO false;
CREATE FUNCTION plpgsql_test_1() RETURNS TABLE(count bigint) AS $$
DECLARE
BEGIN
     RETURN QUERY
         SELECT
			count(*)
		FROM
			orders;

END;
$$ LANGUAGE plpgsql;
CREATE FUNCTION plpgsql_test_2() RETURNS TABLE(count bigint) AS $$
DECLARE
BEGIN
     RETURN QUERY
		SELECT
			count(*)
		FROM
			orders, lineitem
		WHERE
			o_orderkey = l_orderkey;

END;
$$ LANGUAGE plpgsql;
CREATE FUNCTION plpgsql_test_3() RETURNS TABLE(count bigint) AS $$
DECLARE
BEGIN
     RETURN QUERY
		SELECT
			count(*)
		FROM
			orders, customer
		WHERE
			o_custkey = c_custkey;

END;
$$ LANGUAGE plpgsql;
CREATE FUNCTION plpgsql_test_4() RETURNS TABLE(count bigint) AS $$
DECLARE
BEGIN
     RETURN QUERY
		SELECT
			count(*)
		FROM
			orders, customer, lineitem
		WHERE
			o_custkey = c_custkey AND
			o_orderkey = l_orderkey;
END;
$$ LANGUAGE plpgsql;
CREATE FUNCTION plpgsql_test_5() RETURNS TABLE(count bigint) AS $$
DECLARE
BEGIN
     RETURN QUERY
		SELECT
			count(*)
		FROM
			lineitem, customer
		WHERE
			l_partkey = c_nationkey;
END;
$$ LANGUAGE plpgsql;
CREATE FUNCTION plpgsql_test_6(int) RETURNS TABLE(count bigint) AS $$
DECLARE
BEGIN
     RETURN QUERY
		SELECT
			count(*)
		FROM
			orders, lineitem
		WHERE
			o_orderkey = l_orderkey AND
			l_suppkey > $1;
END;
$$ LANGUAGE plpgsql;
CREATE FUNCTION plpgsql_test_7(text, text) RETURNS TABLE(supp_natadsion text, cusasdt_nation text, l_yeasdar int, sasdaum double precision) AS $$
DECLARE
BEGIN
	RETURN  QUERY
		SELECT
			supp_nation::text,
			cust_nation::text,
			l_year::int,
			sum(volume)::double precision AS revenue
		FROM
			(
			SELECT
				supp_nation,
				cust_nation,
				extract(year FROM l_shipdate) AS l_year,
				l_extendedprice * (1 - l_discount) AS volume
			FROM
				supplier,
				lineitem,
				orders,
				customer,
				(
				SELECT
					n1.n_nationkey AS supp_nation_key,
					n2.n_nationkey AS cust_nation_key,
					n1.n_name AS supp_nation,
					n2.n_name AS cust_nation
				FROM
					nation n1,
					nation n2
				WHERE
					(
					(n1.n_name = $1 AND n2.n_name = $2)
					OR (n1.n_name = $2 AND n2.n_name = $1)
					)
				) AS temp
				WHERE
				s_suppkey = l_suppkey
				AND o_orderkey = l_orderkey
				AND c_custkey = o_custkey
				AND s_nationkey = supp_nation_key
				AND c_nationkey = cust_nation_key
				AND l_shipdate between date '1995-01-01' AND date '1996-12-31'
			) AS shipping
		GROUP BY
			supp_nation,
			cust_nation,
			l_year
		ORDER BY
			supp_nation,
			cust_nation,
			l_year;
END;
$$ LANGUAGE plpgsql;
SET citus.task_executor_type TO 'task-tracker';
SET client_min_messages TO INFO;
-- now, run PL/pgsql functions
SELECT plpgsql_test_1();
 plpgsql_test_1
---------------------------------------------------------------------
           2985
(1 row)

SELECT plpgsql_test_2();
 plpgsql_test_2
---------------------------------------------------------------------
          12000
(1 row)

SELECT plpgsql_test_3();
 plpgsql_test_3
---------------------------------------------------------------------
           1956
(1 row)

SELECT plpgsql_test_4();
 plpgsql_test_4
---------------------------------------------------------------------
           7806
(1 row)

SELECT plpgsql_test_5();
 plpgsql_test_5
---------------------------------------------------------------------
             39
(1 row)

-- run PL/pgsql functions with different parameters
SELECT plpgsql_test_6(155);
 plpgsql_test_6
---------------------------------------------------------------------
          11813
(1 row)

SELECT plpgsql_test_6(1555);
 plpgsql_test_6
---------------------------------------------------------------------
          10185
(1 row)

SELECT plpgsql_test_7('UNITED KINGDOM', 'CHINA');
             plpgsql_test_7
---------------------------------------------------------------------
 ("UNITED KINGDOM",CHINA,1996,18560.22)
(1 row)

SELECT plpgsql_test_7('FRANCE', 'GERMANY');
         plpgsql_test_7
---------------------------------------------------------------------
 (GERMANY,FRANCE,1995,2399.2948)
(1 row)

-- now, PL/pgsql functions with random order
SELECT plpgsql_test_6(155);
 plpgsql_test_6
---------------------------------------------------------------------
          11813
(1 row)

SELECT plpgsql_test_3();
 plpgsql_test_3
---------------------------------------------------------------------
           1956
(1 row)

SELECT plpgsql_test_7('FRANCE', 'GERMANY');
         plpgsql_test_7
---------------------------------------------------------------------
 (GERMANY,FRANCE,1995,2399.2948)
(1 row)

SELECT plpgsql_test_5();
 plpgsql_test_5
---------------------------------------------------------------------
             39
(1 row)

SELECT plpgsql_test_1();
 plpgsql_test_1
---------------------------------------------------------------------
           2985
(1 row)

SELECT plpgsql_test_6(1555);
 plpgsql_test_6
---------------------------------------------------------------------
          10185
(1 row)

SELECT plpgsql_test_4();
 plpgsql_test_4
---------------------------------------------------------------------
           7806
(1 row)

SELECT plpgsql_test_7('UNITED KINGDOM', 'CHINA');
             plpgsql_test_7
---------------------------------------------------------------------
 ("UNITED KINGDOM",CHINA,1996,18560.22)
(1 row)

SELECT plpgsql_test_2();
 plpgsql_test_2
---------------------------------------------------------------------
          12000
(1 row)

-- run the tests which do not require re-partition
-- with real-time executor
RESET citus.task_executor_type;
-- now, run PL/pgsql functions
SELECT plpgsql_test_1();
 plpgsql_test_1
---------------------------------------------------------------------
           2985
(1 row)

SELECT plpgsql_test_2();
 plpgsql_test_2
---------------------------------------------------------------------
          12000
(1 row)

-- run PL/pgsql functions with different parameters
SELECT plpgsql_test_6(155);
 plpgsql_test_6
---------------------------------------------------------------------
          11813
(1 row)

SELECT plpgsql_test_6(1555);
 plpgsql_test_6
---------------------------------------------------------------------
          10185
(1 row)

-- test router executor parameterized PL/pgsql functions
CREATE TABLE plpgsql_table (
	key int,
	value int
);
SET citus.shard_replication_factor TO 1;
SELECT create_distributed_table('plpgsql_table','key','hash');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE FUNCTION no_parameter_insert() RETURNS void as $$
BEGIN
	INSERT INTO plpgsql_table (key) VALUES (0);
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT no_parameter_insert();
 no_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT no_parameter_insert();
 no_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT no_parameter_insert();
 no_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT no_parameter_insert();
 no_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT no_parameter_insert();
 no_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT no_parameter_insert();
 no_parameter_insert
---------------------------------------------------------------------

(1 row)

CREATE FUNCTION single_parameter_insert(key_arg int)
	RETURNS void as $$
BEGIN
	INSERT INTO plpgsql_table (key) VALUES (key_arg);
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT single_parameter_insert(1);
 single_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT single_parameter_insert(2);
 single_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT single_parameter_insert(3);
 single_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT single_parameter_insert(4);
 single_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT single_parameter_insert(5);
 single_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT single_parameter_insert(6);
 single_parameter_insert
---------------------------------------------------------------------

(1 row)

CREATE FUNCTION double_parameter_insert(key_arg int, value_arg int)
	RETURNS void as $$
BEGIN
	INSERT INTO plpgsql_table (key, value) VALUES (key_arg, value_arg);
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT double_parameter_insert(1, 10);
 double_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT double_parameter_insert(2, 20);
 double_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT double_parameter_insert(3, 30);
 double_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT double_parameter_insert(4, 40);
 double_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT double_parameter_insert(5, 50);
 double_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT double_parameter_insert(6, 60);
 double_parameter_insert
---------------------------------------------------------------------

(1 row)

CREATE FUNCTION non_partition_parameter_insert(value_arg int)
	RETURNS void as $$
BEGIN
	INSERT INTO plpgsql_table (key, value) VALUES (0, value_arg);
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT non_partition_parameter_insert(10);
 non_partition_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_insert(20);
 non_partition_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_insert(30);
 non_partition_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_insert(40);
 non_partition_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_insert(50);
 non_partition_parameter_insert
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_insert(60);
 non_partition_parameter_insert
---------------------------------------------------------------------

(1 row)

-- check inserted values
SELECT * FROM plpgsql_table ORDER BY key, value;
 key | value
---------------------------------------------------------------------
   0 |    10
   0 |    20
   0 |    30
   0 |    40
   0 |    50
   0 |    60
   0 |
   0 |
   0 |
   0 |
   0 |
   0 |
   1 |    10
   1 |
   2 |    20
   2 |
   3 |    30
   3 |
   4 |    40
   4 |
   5 |    50
   5 |
   6 |    60
   6 |
(24 rows)

-- check router executor select
CREATE FUNCTION router_partition_column_select(key_arg int)
	RETURNS TABLE(key int, value int) AS $$
DECLARE
BEGIN
    RETURN QUERY
	SELECT
		plpgsql_table.key,
		plpgsql_table.value
	FROM
		plpgsql_table
	WHERE
		plpgsql_table.key = key_arg
	ORDER BY
		key,
		value;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT router_partition_column_select(1);
 router_partition_column_select
---------------------------------------------------------------------
 (1,10)
 (1,)
(2 rows)

SELECT router_partition_column_select(2);
 router_partition_column_select
---------------------------------------------------------------------
 (2,20)
 (2,)
(2 rows)

SELECT router_partition_column_select(3);
 router_partition_column_select
---------------------------------------------------------------------
 (3,30)
 (3,)
(2 rows)

SELECT router_partition_column_select(4);
 router_partition_column_select
---------------------------------------------------------------------
 (4,40)
 (4,)
(2 rows)

SELECT router_partition_column_select(5);
 router_partition_column_select
---------------------------------------------------------------------
 (5,50)
 (5,)
(2 rows)

SELECT router_partition_column_select(6);
 router_partition_column_select
---------------------------------------------------------------------
 (6,60)
 (6,)
(2 rows)

CREATE FUNCTION router_non_partition_column_select(value_arg int)
	RETURNS TABLE(key int, value int) AS $$
DECLARE
BEGIN
    RETURN QUERY
	SELECT
		plpgsql_table.key,
		plpgsql_table.value
	FROM
		plpgsql_table
	WHERE
		plpgsql_table.key = 0 AND
		plpgsql_table.value = value_arg
	ORDER BY
		key,
		value;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT router_non_partition_column_select(10);
 router_non_partition_column_select
---------------------------------------------------------------------
 (0,10)
(1 row)

SELECT router_non_partition_column_select(20);
 router_non_partition_column_select
---------------------------------------------------------------------
 (0,20)
(1 row)

SELECT router_non_partition_column_select(30);
 router_non_partition_column_select
---------------------------------------------------------------------
 (0,30)
(1 row)

SELECT router_non_partition_column_select(40);
 router_non_partition_column_select
---------------------------------------------------------------------
 (0,40)
(1 row)

SELECT router_non_partition_column_select(50);
 router_non_partition_column_select
---------------------------------------------------------------------
 (0,50)
(1 row)

SELECT router_non_partition_column_select(60);
 router_non_partition_column_select
---------------------------------------------------------------------
 (0,60)
(1 row)

-- check real-time executor
CREATE FUNCTION real_time_non_partition_column_select(value_arg int)
	RETURNS TABLE(key int, value int) AS $$
DECLARE
BEGIN
    RETURN QUERY
	SELECT
		plpgsql_table.key,
		plpgsql_table.value
	FROM
		plpgsql_table
	WHERE
		plpgsql_table.value = value_arg
	ORDER BY
		key,
		value;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT real_time_non_partition_column_select(10);
 real_time_non_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
(2 rows)

SELECT real_time_non_partition_column_select(20);
 real_time_non_partition_column_select
---------------------------------------------------------------------
 (0,20)
 (2,20)
(2 rows)

SELECT real_time_non_partition_column_select(30);
 real_time_non_partition_column_select
---------------------------------------------------------------------
 (0,30)
 (3,30)
(2 rows)

SELECT real_time_non_partition_column_select(40);
 real_time_non_partition_column_select
---------------------------------------------------------------------
 (0,40)
 (4,40)
(2 rows)

SELECT real_time_non_partition_column_select(50);
 real_time_non_partition_column_select
---------------------------------------------------------------------
 (0,50)
 (5,50)
(2 rows)

SELECT real_time_non_partition_column_select(60);
 real_time_non_partition_column_select
---------------------------------------------------------------------
 (0,60)
 (6,60)
(2 rows)

CREATE FUNCTION real_time_partition_column_select(key_arg int)
	RETURNS TABLE(key int, value int) AS $$
DECLARE
BEGIN
    RETURN QUERY
	SELECT
		plpgsql_table.key,
		plpgsql_table.value
	FROM
		plpgsql_table
	WHERE
		plpgsql_table.key = key_arg OR
		plpgsql_table.value = 10
	ORDER BY
		key,
		value;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT real_time_partition_column_select(1);
 real_time_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (1,)
(3 rows)

SELECT real_time_partition_column_select(2);
 real_time_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (2,20)
 (2,)
(4 rows)

SELECT real_time_partition_column_select(3);
 real_time_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (3,30)
 (3,)
(4 rows)

SELECT real_time_partition_column_select(4);
 real_time_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (4,40)
 (4,)
(4 rows)

SELECT real_time_partition_column_select(5);
 real_time_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (5,50)
 (5,)
(4 rows)

SELECT real_time_partition_column_select(6);
 real_time_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (6,60)
 (6,)
(4 rows)

-- check task-tracker executor
SET citus.task_executor_type TO 'task-tracker';
CREATE FUNCTION task_tracker_non_partition_column_select(value_arg int)
	RETURNS TABLE(key int, value int) AS $$
DECLARE
BEGIN
    RETURN QUERY
	SELECT
		plpgsql_table.key,
		plpgsql_table.value
	FROM
		plpgsql_table
	WHERE
		plpgsql_table.value = value_arg
	ORDER BY
		key,
		value;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT task_tracker_non_partition_column_select(10);
 task_tracker_non_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
(2 rows)

SELECT task_tracker_non_partition_column_select(20);
 task_tracker_non_partition_column_select
---------------------------------------------------------------------
 (0,20)
 (2,20)
(2 rows)

SELECT task_tracker_non_partition_column_select(30);
 task_tracker_non_partition_column_select
---------------------------------------------------------------------
 (0,30)
 (3,30)
(2 rows)

SELECT task_tracker_non_partition_column_select(40);
 task_tracker_non_partition_column_select
---------------------------------------------------------------------
 (0,40)
 (4,40)
(2 rows)

SELECT task_tracker_non_partition_column_select(50);
 task_tracker_non_partition_column_select
---------------------------------------------------------------------
 (0,50)
 (5,50)
(2 rows)

SELECT real_time_non_partition_column_select(60);
 real_time_non_partition_column_select
---------------------------------------------------------------------
 (0,60)
 (6,60)
(2 rows)

CREATE FUNCTION task_tracker_partition_column_select(key_arg int)
	RETURNS TABLE(key int, value int) AS $$
DECLARE
BEGIN
    RETURN QUERY
	SELECT
		plpgsql_table.key,
		plpgsql_table.value
	FROM
		plpgsql_table
	WHERE
		plpgsql_table.key = key_arg OR
		plpgsql_table.value = 10
	ORDER BY
		key,
		value;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT task_tracker_partition_column_select(1);
 task_tracker_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (1,)
(3 rows)

SELECT task_tracker_partition_column_select(2);
 task_tracker_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (2,20)
 (2,)
(4 rows)

SELECT task_tracker_partition_column_select(3);
 task_tracker_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (3,30)
 (3,)
(4 rows)

SELECT task_tracker_partition_column_select(4);
 task_tracker_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (4,40)
 (4,)
(4 rows)

SELECT task_tracker_partition_column_select(5);
 task_tracker_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (5,50)
 (5,)
(4 rows)

SELECT task_tracker_partition_column_select(6);
 task_tracker_partition_column_select
---------------------------------------------------------------------
 (0,10)
 (1,10)
 (6,60)
 (6,)
(4 rows)

RESET citus.task_executor_type;
-- check updates
CREATE FUNCTION partition_parameter_update(int, int) RETURNS void as $$
BEGIN
	UPDATE plpgsql_table SET value = $2 WHERE key = $1;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT partition_parameter_update(1, 11);
 partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_update(2, 21);
 partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_update(3, 31);
 partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_update(4, 41);
 partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_update(5, 51);
 partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_update(6, 61);
 partition_parameter_update
---------------------------------------------------------------------

(1 row)

CREATE FUNCTION non_partition_parameter_update(int, int) RETURNS void as $$
BEGIN
	UPDATE plpgsql_table SET value = $2 WHERE key = 0 AND value = $1;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT non_partition_parameter_update(10, 12);
 non_partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_update(20, 22);
 non_partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_update(30, 32);
 non_partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_update(40, 42);
 non_partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_update(50, 52);
 non_partition_parameter_update
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_update(60, 62);
 non_partition_parameter_update
---------------------------------------------------------------------

(1 row)

-- check table after updates
SELECT * FROM plpgsql_table ORDER BY key, value;
 key | value
---------------------------------------------------------------------
   0 |    12
   0 |    22
   0 |    32
   0 |    42
   0 |    52
   0 |    62
   0 |
   0 |
   0 |
   0 |
   0 |
   0 |
   1 |    11
   1 |    11
   2 |    21
   2 |    21
   3 |    31
   3 |    31
   4 |    41
   4 |    41
   5 |    51
   5 |    51
   6 |    61
   6 |    61
(24 rows)

-- check deletes
CREATE FUNCTION partition_parameter_delete(int, int) RETURNS void as $$
BEGIN
	DELETE FROM plpgsql_table WHERE key = $1 AND value = $2;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT partition_parameter_delete(1, 11);
 partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_delete(2, 21);
 partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_delete(3, 31);
 partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_delete(4, 41);
 partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_delete(5, 51);
 partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT partition_parameter_delete(6, 61);
 partition_parameter_delete
---------------------------------------------------------------------

(1 row)

CREATE FUNCTION  non_partition_parameter_delete(int) RETURNS void as $$
BEGIN
	DELETE FROM plpgsql_table WHERE key = 0 AND value = $1;
END;
$$ LANGUAGE plpgsql;
-- execute 6 times to trigger prepared statement usage
SELECT non_partition_parameter_delete(12);
 non_partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_delete(22);
 non_partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_delete(32);
 non_partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_delete(42);
 non_partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_delete(52);
 non_partition_parameter_delete
---------------------------------------------------------------------

(1 row)

SELECT non_partition_parameter_delete(62);
 non_partition_parameter_delete
---------------------------------------------------------------------

(1 row)

-- check table after deletes
SELECT * FROM plpgsql_table ORDER BY key, value;
 key | value
---------------------------------------------------------------------
   0 |
   0 |
   0 |
   0 |
   0 |
   0 |
(6 rows)

-- check whether we can handle execute parameters
CREATE TABLE execute_parameter_test (key int, val date);
SELECT create_distributed_table('execute_parameter_test', 'key');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

DO $$
BEGIN
 EXECUTE 'INSERT INTO execute_parameter_test VALUES (3, $1)' USING date '2000-01-01';
 EXECUTE 'INSERT INTO execute_parameter_test VALUES (3, $1)' USING NULL::date;
END;
$$;
DROP TABLE execute_parameter_test;
-- check whether we can handle parameters + default
CREATE TABLE func_parameter_test (
    key text NOT NULL,
    seq int4 NOT NULL,
    created_at timestamptz NOT NULL DEFAULT now(),
    updated_at timestamptz NOT NULL DEFAULT now(),
    PRIMARY KEY (key, seq)
);
SELECT create_distributed_table('func_parameter_test', 'key');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE OR REPLACE FUNCTION insert_with_max(pkey text) RETURNS VOID AS
$BODY$
    DECLARE
        max_seq int4;
    BEGIN
        SELECT MAX(seq) INTO max_seq
        FROM func_parameter_test
        WHERE func_parameter_test.key = pkey;

        IF max_seq IS NULL THEN
            max_seq := 0;
        END IF;

        INSERT INTO func_parameter_test(key, seq) VALUES (pkey, max_seq + 1);
    END;
$BODY$
LANGUAGE plpgsql;
SELECT insert_with_max('key');
 insert_with_max
---------------------------------------------------------------------

(1 row)

SELECT insert_with_max('key');
 insert_with_max
---------------------------------------------------------------------

(1 row)

SELECT insert_with_max('key');
 insert_with_max
---------------------------------------------------------------------

(1 row)

SELECT insert_with_max('key');
 insert_with_max
---------------------------------------------------------------------

(1 row)

SELECT insert_with_max('key');
 insert_with_max
---------------------------------------------------------------------

(1 row)

SELECT insert_with_max('key');
 insert_with_max
---------------------------------------------------------------------

(1 row)

SELECT key, seq FROM func_parameter_test ORDER BY seq;
 key | seq
---------------------------------------------------------------------
 key |   1
 key |   2
 key |   3
 key |   4
 key |   5
 key |   6
(6 rows)

DROP FUNCTION insert_with_max(text);
DROP TABLE func_parameter_test;
-- test prepared DDL, mainly to verify we don't mess up the query tree
SET citus.multi_shard_commit_protocol TO '2pc';
CREATE TABLE prepare_ddl (x int, y int);
SELECT create_distributed_table('prepare_ddl', 'x');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE OR REPLACE FUNCTION ddl_in_plpgsql()
RETURNS VOID  AS
$BODY$
BEGIN
    CREATE INDEX prepared_index ON public.prepare_ddl(x);
    DROP INDEX prepared_index;
END;
$BODY$ LANGUAGE plpgsql;
SELECT ddl_in_plpgsql();
 ddl_in_plpgsql
---------------------------------------------------------------------

(1 row)

SELECT ddl_in_plpgsql();
 ddl_in_plpgsql
---------------------------------------------------------------------

(1 row)

-- test prepared ddl with multi search path to make sure the schema name doesn't leak on
-- to the cached statement
CREATE OR REPLACE FUNCTION ddl_in_plpgsql()
RETURNS VOID  AS
$BODY$
BEGIN
    CREATE INDEX prepared_index ON prepare_ddl(x);
END;
$BODY$ LANGUAGE plpgsql;
CREATE SCHEMA otherschema;
SET search_path TO otherschema, public;
SELECT ddl_in_plpgsql();
 ddl_in_plpgsql
---------------------------------------------------------------------

(1 row)

DROP INDEX prepared_index;
-- this creates the same table it 'otherschema'. If there is a leak the index will not be
-- created on this table, but instead on the table in the public schema
CREATE TABLE prepare_ddl (x int, y int);
SELECT create_distributed_table('prepare_ddl', 'x');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT ddl_in_plpgsql();
 ddl_in_plpgsql
---------------------------------------------------------------------

(1 row)

-- verify the index is created in the correct schema
SELECT schemaname, indexrelname FROM pg_stat_all_indexes WHERE indexrelname = 'prepared_index';
 schemaname  |  indexrelname
---------------------------------------------------------------------
 otherschema | prepared_index
(1 row)

-- cleanup
DROP TABLE prepare_ddl;
RESET search_path;
-- test prepared COPY
CREATE OR REPLACE FUNCTION copy_in_plpgsql()
RETURNS VOID  AS
$BODY$
BEGIN
    COPY prepare_ddl (x) FROM PROGRAM 'echo 1' WITH CSV;
END;
$BODY$ LANGUAGE plpgsql;
SELECT copy_in_plpgsql();
 copy_in_plpgsql
---------------------------------------------------------------------

(1 row)

SELECT copy_in_plpgsql();
 copy_in_plpgsql
---------------------------------------------------------------------

(1 row)

-- test prepared COPY on a non-distributed table
CREATE TABLE local_ddl (x int);
CREATE OR REPLACE FUNCTION local_copy_in_plpgsql()
RETURNS VOID  AS
$BODY$
BEGIN
    COPY local_ddl (x) FROM PROGRAM 'echo 1' WITH CSV;
END;
$BODY$ LANGUAGE plpgsql;
SELECT local_copy_in_plpgsql();
 local_copy_in_plpgsql
---------------------------------------------------------------------

(1 row)

SELECT local_copy_in_plpgsql();
 local_copy_in_plpgsql
---------------------------------------------------------------------

(1 row)

-- types statements should not crash nor leak schema specifications on to cached statements
CREATE TYPE prepare_ddl_type AS (x int, y int);
SET search_path TO 'otherschema', public;
CREATE OR REPLACE FUNCTION public.type_ddl_plpgsql()
RETURNS void
LANGUAGE plpgsql
AS $function$
DECLARE
BEGIN
        ALTER TYPE prepare_ddl_type RENAME TO prepare_ddl_type_backup;
END;
$function$;
SELECT type_ddl_plpgsql();
 type_ddl_plpgsql
---------------------------------------------------------------------

(1 row)

-- create same type in new schema, owner of this new type should change
CREATE TYPE prepare_ddl_type AS (x int, y int);
SELECT type_ddl_plpgsql();
 type_ddl_plpgsql
---------------------------------------------------------------------

(1 row)

-- find all renamed types to verify the schema name didn't leak, nor a crash happened
SELECT nspname, typname FROM pg_type JOIN pg_namespace ON pg_namespace.oid = pg_type.typnamespace WHERE typname = 'prepare_ddl_type_backup';
   nspname   |         typname
---------------------------------------------------------------------
 public      | prepare_ddl_type_backup
 otherschema | prepare_ddl_type_backup
(2 rows)

DROP TYPE prepare_ddl_type_backup;
RESET search_path;
DROP TYPE prepare_ddl_type_backup;
DROP FUNCTION type_ddl_plpgsql();
DROP FUNCTION ddl_in_plpgsql();
DROP FUNCTION copy_in_plpgsql();
DROP TABLE prepare_ddl;
DROP TABLE local_ddl;
DROP SCHEMA otherschema;
-- clean-up functions
DROP FUNCTION plpgsql_test_1();
DROP FUNCTION plpgsql_test_2();
DROP FUNCTION plpgsql_test_3();
DROP FUNCTION plpgsql_test_4();
DROP FUNCTION plpgsql_test_5();
DROP FUNCTION plpgsql_test_6(int);
DROP FUNCTION plpgsql_test_7(text, text);
DROP FUNCTION no_parameter_insert();
DROP FUNCTION single_parameter_insert(int);
DROP FUNCTION double_parameter_insert(int, int);
DROP FUNCTION non_partition_parameter_insert(int);
DROP FUNCTION router_partition_column_select(int);
DROP FUNCTION router_non_partition_column_select(int);
DROP FUNCTION real_time_non_partition_column_select(int);
DROP FUNCTION real_time_partition_column_select(int);
DROP FUNCTION task_tracker_non_partition_column_select(int);
DROP FUNCTION task_tracker_partition_column_select(int);
DROP FUNCTION partition_parameter_update(int, int);
DROP FUNCTION non_partition_parameter_update(int, int);
DROP FUNCTION partition_parameter_delete(int, int);
DROP FUNCTION non_partition_parameter_delete(int);