SET citus.enable_repartition_joins to ON;
-- prevent PG 11 - PG 12 outputs to diverge
SET citus.enable_cte_inlining TO false;
SET citus.max_intermediate_result_size TO 2;
-- should fail because the copy size is ~4kB for each cte
WITH cte AS
(
	SELECT * FROM users_table
),
cte2 AS (
	SELECT * FROM events_table
)
SELECT cte.user_id, cte.value_2 FROM cte,cte2 ORDER BY 1,2 LIMIT 10;
ERROR:  the intermediate result size exceeds citus.max_intermediate_result_size (currently 2 kB)
DETAIL:  Citus restricts the size of intermediate results of complex subqueries and CTEs to avoid accidentally pulling large result sets into once place.
HINT:  To run the current query, set citus.max_intermediate_result_size to a higher value or -1 to disable.
SET citus.max_intermediate_result_size TO 9;
-- regular adaptive CTE should fail
WITH cte AS
(
	SELECT
		users_table.user_id, users_table.value_1, users_table.value_2
	FROM
		users_table
		join
		events_table
	on
		(users_table.value_3=events_table.value_3)
),
cte2 AS (
	SELECT * FROM events_table
)
SELECT
	cte.user_id, cte2.value_2
FROM
	cte JOIN cte2 ON (cte.value_1 = cte2.event_type)
ORDER BY
	1,2
LIMIT 10;
 user_id | value_2
---------------------------------------------------------------------
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
(10 rows)

-- router queries should be able to get limitted too
SET citus.max_intermediate_result_size TO 2;
-- this should pass, since we fetch small portions in each subplan
with cte as (select * from users_table where user_id=1),
cte2 as (select * from users_table where user_id=2),
cte3 as (select * from users_table where user_id=3),
cte4 as (select * from users_table where user_id=4),
cte5 as (select * from users_table where user_id=5)
SELECT sum(c) FROM (
(select count(*) as c from cte)
UNION
(select count(*) as c from cte2)
UNION
(select count(*) as c from cte3)
UNION
(select count(*) as c from cte4)
UNION
(select count(*) as c from cte5)
) as foo;
 sum
---------------------------------------------------------------------
  91
(1 row)

-- if we fetch the same amount of data at once, it should fail
WITH cte AS (SELECT * FROM users_table WHERE user_id IN (1,2,3,4,5))
SELECT * FROM cte ORDER BY 1,2,3,4,5 LIMIT 10;
ERROR:  the intermediate result size exceeds citus.max_intermediate_result_size (currently 2 kB)
DETAIL:  Citus restricts the size of intermediate results of complex subqueries and CTEs to avoid accidentally pulling large result sets into once place.
HINT:  To run the current query, set citus.max_intermediate_result_size to a higher value or -1 to disable.
SET citus.max_intermediate_result_size TO 0;
-- this should fail
WITH cte AS (SELECT * FROM users_table WHERE user_id=1),
cte2 AS (SELECT * FROM users_table WHERE user_id=2),
cte3 AS (SELECT * FROM users_table WHERE user_id=3),
cte4 AS (SELECT * FROM users_table WHERE user_id=4),
cte5 AS (SELECT * FROM users_table WHERE user_id=5)
SELECT * FROM (
(SELECT * FROM cte)
UNION
(SELECT * FROM cte2)
UNION
(SELECT * FROM cte3)
UNION
(SELECT * FROM cte4)
UNION
(SELECT * FROM cte5)
)a ORDER BY 1,2,3,4,5 LIMIT 10;
ERROR:  the intermediate result size exceeds citus.max_intermediate_result_size (currently 0 kB)
DETAIL:  Citus restricts the size of intermediate results of complex subqueries and CTEs to avoid accidentally pulling large result sets into once place.
HINT:  To run the current query, set citus.max_intermediate_result_size to a higher value or -1 to disable.
-- this fails since cte-subplan exceeds limit even if cte2 and cte3 don't
-- WHERE EXISTS forces materialization in pg12
SET citus.max_intermediate_result_size TO 4;
WITH cte AS (
	WITH cte2 AS (
		SELECT * FROM users_table
	),
	cte3 AS (
		SELECT * FROM events_table
	)
	SELECT * FROM cte2, cte3 WHERE cte2.user_id = cte3.user_id AND cte2.user_id = 1
	AND EXISTS (select * from cte2, cte3)
)
SELECT count(*) FROM cte WHERE EXISTS (select * from cte);
 count
