-- Tests for modifying CTEs and CTEs in modifications
SET citus.next_shard_id TO 1502000;
CREATE SCHEMA with_modifying;
SET search_path TO with_modifying, public;
CREATE TABLE with_modifying.local_table (id int, val int);
CREATE TABLE with_modifying.modify_table (id int, val int);
SELECT create_distributed_table('modify_table', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE with_modifying.users_table (LIKE public.users_table INCLUDING ALL);
SELECT create_distributed_table('with_modifying.users_table', 'user_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

INSERT INTO with_modifying.users_table SELECT * FROM public.users_table;
CREATE TABLE with_modifying.summary_table (id int, counter int);
SELECT create_distributed_table('summary_table', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE with_modifying.anchor_table (id int);
SELECT create_reference_table('anchor_table');
 create_reference_table
---------------------------------------------------------------------

(1 row)

-- basic insert query in CTE
WITH basic_insert AS (
	INSERT INTO users_table VALUES (1), (2), (3) RETURNING *
)
SELECT
	*
FROM
	basic_insert
ORDER BY
	user_id;
 user_id | time | value_1 | value_2 | value_3 | value_4
---------------------------------------------------------------------
       1 |      |         |         |         |
       2 |      |         |         |         |
       3 |      |         |         |         |
(3 rows)

-- single-shard UPDATE in CTE
WITH basic_update AS (
	UPDATE users_table SET value_3=41 WHERE user_id=1 RETURNING *
)
SELECT
	*
FROM
	basic_update
ORDER BY
	user_id,
	time
LIMIT 10;
 user_id |              time               | value_1 | value_2 | value_3 | value_4
---------------------------------------------------------------------
       1 | Wed Nov 22 22:51:43.132261 2017 |       4 |       0 |      41 |
       1 | Thu Nov 23 03:32:50.803031 2017 |       3 |       2 |      41 |
       1 | Thu Nov 23 09:26:42.145043 2017 |       1 |       3 |      41 |
       1 | Thu Nov 23 11:11:24.40789 2017  |       3 |       4 |      41 |
       1 | Thu Nov 23 11:44:57.515981 2017 |       4 |       3 |      41 |
       1 | Thu Nov 23 17:23:03.441394 2017 |       5 |       4 |      41 |
       1 | Thu Nov 23 17:30:34.635085 2017 |       3 |       4 |      41 |
       1 |                                 |         |         |      41 |
       1 |                                 |         |         |      41 |
(9 rows)

-- multi-shard UPDATE in CTE
WITH basic_update AS (
	UPDATE users_table SET value_3=42 WHERE value_2=1 RETURNING *
)
SELECT
	*
FROM
	basic_update
ORDER BY
	user_id,
	time
LIMIT 10;
 user_id |              time               | value_1 | value_2 | value_3 | value_4
---------------------------------------------------------------------
       2 | Thu Nov 23 06:50:30.797805 2017 |       1 |       1 |      42 |
       2 | Thu Nov 23 06:56:38.46819 2017  |       0 |       1 |      42 |
       2 | Thu Nov 23 13:52:54.83829 2017  |       3 |       1 |      42 |
       3 | Wed Nov 22 18:43:51.450263 2017 |       1 |       1 |      42 |
       3 | Thu Nov 23 00:15:45.610845 2017 |       1 |       1 |      42 |
       4 | Wed Nov 22 23:59:46.493416 2017 |       3 |       1 |      42 |
       4 | Thu Nov 23 01:55:21.824618 2017 |       3 |       1 |      42 |
       4 | Thu Nov 23 06:50:08.101207 2017 |       2 |       1 |      42 |
       4 | Thu Nov 23 07:09:37.382372 2017 |       4 |       1 |      42 |
       4 | Thu Nov 23 08:38:45.877401 2017 |       4 |       1 |      42 |
(10 rows)

-- single-shard DELETE in CTE
WITH basic_delete AS (
	DELETE FROM users_table WHERE user_id=6 RETURNING *
)
SELECT
	*
FROM
	basic_delete
ORDER BY
	user_id,
	time
LIMIT 10;
 user_id |              time               | value_1 | value_2 | value_3 | value_4
---------------------------------------------------------------------
       6 | Wed Nov 22 20:15:53.317797 2017 |       1 |       1 |      42 |
       6 | Wed Nov 22 23:01:24.82289 2017  |       2 |       4 |       1 |
       6 | Thu Nov 23 00:07:11.068353 2017 |       1 |       1 |      42 |
       6 | Thu Nov 23 00:09:44.19812 2017  |       5 |       2 |       0 |
       6 | Thu Nov 23 01:13:50.526322 2017 |       2 |       4 |       1 |
       6 | Thu Nov 23 01:14:55.769581 2017 |       0 |       0 |       5 |
       6 | Thu Nov 23 10:22:11.02918 2017  |       5 |       0 |       5 |
       6 | Thu Nov 23 11:08:04.244582 2017 |       2 |       3 |       2 |
       6 | Thu Nov 23 13:51:16.92838 2017  |       0 |       4 |       2 |
       6 | Thu Nov 23 14:43:18.024104 2017 |       3 |       2 |       5 |
(10 rows)

-- multi-shard DELETE in CTE
WITH basic_delete AS (
	DELETE FROM users_table WHERE value_3=41 RETURNING *
)
SELECT
	*
FROM
	basic_delete
ORDER BY
	user_id,
	time
LIMIT 10;
 user_id |              time               | value_1 | value_2 | value_3 | value_4
---------------------------------------------------------------------
       1 | Wed Nov 22 22:51:43.132261 2017 |       4 |       0 |      41 |
       1 | Thu Nov 23 03:32:50.803031 2017 |       3 |       2 |      41 |
       1 | Thu Nov 23 09:26:42.145043 2017 |       1 |       3 |      41 |
       1 | Thu Nov 23 11:11:24.40789 2017  |       3 |       4 |      41 |
       1 | Thu Nov 23 11:44:57.515981 2017 |       4 |       3 |      41 |
       1 | Thu Nov 23 17:23:03.441394 2017 |       5 |       4 |      41 |
       1 | Thu Nov 23 17:30:34.635085 2017 |       3 |       4 |      41 |
       1 |                                 |         |         |      41 |
       1 |                                 |         |         |      41 |
(9 rows)

-- INSERT...SELECT query in CTE
WITH copy_table AS (
	INSERT INTO users_table SELECT * FROM users_table WHERE user_id = 0 OR user_id = 3 RETURNING *
)
SELECT
	*
FROM
	copy_table
ORDER BY
	user_id,
	time
LIMIT 10;
 user_id |              time               | value_1 | value_2 | value_3 | value_4
---------------------------------------------------------------------
       3 | Wed Nov 22 18:43:51.450263 2017 |       1 |       1 |      42 |
       3 | Wed Nov 22 20:43:31.008625 2017 |       1 |       3 |       2 |
       3 | Wed Nov 22 23:24:32.080584 2017 |       3 |       2 |       5 |
       3 | Thu Nov 23 00:15:45.610845 2017 |       1 |       1 |      42 |
       3 | Thu Nov 23 03:23:24.702501 2017 |       1 |       2 |       5 |
       3 | Thu Nov 23 03:52:32.008895 2017 |       4 |       2 |       0 |
       3 | Thu Nov 23 04:01:08.04806 2017  |       5 |       5 |       3 |
       3 | Thu Nov 23 05:01:44.885505 2017 |       3 |       5 |       4 |
       3 | Thu Nov 23 06:20:05.854857 2017 |       1 |       4 |       2 |
       3 | Thu Nov 23 09:57:41.540228 2017 |       2 |       2 |       3 |
(10 rows)

-- CTEs prior to INSERT...SELECT via the coordinator should work
WITH cte AS (
	SELECT user_id FROM users_table WHERE value_2 IN (1, 2)
)
INSERT INTO modify_table (SELECT * FROM cte);
WITH cte_1 AS (
	SELECT user_id, value_2 FROM users_table WHERE value_2 IN (1, 2, 3, 4)
),
cte_2 AS (
	SELECT user_id, value_2 FROM users_table WHERE value_2 IN (3, 4, 5, 6)
)
INSERT INTO modify_table (SELECT cte_1.user_id FROM cte_1 join cte_2 on cte_1.value_2=cte_2.value_2);
-- we execute the query within a function to consolidate the error messages
-- between different executors
CREATE FUNCTION raise_failed_execution_cte(query text) RETURNS void AS $$
BEGIN
	EXECUTE query;
	EXCEPTION WHEN OTHERS THEN
	IF SQLERRM LIKE '%more than one row returned by a subquery used as an expression%' THEN
		RAISE 'Task failed to execute';
	ELSIF SQLERRM LIKE '%could not receive query results%' THEN
		RAISE 'Task failed to execute';
	END IF;
END;
$$LANGUAGE plpgsql;
SET client_min_messages TO ERROR;
\set VERBOSITY terse
-- even if this is an INSERT...SELECT, the CTE is under SELECT
-- function joins in CTE results can create lateral joins that are not supported
SELECT raise_failed_execution_cte($$
	WITH cte AS (
		SELECT user_id, value_2 FROM users_table WHERE value_2 IN (1, 2)
	)
	INSERT INTO modify_table (SELECT (SELECT value_2 FROM cte GROUP BY value_2));
$$);
ERROR:  Task failed to execute
SET client_min_messages TO DEFAULT;
\set VERBOSITY DEFAULT
-- CTEs prior to any other modification should error out
WITH cte AS (
	SELECT value_2 FROM users_table WHERE user_id IN (1, 2, 3)
)
DELETE FROM modify_table WHERE id IN (SELECT value_2 FROM cte);
WITH cte AS (
	SELECT value_2 FROM users_table WHERE user_id IN (1, 2, 3)
)
UPDATE modify_table SET val=-1 WHERE val IN (SELECT * FROM cte);
WITH user_data AS (
	SELECT user_id, value_2 FROM users_table
)
INSERT INTO modify_table SELECT * FROM user_data;
WITH raw_data AS (
	DELETE FROM modify_table RETURNING *
)
INSERT INTO summary_table SELECT id, COUNT(*) AS counter FROM raw_data GROUP BY id;
SELECT * FROM summary_table ORDER BY id;
 id | counter
---------------------------------------------------------------------
  2 |      20
  3 |      38
  4 |      24
  5 |      27
(4 rows)

SELECT COUNT(*) FROM modify_table;
 count
---------------------------------------------------------------------
     0
(1 row)

INSERT INTO modify_table VALUES (1,1), (2, 2), (3,3);
WITH raw_data AS (
	DELETE FROM modify_table RETURNING *
)
INSERT INTO summary_table SELECT id, COUNT(*) AS counter FROM raw_data GROUP BY id;
SELECT * FROM summary_table ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
  1 |       1
  2 |       1
  2 |      20
  3 |       1
  3 |      38
  4 |      24
  5 |      27
(7 rows)

SELECT COUNT(*) FROM modify_table;
 count
---------------------------------------------------------------------
     0
(1 row)

WITH insert_reference AS (
	INSERT INTO anchor_table VALUES (1), (2) RETURNING *
)
SELECT id FROM insert_reference ORDER BY id;
 id
---------------------------------------------------------------------
  1
  2
(2 rows)

WITH anchor_data AS (
	SELECT * FROM anchor_table
),
raw_data AS (
	DELETE FROM modify_table RETURNING *
),
summary_data AS (
	DELETE FROM summary_table RETURNING *
)
INSERT INTO
	summary_table
SELECT id, SUM(counter) FROM (
	(SELECT raw_data.id, COUNT(*) AS counter FROM raw_data, anchor_data
		WHERE raw_data.id = anchor_data.id GROUP BY raw_data.id)
	UNION ALL
	(SELECT * FROM summary_data)) AS all_rows
GROUP BY
	id;
SELECT COUNT(*) FROM modify_table;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT * FROM summary_table ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
  1 |       1
  2 |      21
  3 |      39
  4 |      24
  5 |      27
(5 rows)

WITH added_data AS (
	INSERT INTO modify_table VALUES (1,2), (1,6), (2,4), (3,6) RETURNING *
),
raw_data AS (
	DELETE FROM modify_table WHERE id = 1 AND val = (SELECT MAX(val) FROM added_data) RETURNING *
)
INSERT INTO summary_table SELECT id, COUNT(*) AS counter FROM raw_data GROUP BY id;
SELECT COUNT(*) FROM modify_table;
 count
---------------------------------------------------------------------
     3
(1 row)

SELECT * FROM summary_table ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
  1 |       1
  1 |       1
  2 |      21
  3 |      39
  4 |      24
  5 |      27
(6 rows)

-- Merge rows in the summary_table
WITH summary_data AS (
	DELETE FROM summary_table RETURNING *
)
INSERT INTO summary_table SELECT id, SUM(counter) AS counter FROM summary_data GROUP BY id;
SELECT * FROM summary_table ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
  1 |       2
  2 |      21
  3 |      39
  4 |      24
  5 |      27
(5 rows)

SELECT * FROM modify_table ORDER BY id, val;
 id | val
---------------------------------------------------------------------
  1 |   2
  2 |   4
  3 |   6
(3 rows)

SELECT * FROM anchor_table ORDER BY id;
 id
---------------------------------------------------------------------
  1
  2
(2 rows)

INSERT INTO modify_table VALUES (11, 1), (12, 2), (13, 3);
WITH select_data AS (
	SELECT * FROM modify_table
),
raw_data AS (
	DELETE FROM modify_table WHERE id >= (SELECT min(id) FROM select_data WHERE id > 10) RETURNING *
)
INSERT INTO summary_table SELECT id, COUNT(*) AS counter FROM raw_data GROUP BY id;
INSERT INTO modify_table VALUES (21, 1), (22, 2), (23, 3);
-- read ids from the same table
WITH distinct_ids AS (
	SELECT DISTINCT id FROM modify_table
),
update_data AS (
	UPDATE modify_table SET val = 100 WHERE id > 10 AND
		id IN (SELECT * FROM distinct_ids) RETURNING *
)
SELECT count(*) FROM update_data;
 count
---------------------------------------------------------------------
     3
(1 row)

-- read ids from a different table
WITH distinct_ids AS (
	SELECT DISTINCT id FROM summary_table
),
update_data AS (
	UPDATE modify_table SET val = 100 WHERE id > 10 AND
		id IN (SELECT * FROM distinct_ids) RETURNING *
)
SELECT count(*) FROM update_data;
 count
---------------------------------------------------------------------
     0
(1 row)

-- test update with generate series
UPDATE modify_table SET val = 200 WHERE id > 10 AND
	id IN (SELECT 2*s FROM generate_series(1,20) s);
-- test update with generate series in CTE
WITH update_data AS (
	UPDATE modify_table SET val = 300 WHERE id > 10 AND
	id IN (SELECT 3*s FROM generate_series(1,20) s) RETURNING *
)
SELECT COUNT(*) FROM update_data;
 count
---------------------------------------------------------------------
     1
(1 row)

WITH delete_rows AS (
	DELETE FROM modify_table WHERE id > 10 RETURNING *
)
SELECT * FROM delete_rows ORDER BY id, val;
 id | val
---------------------------------------------------------------------
 21 | 300
 22 | 200
 23 | 100
(3 rows)

WITH delete_rows AS (
	DELETE FROM summary_table WHERE id > 10 RETURNING *
)
SELECT * FROM delete_rows ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
 11 |       1
 12 |       1
 13 |       1
(3 rows)

-- Check modifiying CTEs inside a transaction
BEGIN;
WITH raw_data AS (
	DELETE FROM modify_table RETURNING *
)
INSERT INTO summary_table SELECT id, COUNT(*) AS counter FROM raw_data GROUP BY id;
WITH insert_reference AS (
	INSERT INTO anchor_table VALUES (3), (4) RETURNING *
)
SELECT id FROM insert_reference ORDER BY id;
 id
---------------------------------------------------------------------
  3
  4
(2 rows)

SELECT * FROM summary_table ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
  1 |       1
  1 |       2
  2 |       1
  2 |      21
  3 |       1
  3 |      39
  4 |      24
  5 |      27
(8 rows)

SELECT * FROM modify_table ORDER BY id, val;
 id | val
---------------------------------------------------------------------
(0 rows)

SELECT * FROM anchor_table ORDER BY id;
 id
---------------------------------------------------------------------
  1
  2
  3
  4
(4 rows)

ROLLBACK;
SELECT * FROM summary_table ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
  1 |       2
  2 |      21
  3 |      39
  4 |      24
  5 |      27
(5 rows)

SELECT * FROM modify_table ORDER BY id, val;
 id | val
---------------------------------------------------------------------
  1 |   2
  2 |   4
  3 |   6
(3 rows)

SELECT * FROM anchor_table ORDER BY id;
 id
---------------------------------------------------------------------
  1
  2
(2 rows)

-- Test delete with subqueries
WITH deleted_rows AS (
	DELETE FROM modify_table WHERE id IN (SELECT id FROM modify_table WHERE id = 1) RETURNING *
)
SELECT * FROM deleted_rows;
 id | val
---------------------------------------------------------------------
  1 |   2
(1 row)

WITH deleted_rows AS (
	DELETE FROM modify_table WHERE id IN (SELECT id FROM modify_table WHERE val = 4) RETURNING *
)
SELECT * FROM deleted_rows;
 id | val
---------------------------------------------------------------------
  2 |   4
(1 row)

WITH select_rows AS (
	SELECT id FROM modify_table WHERE val = 4
),
deleted_rows AS (
	DELETE FROM modify_table WHERE id IN (SELECT id FROM select_rows) RETURNING *
)
SELECT * FROM deleted_rows;
 id | val
---------------------------------------------------------------------
(0 rows)

WITH deleted_rows AS (
	DELETE FROM modify_table WHERE val IN (SELECT val FROM modify_table WHERE id = 3) RETURNING *
)
SELECT * FROM deleted_rows;
 id | val
---------------------------------------------------------------------
  3 |   6
(1 row)

WITH select_rows AS (
	SELECT val FROM modify_table WHERE id = 3
),
deleted_rows AS (
	DELETE FROM modify_table WHERE val IN (SELECT val FROM select_rows) RETURNING *
)
SELECT * FROM deleted_rows;
 id | val
---------------------------------------------------------------------
(0 rows)

WITH deleted_rows AS (
	DELETE FROM modify_table WHERE ctid IN (SELECT ctid FROM modify_table WHERE id = 1) RETURNING *
)
SELECT * FROM deleted_rows;
ERROR:  cannot perform distributed planning for the given modification
DETAIL:  Recursively planned distributed modifications with ctid on where clause are not supported.
WITH select_rows AS (
	SELECT ctid FROM modify_table WHERE id = 1
),
deleted_rows AS (
	DELETE FROM modify_table WHERE ctid IN (SELECT ctid FROM select_rows) RETURNING *
)
SELECT * FROM deleted_rows;
ERROR:  cannot perform distributed planning for the given modification
DETAIL:  Recursively planned distributed modifications with ctid on where clause are not supported.
WITH added_data AS (
	INSERT INTO modify_table VALUES (1,2), (1,6) RETURNING *
),
select_data AS (
	SELECT * FROM added_data WHERE id = 1
),
raw_data AS (
	DELETE FROM modify_table WHERE id = 1 AND ctid IN (SELECT ctid FROM select_data) RETURNING val
)
SELECT * FROM raw_data ORDER BY val;
ERROR:  cannot perform distributed planning for the given modification
DETAIL:  Recursively planned distributed modifications with ctid on where clause are not supported.
-- Needed becaues of CTE inlining triggering https://github.com/citusdata/citus/issues/3189
SET citus.enable_cte_inlining TO FALSE;
WITH added_data AS (
	INSERT INTO modify_table VALUES (1, trunc(10 * random())), (1, trunc(random())) RETURNING *
),
select_data AS (
	SELECT val, now() FROM added_data WHERE id = 1
),
raw_data AS (
	DELETE FROM modify_table WHERE id = 1 AND val IN (SELECT val FROM select_data) RETURNING *
)
SELECT COUNT(*) FROM raw_data;
 count
---------------------------------------------------------------------
     2
(1 row)

SET citus.enable_cte_inlining TO TRUE;
WITH added_data AS (
	INSERT INTO modify_table VALUES (1, trunc(10 * random())), (1, trunc(random())) RETURNING *
),
select_data AS (
	SELECT val, '2011-01-01' FROM added_data WHERE id = 1
),
raw_data AS (
	DELETE FROM modify_table WHERE id = 1 AND val IN (SELECT val FROM select_data) RETURNING *
)
SELECT COUNT(*) FROM raw_data;
 count
---------------------------------------------------------------------
     2
(1 row)

INSERT INTO modify_table VALUES (1,2), (1,6), (2, 3), (3, 5);
WITH select_data AS (
	SELECT * FROM modify_table
),
raw_data AS (
	DELETE FROM modify_table WHERE id IN (SELECT id FROM select_data WHERE val > 5) RETURNING id, val
)
SELECT * FROM raw_data ORDER BY val;
 id | val
---------------------------------------------------------------------
  1 |   2
  1 |   6
(2 rows)

WITH select_data AS (
	SELECT * FROM modify_table
),
raw_data AS (
	UPDATE modify_table SET val = 0 WHERE id IN (SELECT id FROM select_data WHERE val < 5) RETURNING id, val
)
SELECT * FROM raw_data ORDER BY val;
 id | val
---------------------------------------------------------------------
  2 |   0
(1 row)

SELECT * FROM modify_table ORDER BY id, val;
 id | val
---------------------------------------------------------------------
  2 |   0
  3 |   5
(2 rows)

-- Test with joins
WITH select_data AS (
	SELECT * FROM modify_table
),
raw_data AS (
	UPDATE modify_table SET val = 0 WHERE
		id IN (SELECT id FROM select_data) AND
		val IN (SELECT counter FROM summary_table)
	RETURNING id, val
)
SELECT * FROM raw_data ORDER BY val;
 id | val
---------------------------------------------------------------------
(0 rows)

-- Test that local tables are barred
UPDATE local_table lt SET val = mt.val
FROM modify_table mt WHERE mt.id = lt.id;
ERROR:  cannot plan modifications of local tables involving distributed tables
-- Including inside CTEs
WITH cte AS (
	UPDATE local_table lt SET val = mt.val
	FROM modify_table mt WHERE mt.id = lt.id
	RETURNING lt.id, lt.val
) SELECT * FROM cte JOIN modify_table mt ON mt.id = cte.id ORDER BY 1,2;
ERROR:  cannot plan modifications of local tables involving distributed tables
-- Make sure checks for volatile functions apply to CTEs too
WITH cte AS (UPDATE modify_table SET val = random() WHERE id = 3 RETURNING *)
SELECT * FROM cte JOIN modify_table mt ON mt.id = 3 AND mt.id = cte.id ORDER BY 1,2;
ERROR:  functions used in UPDATE queries on distributed tables must not be VOLATILE
-- Two queries from HammerDB:
-- 1
CREATE TABLE with_modifying.stock (s_i_id numeric(6,0) NOT NULL, s_w_id numeric(4,0) NOT NULL, s_quantity numeric(6,0), s_dist_01 character(24)) WITH (fillfactor='50');
ALTER TABLE with_modifying.stock ADD CONSTRAINT stock_i1 PRIMARY KEY (s_i_id, s_w_id);
SELECT create_distributed_table('stock', 's_w_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

INSERT INTO with_modifying.stock VALUES
	(64833, 10, 3, 'test1'),
	(64834, 10, 3, 'test2'),
	(63867, 10, 3, 'test3');
PREPARE su_after(INT[], SMALLINT[], SMALLINT[], NUMERIC(5,2)[], NUMERIC, NUMERIC, NUMERIC) AS
	WITH stock_update AS (
		UPDATE stock
		SET s_quantity = ( CASE WHEN s_quantity < (item_stock.quantity + 10) THEN s_quantity + 91 ELSE s_quantity END) - item_stock.quantity
		FROM UNNEST($1, $2, $3, $4) AS item_stock (item_id, supply_wid, quantity, price)
		WHERE stock.s_i_id = item_stock.item_id
			AND stock.s_w_id = item_stock.supply_wid
			AND stock.s_w_id = ANY ($2)
		RETURNING stock.s_dist_01 as s_dist, stock.s_quantity, ( item_stock.quantity + item_stock.price * ( 1 + $5 + $6 ) * ( 1 - $7) ) amount
	)
	SELECT array_agg ( s_dist ), array_agg ( s_quantity ), array_agg ( amount )
	FROM stock_update;
EXECUTE su_after('{64833,63857,13941,76514,35858,10004,88553,34483,91251,28144,51687,36407,54436,72873}', '{10,10,10,10,10,10,10,10,10,10,10,10,10,10}', '{8,2,2,6,7,4,6,1,1,5,6,7,6,2}', '{26.04,4.79,67.84,77.66,47.06,23.12,32.74,56.99,84.75,37.52,73.52,98.86,49.96,29.47}', 0.1800, 0.1100, 0.5000);
          array_agg           | array_agg |    array_agg
---------------------------------------------------------------------
 {"test1                   "} | {86}      | {24.7958000000}
(1 row)

EXECUTE su_after('{64833,63857,13941,76514,35858,10004,88553,34483,91251,28144,51687,36407,54436,72873}', '{10,10,10,10,10,10,10,10,10,10,10,10,10,10}', '{8,2,2,6,7,4,6,1,1,5,6,7,6,2}', '{26.04,4.79,67.84,77.66,47.06,23.12,32.74,56.99,84.75,37.52,73.52,98.86,49.96,29.47}', 0.1800, 0.1100, 0.5000);
          array_agg           | array_agg |    array_agg
---------------------------------------------------------------------
 {"test1                   "} | {78}      | {24.7958000000}
(1 row)

EXECUTE su_after('{64833,63857,13941,76514,35858,10004,88553,34483,91251,28144,51687,36407,54436,72873}', '{10,10,10,10,10,10,10,10,10,10,10,10,10,10}', '{8,2,2,6,7,4,6,1,1,5,6,7,6,2}', '{26.04,4.79,67.84,77.66,47.06,23.12,32.74,56.99,84.75,37.52,73.52,98.86,49.96,29.47}', 0.1800, 0.1100, 0.5000);
          array_agg           | array_agg |    array_agg
---------------------------------------------------------------------
 {"test1                   "} | {70}      | {24.7958000000}
(1 row)

EXECUTE su_after('{64833,63857,13941,76514,35858,10004,88553,34483,91251,28144,51687,36407,54436,72873}', '{10,10,10,10,10,10,10,10,10,10,10,10,10,10}', '{8,2,2,6,7,4,6,1,1,5,6,7,6,2}', '{26.04,4.79,67.84,77.66,47.06,23.12,32.74,56.99,84.75,37.52,73.52,98.86,49.96,29.47}', 0.1800, 0.1100, 0.5000);
          array_agg           | array_agg |    array_agg
---------------------------------------------------------------------
 {"test1                   "} | {62}      | {24.7958000000}
(1 row)

EXECUTE su_after('{64833,63857,13941,76514,35858,10004,88553,34483,91251,28144,51687,36407,54436,72873}', '{10,10,10,10,10,10,10,10,10,10,10,10,10,10}', '{8,2,2,6,7,4,6,1,1,5,6,7,6,2}', '{26.04,4.79,67.84,77.66,47.06,23.12,32.74,56.99,84.75,37.52,73.52,98.86,49.96,29.47}', 0.1800, 0.1100, 0.5000);
          array_agg           | array_agg |    array_agg
---------------------------------------------------------------------
 {"test1                   "} | {54}      | {24.7958000000}
(1 row)

EXECUTE su_after('{64833,63857,13941,76514,35858,10004,88553,34483,91251,28144,51687,36407,54436,72873}', '{10,10,10,10,10,10,10,10,10,10,10,10,10,10}', '{8,2,2,6,7,4,6,1,1,5,6,7,6,2}', '{26.04,4.79,67.84,77.66,47.06,23.12,32.74,56.99,84.75,37.52,73.52,98.86,49.96,29.47}', 0.1800, 0.1100, 0.5000);
          array_agg           | array_agg |    array_agg
---------------------------------------------------------------------
 {"test1                   "} | {46}      | {24.7958000000}
(1 row)

-- 2
CREATE TABLE with_modifying.orders (o_id numeric NOT NULL, o_w_id numeric NOT NULL, o_d_id numeric NOT NULL, o_c_id numeric) WITH (fillfactor='50');
CREATE UNIQUE INDEX orders_i2 ON with_modifying.orders USING btree (o_w_id, o_d_id, o_c_id, o_id) TABLESPACE pg_default;
ALTER TABLE with_modifying.orders ADD CONSTRAINT orders_i1 PRIMARY KEY (o_w_id, o_d_id, o_id);
CREATE TABLE with_modifying.order_line (ol_w_id numeric NOT NULL, ol_d_id numeric NOT NULL, ol_o_id numeric NOT NULL, ol_number numeric NOT NULL, ol_delivery_d timestamp without time zone, ol_amount numeric) WITH (fillfactor='50');
ALTER TABLE with_modifying.order_line ADD CONSTRAINT order_line_i1 PRIMARY KEY (ol_w_id, ol_d_id, ol_o_id, ol_number);
SELECT create_distributed_table('orders', 'o_w_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('order_line', 'ol_w_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

INSERT INTO orders VALUES (1, 1, 1, 1), (2, 2, 2, 2), (3, 3, 3, 3);
INSERT INTO order_line VALUES (1, 1, 1, 10), (2, 2, 2, 20), (3, 3, 3, 30);
PREPARE olu(int,int[],int[]) AS
	WITH order_line_update AS (
		UPDATE order_line
		SET ol_delivery_d = current_timestamp
		FROM UNNEST($2, $3) AS ids(o_id, d_id)
		WHERE ol_o_id = ids.o_id
			AND ol_d_id = ids.d_id
			AND ol_w_id = $1
		RETURNING ol_d_id, ol_o_id, ol_amount
	)
	SELECT array_agg(ol_d_id), array_agg(c_id), array_agg(sum_amount)
	FROM (
		SELECT ol_d_id,
			(SELECT DISTINCT o_c_id FROM orders WHERE o_id = ol_o_id AND o_d_id = ol_d_id AND o_w_id = $1) AS c_id,
			sum(ol_amount) AS sum_amount
		FROM order_line_update
		GROUP BY ol_d_id, ol_o_id
	) AS inner_sum;
EXECUTE olu(1,ARRAY[1,2],ARRAY[1,2]);
 array_agg | array_agg | array_agg
---------------------------------------------------------------------
 {1}       | {1}       | {NULL}
(1 row)

EXECUTE olu(1,ARRAY[1,2],ARRAY[1,2]);
 array_agg | array_agg | array_agg
---------------------------------------------------------------------
 {1}       | {1}       | {NULL}
(1 row)

EXECUTE olu(1,ARRAY[1,2],ARRAY[1,2]);
 array_agg | array_agg | array_agg
---------------------------------------------------------------------
 {1}       | {1}       | {NULL}
(1 row)

EXECUTE olu(1,ARRAY[1,2],ARRAY[1,2]);
 array_agg | array_agg | array_agg
---------------------------------------------------------------------
 {1}       | {1}       | {NULL}
(1 row)

EXECUTE olu(1,ARRAY[1,2],ARRAY[1,2]);
 array_agg | array_agg | array_agg
---------------------------------------------------------------------
 {1}       | {1}       | {NULL}
(1 row)

EXECUTE olu(1,ARRAY[1,2],ARRAY[1,2]);
 array_agg | array_agg | array_agg
---------------------------------------------------------------------
 {1}       | {1}       | {NULL}
(1 row)

-- Test with replication factor 2
SET citus.shard_replication_factor to 2;
DROP TABLE modify_table;
CREATE TABLE with_modifying.modify_table (id int, val int);
SELECT create_distributed_table('modify_table', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

INSERT INTO with_modifying.modify_table SELECT user_id, value_1 FROM public.users_table;
DROP TABLE summary_table;
CREATE TABLE with_modifying.summary_table (id int, counter int);
SELECT create_distributed_table('summary_table', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT COUNT(*) FROM modify_table;
 count
---------------------------------------------------------------------
   107
(1 row)

SELECT * FROM summary_table ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
(0 rows)

WITH raw_data AS (
	DELETE FROM modify_table RETURNING *
)
INSERT INTO summary_table SELECT id, COUNT(*) AS counter FROM raw_data GROUP BY id;
SELECT COUNT(*) FROM modify_table;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT * FROM summary_table ORDER BY id, counter;
 id | counter
---------------------------------------------------------------------
  1 |       8
  2 |      19
  3 |      18
  4 |      24
  5 |      27
  6 |      11
(6 rows)

-- make sure that the intermediate result uses a connection
-- that does not interfere with placement connections
BEGIN;
	INSERT INTO modify_table (id) VALUES (10000);
	WITH test_cte AS (SELECT count(*) FROM modify_table) SELECT * FROM test_cte;
 count
---------------------------------------------------------------------
     1
(1 row)

ROLLBACK;
-- similarly, make sure that the intermediate result uses a seperate connection
WITH first_query AS (INSERT INTO modify_table (id) VALUES (10001)),
 	second_query AS (SELECT * FROM modify_table) SELECT count(*) FROM second_query;
 count
---------------------------------------------------------------------
     1
(1 row)

SET client_min_messages TO debug2;
-- pushed down without the insert
WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING NULL) INSERT INTO modify_table WITH ma AS (SELECT * FROM modify_table LIMIT 10) SELECT count(*) FROM mb;
DEBUG:  LIMIT clauses are not allowed in distributed INSERT ... SELECT queries
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
DETAIL:  distribution column value: 3
DEBUG:  Collecting INSERT ... SELECT results on coordinator
-- not pushed down due to volatile
WITH ma AS (SELECT count(*) FROM modify_table where id = 1), mu AS (WITH allref AS (SELECT random() a FROM modify_table limit 4) UPDATE modify_table SET val = 3 WHERE id = 1 AND val IN (SELECT a FROM allref) RETURNING id+1) SELECT count(*) FROM mu, ma;
DEBUG:  CTE ma is going to be inlined via distributed planning
DEBUG:  Router planner doesn't support VOLATILE functions in common table expressions.
DEBUG:  generating subplan XXX_1 for CTE mu: WITH allref AS (SELECT random() AS a FROM with_modifying.modify_table modify_table_1 LIMIT 4) UPDATE with_modifying.modify_table SET val = 3 WHERE ((id OPERATOR(pg_catalog.=) 1) AND ((val)::double precision OPERATOR(pg_catalog.=) ANY (SELECT allref.a FROM allref))) RETURNING (id OPERATOR(pg_catalog.+) 1)
DEBUG:  Router planner doesn't support VOLATILE functions in common table expressions.
DEBUG:  generating subplan XXX_1 for CTE allref: SELECT random() AS a FROM with_modifying.modify_table LIMIT 4
DEBUG:  Router planner cannot handle multi-shard select queries
DEBUG:  push down of limit count: 4
DEBUG:  Plan XXX query after replacing subqueries and CTEs: UPDATE with_modifying.modify_table SET val = 3 WHERE ((id OPERATOR(pg_catalog.=) 1) AND ((val)::double precision OPERATOR(pg_catalog.=) ANY (SELECT allref.a FROM (SELECT intermediate_result.a FROM read_intermediate_result('XXX_1'::text, 'binary'::citus_copy_format) intermediate_result(a double precision)) allref))) RETURNING (id OPERATOR(pg_catalog.+) 1)
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
DETAIL:  distribution column value: 1
DEBUG:  Distributed planning for a fast-path router query
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
DETAIL:  distribution column value: 1
DEBUG:  generating subplan XXX_2 for subquery SELECT count(*) AS count FROM with_modifying.modify_table WHERE (id OPERATOR(pg_catalog.=) 1)
DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT count(*) AS count FROM (SELECT intermediate_result."?column?" FROM read_intermediate_result('XXX_1'::text, 'binary'::citus_copy_format) intermediate_result("?column?" integer)) mu("?column?"), (SELECT intermediate_result.count FROM read_intermediate_result('XXX_2'::text, 'binary'::citus_copy_format) intermediate_result(count bigint)) ma
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
 count
---------------------------------------------------------------------
     0
(1 row)

WITH mu AS (WITH allref AS (SELECT random() a FROM anchor_table) UPDATE modify_table SET val = 3 WHERE id = 1 AND val IN (SELECT a FROM allref) RETURNING id+1) SELECT count(*) FROM mu;
DEBUG:  Router planner doesn't support VOLATILE functions in common table expressions.
DEBUG:  generating subplan XXX_1 for CTE mu: WITH allref AS (SELECT random() AS a FROM with_modifying.anchor_table) UPDATE with_modifying.modify_table SET val = 3 WHERE ((id OPERATOR(pg_catalog.=) 1) AND ((val)::double precision OPERATOR(pg_catalog.=) ANY (SELECT allref.a FROM allref))) RETURNING (id OPERATOR(pg_catalog.+) 1)
DEBUG:  Router planner doesn't support VOLATILE functions in common table expressions.
DEBUG:  generating subplan XXX_1 for CTE allref: SELECT random() AS a FROM with_modifying.anchor_table
DEBUG:  Distributed planning for a fast-path router query
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
DEBUG:  Plan XXX query after replacing subqueries and CTEs: UPDATE with_modifying.modify_table SET val = 3 WHERE ((id OPERATOR(pg_catalog.=) 1) AND ((val)::double precision OPERATOR(pg_catalog.=) ANY (SELECT allref.a FROM (SELECT intermediate_result.a FROM read_intermediate_result('XXX_1'::text, 'binary'::citus_copy_format) intermediate_result(a double precision)) allref))) RETURNING (id OPERATOR(pg_catalog.+) 1)
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
DETAIL:  distribution column value: 1
DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT count(*) AS count FROM (SELECT intermediate_result."?column?" FROM read_intermediate_result('XXX_1'::text, 'binary'::citus_copy_format) intermediate_result("?column?" integer)) mu("?column?")
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
 count
---------------------------------------------------------------------
     0
(1 row)

-- pushed down
WITH mu AS (WITH allref AS (SELECT id a FROM anchor_table) UPDATE modify_table SET val = 3 WHERE id = 1 AND val IN (SELECT a FROM allref) RETURNING id+1) SELECT count(*) FROM mu;
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
DETAIL:  distribution column value: 1
 count
---------------------------------------------------------------------
     0
(1 row)

-- pushed down and stable function evaluated
WITH mu AS (WITH allref AS (SELECT now() a FROM anchor_table) UPDATE modify_table SET val = 3 WHERE id = 1 AND now() IN (SELECT a FROM allref) RETURNING id+1) SELECT count(*) FROM mu;
DEBUG:  Creating router plan
DEBUG:  Plan is router executable
DETAIL:  distribution column value: 1
 count
---------------------------------------------------------------------
     0
(1 row)

RESET client_min_messages;
-- https://github.com/citusdata/citus/issues/3975
WITH mb AS (INSERT INTO modify_table VALUES (3, 3) RETURNING NULL, NULL) SELECT * FROM modify_table WHERE id = 3;
 id | val
---------------------------------------------------------------------
  3 |   3
(1 row)

WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING NULL) SELECT * FROM modify_table WHERE id = 3;
 id | val
---------------------------------------------------------------------
  3 |   3
(1 row)

WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING NULL) SELECT * FROM modify_table, mb WHERE id = 3;
 id | val | ?column?
---------------------------------------------------------------------
  3 |   3 |
(1 row)

WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING NULL, NULL) SELECT * FROM modify_table WHERE id = 3;
 id | val
---------------------------------------------------------------------
  3 |   3
(1 row)

WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING NULL, NULL) SELECT * FROM modify_table, mb WHERE id = 3;
 id | val | ?column? | ?column?
---------------------------------------------------------------------
  3 |   3 |          |
(1 row)

WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING NULL alias) SELECT * FROM modify_table WHERE id = 3;
 id | val
---------------------------------------------------------------------
  3 |   3
(1 row)

WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING NULL alias) SELECT * FROM modify_table, mb WHERE id = 3;
 id | val | alias
---------------------------------------------------------------------
  3 |   3 |
(1 row)

WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING val) SELECT * FROM modify_table WHERE id = 3;
 id | val
---------------------------------------------------------------------
  3 |   3
(1 row)

WITH mb AS (UPDATE modify_table SET val = 3 WHERE id = 3 RETURNING val) SELECT * FROM modify_table, mb WHERE id = 3;
 id | val | val
---------------------------------------------------------------------
  3 |   3 |   3
(1 row)

WITH mb AS (DELETE FROM modify_table WHERE id = 3 RETURNING NULL, NULL) SELECT * FROM modify_table WHERE id = 3;
 id | val
---------------------------------------------------------------------
  3 |   3
(1 row)

\set VERBOSITY terse
DROP SCHEMA with_modifying CASCADE;
NOTICE:  drop cascades to 9 other objects