1-- this test file aims to test UPSERT feature on Citus
2
3
4SET citus.next_shard_id TO 980000;
5
6
7CREATE TABLE upsert_test
8(
9	part_key int UNIQUE,
10	other_col int,
11	third_col int
12);
13
14-- distribute the table and create shards
15SELECT create_distributed_table('upsert_test', 'part_key', 'hash');
16
17-- do a regular insert
18INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1), (2, 2);
19
20-- observe that there is a conflict and the following query does nothing
21INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT DO NOTHING;
22
23-- same as the above with different syntax
24INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key) DO NOTHING;
25
26--again the same query with another syntax
27INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT ON CONSTRAINT upsert_test_part_key_key DO NOTHING;
28
29-- now, update the columns
30INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1)
31	ON CONFLICT (part_key) DO UPDATE SET other_col = 2,  third_col = 4;
32
33-- see the results
34SELECT * FROM upsert_test ORDER BY part_key ASC;
35
36-- do a multi-row DO NOTHING insert
37INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1), (2, 2)
38ON CONFLICT DO NOTHING;
39
40-- do a multi-row DO UPDATE insert
41INSERT INTO upsert_test (part_key, other_col) VALUES (1, 10), (2, 20), (3, 30)
42ON CONFLICT (part_key) DO
43UPDATE SET other_col = EXCLUDED.other_col WHERE upsert_test.part_key != 1;
44
45-- see the results
46SELECT * FROM upsert_test ORDER BY part_key ASC;
47
48DELETE FROM upsert_test WHERE part_key = 2;
49DELETE FROM upsert_test WHERE part_key = 3;
50
51-- use a WHERE clause, so that SET doesn't have an affect
52INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key)
53	DO UPDATE SET other_col = 30 WHERE upsert_test.other_col = 3;
54
55-- see the results
56SELECT * FROM upsert_test;
57
58-- use a WHERE clause, that hits the row and updates it
59INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key)
60	DO UPDATE SET other_col = 30 WHERE upsert_test.other_col = 2;
61
62-- see the results
63SELECT * FROM upsert_test;
64
65-- use two elements in the WHERE, that doesn't hit the row and updates it
66INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key)
67	DO UPDATE SET other_col = 30 WHERE upsert_test.other_col = 2 AND upsert_test.other_col = 3;
68
69-- use EXCLUDED keyword
70INSERT INTO upsert_test (part_key, other_col, third_col) VALUES (1, 1, 100) ON CONFLICT (part_key)
71	DO UPDATE SET other_col = EXCLUDED.third_col;
72
73-- see the results
74SELECT * FROM upsert_test;
75
76-- now update multiple columns with ALIAS table and reference to the row itself
77INSERT INTO upsert_test as ups_test (part_key) VALUES (1)
78	ON CONFLICT (part_key) DO UPDATE SET other_col = ups_test.other_col + 50, third_col = 200;
79
80-- see the results
81SELECT * FROM upsert_test;
82
83-- now, do some more complex assignments
84INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key)
85		DO UPDATE SET other_col = upsert_test.other_col + 1,
86		third_col = upsert_test.third_col + (EXCLUDED.part_key + EXCLUDED.other_col) + 670;
87
88-- see the results
89SELECT * FROM upsert_test;
90
91-- now, WHERE clause also has table reference
92INSERT INTO upsert_test as ups_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key)
93		DO UPDATE SET other_col = (ups_test.other_col + ups_test.third_col + (EXCLUDED.part_key + EXCLUDED.other_col)) % 15
94		WHERE ups_test.third_col < 1000 + ups_test.other_col;
95
96-- see the results
97SELECT * FROM upsert_test;
98
99-- Test upsert, with returning:
100INSERT INTO upsert_test (part_key, other_col) VALUES (2, 2)
101	ON CONFLICT (part_key) DO UPDATE SET other_col = 3
102        RETURNING *;
103
104INSERT INTO upsert_test (part_key, other_col) VALUES (2, 2)
105	ON CONFLICT (part_key) DO UPDATE SET other_col = 3
106        RETURNING *;
107
108-- create another table
109CREATE TABLE upsert_test_2
110(
111	part_key int,
112	other_col int,
113	third_col int,
114	PRIMARY KEY (part_key, other_col)
115);
116
117-- distribute the table and create shards
118SELECT create_distributed_table('upsert_test_2', 'part_key', 'hash');
119
120-- now show that Citus works with multiple columns as the PRIMARY KEY, including the partiton key
121INSERT INTO upsert_test_2 (part_key, other_col) VALUES (1, 1);
122INSERT INTO upsert_test_2 (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key, other_col) DO NOTHING;
123
124-- this errors out since there is no unique constraint on partition key
125INSERT INTO upsert_test_2 (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key) DO NOTHING;
126
127-- create another table
128CREATE TABLE upsert_test_3
129(
130	part_key int,
131	count int
132);
133
134-- note that this is not a unique index
135CREATE INDEX idx_ups_test ON upsert_test_3(part_key);
136
137-- distribute the table and create shards
138SELECT create_distributed_table('upsert_test_3', 'part_key', 'hash');
139
140-- since there are no unique indexes, error-out
141INSERT INTO upsert_test_3 VALUES (1, 0) ON CONFLICT(part_key) DO UPDATE SET count = upsert_test_3.count + 1;
142
143-- create another table
144CREATE TABLE upsert_test_4
145(
146	part_key int UNIQUE,
147	count int
148);
149
150-- distribute the table and create shards
151SELECT create_distributed_table('upsert_test_4', 'part_key', 'hash');
152
153-- a single row insert
154INSERT INTO upsert_test_4 VALUES (1, 0);
155
156-- show a simple count example use case
157INSERT INTO upsert_test_4 VALUES (1, 0) ON CONFLICT(part_key) DO UPDATE SET count = upsert_test_4.count + 1;
158INSERT INTO upsert_test_4 VALUES (1, 0) ON CONFLICT(part_key) DO UPDATE SET count = upsert_test_4.count + 1;
159INSERT INTO upsert_test_4 VALUES (1, 0) ON CONFLICT(part_key) DO UPDATE SET count = upsert_test_4.count + 1;
160INSERT INTO upsert_test_4 VALUES (1, 0) ON CONFLICT(part_key) DO UPDATE SET count = upsert_test_4.count + 1;
161INSERT INTO upsert_test_4 VALUES (1, 0) ON CONFLICT(part_key) DO UPDATE SET count = upsert_test_4.count + 1;
162INSERT INTO upsert_test_4 VALUES (1, 0) ON CONFLICT(part_key) DO UPDATE SET count = upsert_test_4.count + 1;
163
164-- now see the results
165SELECT * FROM upsert_test_4;
166
167-- now test dropped columns
168SET citus.shard_replication_factor TO 1;
169CREATE TABLE dropcol_distributed(key int primary key, drop1 int, keep1 text, drop2 numeric, keep2 float);
170SELECT create_distributed_table('dropcol_distributed', 'key', 'hash');
171
172INSERT INTO dropcol_distributed AS dropcol (key, keep1, keep2) VALUES (1, '5', 5) ON CONFLICT(key)
173	DO UPDATE SET keep1 = dropcol.keep1;
174
175ALTER TABLE dropcol_distributed DROP COLUMN drop2;
176
177INSERT INTO dropcol_distributed (key, keep1, keep2) VALUES (1, '5', 5) ON CONFLICT(key)
178	DO UPDATE SET keep1 = dropcol_distributed.keep1;
179
180ALTER TABLE dropcol_distributed DROP COLUMN keep2;
181
182INSERT INTO dropcol_distributed AS dropcol (key, keep1) VALUES (1, '5') ON CONFLICT(key)
183	DO UPDATE SET keep1 = dropcol.keep1;
184
185ALTER TABLE dropcol_distributed DROP COLUMN drop1;
186
187INSERT INTO dropcol_distributed AS dropcol (key, keep1) VALUES (1, '5') ON CONFLICT(key)
188	DO UPDATE SET keep1 = dropcol.keep1;
189
190-- below we test the cases that Citus does not support
191-- subquery in the SET clause
192INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key) DO
193	UPDATE SET other_col = (SELECT count(*) from upsert_test);
194
195-- non mutable function call in the SET
196INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key) DO
197	UPDATE SET other_col = random()::int;
198
199-- non mutable function call in the WHERE
200INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key) DO
201	UPDATE SET other_col = 5 WHERE upsert_test.other_col = random()::int;
202
203-- non mutable function call in the arbiter WHERE
204INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key)  WHERE part_key = random()::int
205	DO UPDATE SET other_col = 5;
206
207-- error out on attempt to update the partition key
208INSERT INTO upsert_test (part_key, other_col) VALUES (1, 1) ON CONFLICT (part_key) DO
209	UPDATE SET part_key = 15;
210