---------------------------------------------------------------------
   105
(1 row)

SET citus.max_intermediate_result_size TO 3;
-- this should fail since the cte-subplan exceeds the limit even if the
-- cte2 and cte3 does not
WITH cte AS (
	WITH cte2 AS (
		SELECT * FROM users_table WHERE user_id IN (3,4,5,6)
	),
	cte3 AS (
		SELECT * FROM events_table WHERE event_type = 1
	)
	SELECT * FROM cte2, cte3 WHERE cte2.value_1 IN (SELECT value_2 FROM cte3)
)
SELECT count(*) FROM cte;
 count
---------------------------------------------------------------------
  1824
(1 row)

-- this will fail in real_time_executor
SET citus.max_intermediate_result_size TO 2;
WITH cte AS (
	WITH cte2 AS (
		SELECT * FROM users_table WHERE user_id IN (1, 2)
	),
	cte3 AS (
		SELECT * FROM users_table WHERE user_id = 3
	)
	SELECT * FROM cte2 UNION (SELECT * FROM cte3)
),
cte4 AS (
	SELECT * FROM events_table
)
SELECT * FROM cte UNION ALL
SELECT * FROM cte4 ORDER BY 1,2,3,4,5 LIMIT 5;
ERROR:  the intermediate result size exceeds citus.max_intermediate_result_size (currently 2 kB)
DETAIL:  Citus restricts the size of intermediate results of complex subqueries and CTEs to avoid accidentally pulling large result sets into once place.
HINT:  To run the current query, set citus.max_intermediate_result_size to a higher value or -1 to disable.
SET citus.max_intermediate_result_size TO 1;
-- this will fail in router_executor
WITH cte AS (
	WITH cte2 AS (
		SELECT * FROM users_table WHERE user_id IN (1, 2)
	),
	cte3 AS (
		SELECT * FROM users_table WHERE user_id = 3
	)
	SELECT * FROM cte2 UNION (SELECT * FROM cte3)
),
cte4 AS (
	SELECT * FROM events_table
)
SELECT * FROM cte UNION ALL
SELECT * FROM cte4 ORDER BY 1,2,3,4,5 LIMIT 5;
ERROR:  the intermediate result size exceeds citus.max_intermediate_result_size (currently 1 kB)
DETAIL:  Citus restricts the size of intermediate results of complex subqueries and CTEs to avoid accidentally pulling large result sets into once place.
HINT:  To run the current query, set citus.max_intermediate_result_size to a higher value or -1 to disable.
-- Below that, all should pass since -1 disables the limit
SET citus.max_intermediate_result_size TO -1;
-- real_time_executor + router_executor + real_time_executor will pass
WITH cte AS (
	WITH cte2 AS (
		SELECT * FROM users_table WHERE user_id IN (1, 2)
	),
	cte3 AS (
		SELECT * FROM users_table WHERE user_id = 3
	)
	SELECT * FROM cte2 UNION (SELECT * FROM cte3)
),
cte4 AS (
	SELECT * FROM events_table
)
SELECT * FROM cte UNION ALL
SELECT * FROM cte4 ORDER BY 1,2,3,4,5 LIMIT 5;
 user_id |              time               | value_1 | value_2 | value_3 | value_4
---------------------------------------------------------------------
       1 | Wed Nov 22 18:49:42.327403 2017 |       3 |       2 |       1 |
       1 | Wed Nov 22 19:03:01.772353 2017 |       4 |       1 |       2 |
       1 | Wed Nov 22 19:07:03.846437 2017 |       1 |       2 |       5 |
       1 | Wed Nov 22 20:56:21.122638 2017 |       2 |       4 |       4 |
       1 | Wed Nov 22 21:06:57.457147 2017 |       4 |       3 |       2 |
