CREATE SCHEMA cte_failure;
SET SEARCH_PATH=cte_failure;
SET citus.shard_count to 2;
SET citus.shard_replication_factor to 1;
SET citus.next_shard_id TO 16000000;
-- CTE inlining should not happen because
-- the tests rely on intermediate results
-- That's why we use MATERIALIZED CTEs in the test file
-- prevent using locally executing the intermediate results
SET citus.task_assignment_policy TO "round-robin";
SELECT pg_backend_pid() as pid \gset
CREATE TABLE users_table (user_id int, user_name text);
CREATE TABLE events_table(user_id int, event_id int, event_type int);
SELECT create_distributed_table('users_table', 'user_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('events_table', 'user_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE users_table_local AS SELECT * FROM users_table;
-- kill at the first copy (push)
SELECT citus.mitmproxy('conn.onQuery(query="^COPY").kill()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte AS MATERIALIZED (
	WITH local_cte AS MATERIALIZED (
		SELECT * FROM users_table_local
	),
	dist_cte AS MATERIALIZED (
		SELECT user_id FROM events_table
	)
	SELECT dist_cte.user_id FROM local_cte join dist_cte on dist_cte.user_id=local_cte.user_id
)
SELECT
	count(*)
FROM
	cte,
	  (SELECT
    	DISTINCT users_table.user_id
     FROM
     	users_table, events_table
     WHERE
     	users_table.user_id = events_table.user_id AND
     event_type IN (1,2,3,4)
     ORDER BY 1 DESC LIMIT 5
     ) as foo
	  WHERE foo.user_id = cte.user_id;
ERROR:  connection not open
CONTEXT:  while executing command on localhost:xxxxx
-- kill at the second copy (pull)
SELECT citus.mitmproxy('conn.onQuery(query="SELECT user_id FROM cte_failure.events_table_16000002").kill()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte AS MATERIALIZED (
	WITH local_cte AS MATERIALIZED (
		SELECT * FROM users_table_local
	),
	dist_cte AS MATERIALIZED (
		SELECT user_id FROM events_table
	)
	SELECT dist_cte.user_id FROM local_cte join dist_cte on dist_cte.user_id=local_cte.user_id
)
SELECT
	count(*)
FROM
	cte,
	  (SELECT
    	DISTINCT users_table.user_id
     FROM
     	users_table, events_table
     WHERE
     	users_table.user_id = events_table.user_id AND
     event_type IN (1,2,3,4)
     ORDER BY 1 DESC LIMIT 5
     ) as foo
	  WHERE foo.user_id = cte.user_id;
ERROR:  connection to the remote node postgres@localhost:xxxxx failed with the following error: connection not open
-- kill at the third copy (pull)
SELECT citus.mitmproxy('conn.onQuery(query="SELECT DISTINCT users_table.user").kill()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte AS MATERIALIZED (
	WITH local_cte AS MATERIALIZED (
		SELECT * FROM users_table_local
	),
	dist_cte AS MATERIALIZED (
		SELECT user_id FROM events_table
	)
	SELECT dist_cte.user_id FROM local_cte join dist_cte on dist_cte.user_id=local_cte.user_id
)
SELECT
	count(*)
FROM
	cte,
	  (SELECT
    	DISTINCT users_table.user_id
     FROM
     	users_table, events_table
     WHERE
     	users_table.user_id = events_table.user_id AND
     event_type IN (1,2,3,4)
     ORDER BY 1 DESC LIMIT 5
     ) as foo
	  WHERE foo.user_id = cte.user_id;
ERROR:  connection to the remote node postgres@localhost:xxxxx failed with the following error: connection not open
-- cancel at the first copy (push)
SELECT citus.mitmproxy('conn.onQuery(query="^COPY").cancel(' || :pid || ')');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte AS MATERIALIZED (
	WITH local_cte AS MATERIALIZED (
		SELECT * FROM users_table_local
	),
	dist_cte AS MATERIALIZED (
		SELECT user_id FROM events_table
	)
	SELECT dist_cte.user_id FROM local_cte join dist_cte on dist_cte.user_id=local_cte.user_id
)
SELECT
	count(*)
FROM
	cte,
	  (SELECT
    	DISTINCT users_table.user_id
     FROM
     	users_table, events_table
     WHERE
     	users_table.user_id = events_table.user_id AND
     event_type IN (1,2,3,4)
     ORDER BY 1 DESC LIMIT 5
     ) as foo
	  WHERE foo.user_id = cte.user_id;
ERROR:  canceling statement due to user request
-- cancel at the second copy (pull)
SELECT citus.mitmproxy('conn.onQuery(query="SELECT user_id FROM").cancel(' || :pid || ')');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte AS MATERIALIZED (
	WITH local_cte AS MATERIALIZED (
		SELECT * FROM users_table_local
	),
	dist_cte AS MATERIALIZED (
		SELECT user_id FROM events_table
	)
	SELECT dist_cte.user_id FROM local_cte join dist_cte on dist_cte.user_id=local_cte.user_id
)
SELECT
	count(*)
FROM
	cte,
	  (SELECT
    	DISTINCT users_table.user_id
     FROM
     	users_table, events_table
     WHERE
     	users_table.user_id = events_table.user_id AND
     event_type IN (1,2,3,4)
     ORDER BY 1 DESC LIMIT 5
     ) as foo
	  WHERE foo.user_id = cte.user_id;
ERROR:  canceling statement due to user request
-- cancel at the third copy (pull)
SELECT citus.mitmproxy('conn.onQuery(query="SELECT DISTINCT users_table.user").cancel(' || :pid || ')');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte AS MATERIALIZED (
	WITH local_cte AS MATERIALIZED (
		SELECT * FROM users_table_local
	),
	dist_cte AS MATERIALIZED (
		SELECT user_id FROM events_table
	)
	SELECT dist_cte.user_id FROM local_cte join dist_cte on dist_cte.user_id=local_cte.user_id
)
SELECT
	count(*)
FROM
	cte,
	  (SELECT
    	DISTINCT users_table.user_id
     FROM
     	users_table, events_table
     WHERE
     	users_table.user_id = events_table.user_id AND
     event_type IN (1,2,3,4)
     ORDER BY 1 DESC LIMIT 5
     ) as foo
	  WHERE foo.user_id = cte.user_id;
ERROR:  canceling statement due to user request
-- distributed update tests
SELECT citus.mitmproxy('conn.allow()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

-- insert some rows
INSERT INTO users_table VALUES (1, 'A'), (2, 'B'), (3, 'C'), (4, 'D'), (5, 'E');
INSERT INTO events_table VALUES (1,1,1), (1,2,1), (1,3,1), (2,1, 4), (3, 4,1), (5, 1, 2), (5, 2, 1), (5, 2,2);
SELECT * FROM users_table ORDER BY 1, 2;
 user_id | user_name
---------------------------------------------------------------------
       1 | A
       2 | B
       3 | C
       4 | D
       5 | E
(5 rows)

-- following will delete and insert the same rows
WITH cte_delete AS MATERIALIZED (DELETE FROM users_table WHERE user_name in ('A', 'D') RETURNING *)
INSERT INTO users_table SELECT * FROM cte_delete;
-- verify contents are the same
SELECT * FROM users_table ORDER BY 1, 2;
 user_id | user_name
---------------------------------------------------------------------
       1 | A
       2 | B
       3 | C
       4 | D
       5 | E
(5 rows)

-- kill connection during deletion
SELECT citus.mitmproxy('conn.onQuery(query="^DELETE FROM").kill()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte_delete AS MATERIALIZED (DELETE FROM users_table WHERE user_name in ('A', 'D') RETURNING *)
INSERT INTO users_table SELECT * FROM cte_delete;
ERROR:  connection to the remote node postgres@localhost:xxxxx failed with the following error: connection not open
-- verify contents are the same
SELECT citus.mitmproxy('conn.allow()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

SELECT * FROM users_table ORDER BY 1, 2;
 user_id | user_name
---------------------------------------------------------------------
       1 | A
       2 | B
       3 | C
       4 | D
       5 | E
(5 rows)

-- kill connection during insert
SELECT citus.mitmproxy('conn.onQuery(query="^COPY").kill()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte_delete AS MATERIALIZED (DELETE FROM users_table WHERE user_name in ('A', 'D') RETURNING *)
INSERT INTO users_table SELECT * FROM cte_delete;
ERROR:  connection not open
CONTEXT:  while executing command on localhost:xxxxx
-- verify contents are the same
SELECT citus.mitmproxy('conn.allow()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

SELECT * FROM users_table ORDER BY 1, 2;
 user_id | user_name
---------------------------------------------------------------------
       1 | A
       2 | B
       3 | C
       4 | D
       5 | E
(5 rows)

-- cancel during deletion
SELECT citus.mitmproxy('conn.onQuery(query="^DELETE FROM").cancel(' || :pid || ')');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte_delete AS MATERIALIZED (DELETE FROM users_table WHERE user_name in ('A', 'D') RETURNING *)
INSERT INTO users_table SELECT * FROM cte_delete;
ERROR:  canceling statement due to user request
-- verify contents are the same
SELECT citus.mitmproxy('conn.allow()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

SELECT * FROM users_table ORDER BY 1, 2;
 user_id | user_name
---------------------------------------------------------------------
       1 | A
       2 | B
       3 | C
       4 | D
       5 | E
(5 rows)

-- cancel during insert
SELECT citus.mitmproxy('conn.onQuery(query="^COPY").cancel(' || :pid || ')');
 mitmproxy
---------------------------------------------------------------------

(1 row)

WITH cte_delete AS MATERIALIZED (DELETE FROM users_table WHERE user_name in ('A', 'D') RETURNING *)
INSERT INTO users_table SELECT * FROM cte_delete;
ERROR:  canceling statement due to user request
-- verify contents are the same
SELECT citus.mitmproxy('conn.allow()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

SELECT * FROM users_table ORDER BY 1, 2;
 user_id | user_name
---------------------------------------------------------------------
       1 | A
       2 | B
       3 | C
       4 | D
       5 | E
(5 rows)

-- test sequential delete/insert
SELECT citus.mitmproxy('conn.onQuery(query="^DELETE FROM").kill()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

BEGIN;
SET LOCAL citus.multi_shard_modify_mode = 'sequential';
WITH cte_delete AS MATERIALIZED (DELETE FROM users_table WHERE user_name in ('A', 'D') RETURNING *)
INSERT INTO users_table SELECT * FROM cte_delete;
ERROR:  connection to the remote node postgres@localhost:xxxxx failed with the following error: connection not open
END;
RESET SEARCH_PATH;
SELECT citus.mitmproxy('conn.allow()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

DROP SCHEMA cte_failure CASCADE;
NOTICE:  drop cascades to 3 other objects
DETAIL:  drop cascades to table cte_failure.users_table
drop cascades to table cte_failure.events_table
drop cascades to table cte_failure.users_table_local