(5 rows)

-- regular adaptive CTE, should work since -1 disables the limit
WITH cte AS
(
	SELECT
		users_table.user_id, users_table.value_1, users_table.value_2
	FROM
		users_table
		join
		events_table
	on
		(users_table.value_2=events_table.value_2)
),
cte2 AS (
	SELECT * FROM events_table
)
SELECT
	cte.user_id, cte2.value_2
FROM
	cte JOIN cte2 ON (cte.value_1 = cte2.event_type)
ORDER BY
	1,2
LIMIT 10;
 user_id | value_2
---------------------------------------------------------------------
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
(10 rows)

-- regular real-time CTE fetches around ~4kb data in each subplan
WITH cte AS
(
	SELECT * FROM users_table
),
cte2 AS (
	SELECT * FROM events_table
)
SELECT cte.user_id, cte.value_2 FROM cte,cte2 ORDER BY 1,2 LIMIT 10;
 user_id | value_2
---------------------------------------------------------------------
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
       1 |       0
(10 rows)

-- regular real-time query fetches ~4kB
WITH cte AS
(
	SELECT * FROM users_table WHERE user_id IN (1,2,3,4,5)
)
SELECT * FROM cte ORDER BY 1,2,3,4,5 LIMIT 10;
 user_id |              time               | value_1 | value_2 | value_3 | value_4
---------------------------------------------------------------------
       1 | Wed Nov 22 22:51:43.132261 2017 |       4 |       0 |       3 |
       1 | Thu Nov 23 03:32:50.803031 2017 |       3 |       2 |       1 |
       1 | Thu Nov 23 09:26:42.145043 2017 |       1 |       3 |       3 |
       1 | Thu Nov 23 11:11:24.40789 2017  |       3 |       4 |       0 |
       1 | Thu Nov 23 11:44:57.515981 2017 |       4 |       3 |       4 |
       1 | Thu Nov 23 17:23:03.441394 2017 |       5 |       4 |       3 |
       1 | Thu Nov 23 17:30:34.635085 2017 |       3 |       4 |       4 |
       2 | Wed Nov 22 18:19:49.944985 2017 |       3 |       5 |       1 |
       2 | Thu Nov 23 00:19:14.138058 2017 |       3 |       4 |       0 |
       2 | Thu Nov 23 01:04:26.198826 2017 |       4 |       3 |       4 |
(10 rows)

-- nested CTEs
WITH cte AS (
	WITH cte2 AS (
		SELECT * FROM users_table
	),
	cte3 AS (
		SELECT * FROM events_table
	)
	SELECT
		cte2.user_id, cte2.time, cte3.event_type, cte3.value_2, cte3.value_3
	FROM
		cte2, cte3
	WHERE
		cte2.user_id = cte3.user_id AND cte2.user_id = 1
)
SELECT * FROM cte ORDER BY 1,2,3,4,5 LIMIT 10;
 user_id |              time               | event_type | value_2 | value_3
---------------------------------------------------------------------
       1 | Wed Nov 22 22:51:43.132261 2017 |          0 |       2 |       0
       1 | Wed Nov 22 22:51:43.132261 2017 |          0 |       5 |       1
       1 | Wed Nov 22 22:51:43.132261 2017 |          1 |       1 |       1
       1 | Wed Nov 22 22:51:43.132261 2017 |          1 |       2 |       5
       1 | Wed Nov 22 22:51:43.132261 2017 |          2 |       4 |       3
       1 | Wed Nov 22 22:51:43.132261 2017 |          2 |       4 |       4
       1 | Wed Nov 22 22:51:43.132261 2017 |          3 |       1 |       1
       1 | Wed Nov 22 22:51:43.132261 2017 |          3 |       2 |       1
       1 | Wed Nov 22 22:51:43.132261 2017 |          3 |       4 |       1
       1 | Wed Nov 22 22:51:43.132261 2017 |          4 |       1 |       2
(10 rows)