Merge pull request #2461 from citusdata/cancellation_multi_shard

Add cancellation tests for multi shard modification queries

src/test/regress/expected/failure_multi_dml.out

@@ -56,6 +56,35 @@ SELECT * FROM dml_test ORDER BY id ASC;
   4 | Delta
 (4 rows)
 
+-- cancel at DELETE
+SELECT citus.mitmproxy('conn.onQuery(query="^DELETE").cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+ERROR:  canceling statement due to user request
+DELETE FROM dml_test WHERE id = 2;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+COMMIT;
+--- shouldn't see any changes performed in failed transaction
+SELECT * FROM dml_test ORDER BY id ASC;
+ id | name  
+----+-------
+  1 | Alpha
+  2 | Beta
+  3 | Gamma
+  4 | Delta
+(4 rows)
+
 -- fail at INSERT
 SELECT citus.mitmproxy('conn.onQuery(query="^INSERT").kill()');
  mitmproxy 
@@ -86,6 +115,33 @@ SELECT * FROM dml_test ORDER BY id ASC;
   4 | Delta
 (4 rows)
 
+-- cancel at INSERT
+SELECT citus.mitmproxy('conn.onQuery(query="^INSERT").cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+ERROR:  canceling statement due to user request
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+COMMIT;
+--- shouldn't see any changes before failed INSERT
+SELECT * FROM dml_test ORDER BY id ASC;
+ id | name  
+----+-------
+  1 | Alpha
+  2 | Beta
+  3 | Gamma
+  4 | Delta
+(4 rows)
+
 -- fail at UPDATE
 SELECT citus.mitmproxy('conn.onQuery(query="^UPDATE").kill()');
  mitmproxy 
@@ -115,6 +171,32 @@ SELECT * FROM dml_test ORDER BY id ASC;
   4 | Delta
 (4 rows)
 
+-- cancel at UPDATE
+SELECT citus.mitmproxy('conn.onQuery(query="^UPDATE").cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+ERROR:  canceling statement due to user request
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+COMMIT;
+--- shouldn't see any changes after failed UPDATE
+SELECT * FROM dml_test ORDER BY id ASC;
+ id | name  
+----+-------
+  1 | Alpha
+  2 | Beta
+  3 | Gamma
+  4 | Delta
+(4 rows)
+
 -- fail at PREPARE TRANSACTION
 SELECT citus.mitmproxy('conn.onQuery(query="^PREPARE TRANSACTION").kill()');
  mitmproxy 
@@ -160,6 +242,51 @@ SELECT * FROM dml_test ORDER BY id ASC;
   4 | Delta
 (4 rows)
 
+-- cancel at PREPARE TRANSACTION
+SELECT citus.mitmproxy('conn.onQuery(query="^PREPARE TRANSACTION").cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+-- hide the error message (it has the PID)...
+-- we'll test for the txn side-effects to ensure it didn't run
+SET client_min_messages TO FATAL;
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+SET client_min_messages TO DEFAULT;
+SELECT citus.mitmproxy('conn.allow()');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+SELECT shardid FROM pg_dist_shard_placement WHERE shardstate = 3;
+ shardid 
+---------
+(0 rows)
+
+SELECT recover_prepared_transactions();
+ recover_prepared_transactions 
+-------------------------------
+                             0
+(1 row)
+
+-- shouldn't see any changes after failed PREPARE
+SELECT * FROM dml_test ORDER BY id ASC;
+ id | name  
+----+-------
+  1 | Alpha
+  2 | Beta
+  3 | Gamma
+  4 | Delta
+(4 rows)
+
 -- fail at COMMIT
 SELECT citus.mitmproxy('conn.onQuery(query="^COMMIT").kill()');
  mitmproxy 
@@ -204,6 +331,30 @@ SELECT * FROM dml_test ORDER BY id ASC;
   5 | Epsilon
 (3 rows)
 
+-- cancel at COMMITs are ignored by Postgres
+SELECT citus.mitmproxy('conn.onQuery(query="^COMMIT").cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+-- should see changes, because cancellation is ignored
+SELECT * FROM dml_test ORDER BY id ASC;
+ id |  name   
+----+---------
+  3 | gamma
+  4 | Delta
+  5 | Epsilon
+  5 | Epsilon
+(4 rows)
+
 -- drop table and recreate with different replication/sharding
 DROP TABLE dml_test;
 SET citus.shard_count = 1;
@@ -295,6 +446,32 @@ SELECT * FROM dml_test ORDER BY id ASC;
   4 | Delta
 (4 rows)
 
+-- cancel at COMMIT (by cancelling on PREPARE)
+SELECT citus.mitmproxy('conn.onQuery(query="^PREPARE").cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+ERROR:  canceling statement due to user request
+--- shouldn't see any changes after cancelled PREPARE
+SELECT * FROM dml_test ORDER BY id ASC;
+ id | name  
+----+-------
+  1 | Alpha
+  2 | Beta
+  3 | Gamma
+  4 | Delta
+(4 rows)
+
+-- allow connection to allow DROP
 SELECT citus.mitmproxy('conn.allow()');
  mitmproxy 
 -----------

src/test/regress/sql/failure_multi_dml.sql

@@ -34,6 +34,20 @@ COMMIT;
 --- shouldn't see any changes performed in failed transaction
 SELECT * FROM dml_test ORDER BY id ASC;
 
+-- cancel at DELETE
+SELECT citus.mitmproxy('conn.onQuery(query="^DELETE").cancel(' ||  pg_backend_pid() || ')');
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+
+--- shouldn't see any changes performed in failed transaction
+SELECT * FROM dml_test ORDER BY id ASC;
+
 -- fail at INSERT
 SELECT citus.mitmproxy('conn.onQuery(query="^INSERT").kill()');
 
@@ -48,6 +62,20 @@ COMMIT;
 --- shouldn't see any changes before failed INSERT
 SELECT * FROM dml_test ORDER BY id ASC;
 
+-- cancel at INSERT
+SELECT citus.mitmproxy('conn.onQuery(query="^INSERT").cancel(' ||  pg_backend_pid() || ')');
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+
+--- shouldn't see any changes before failed INSERT
+SELECT * FROM dml_test ORDER BY id ASC;
+
 -- fail at UPDATE
 SELECT citus.mitmproxy('conn.onQuery(query="^UPDATE").kill()');
 
@@ -62,6 +90,20 @@ COMMIT;
 --- shouldn't see any changes after failed UPDATE
 SELECT * FROM dml_test ORDER BY id ASC;
 
+-- cancel at UPDATE
+SELECT citus.mitmproxy('conn.onQuery(query="^UPDATE").cancel(' ||  pg_backend_pid() || ')');
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+
+--- shouldn't see any changes after failed UPDATE
+SELECT * FROM dml_test ORDER BY id ASC;
+
 -- fail at PREPARE TRANSACTION
 SELECT citus.mitmproxy('conn.onQuery(query="^PREPARE TRANSACTION").kill()');
 
@@ -86,6 +128,31 @@ SELECT recover_prepared_transactions();
 -- shouldn't see any changes after failed PREPARE
 SELECT * FROM dml_test ORDER BY id ASC;
 
+
+-- cancel at PREPARE TRANSACTION
+SELECT citus.mitmproxy('conn.onQuery(query="^PREPARE TRANSACTION").cancel(' ||  pg_backend_pid() || ')');
+
+-- hide the error message (it has the PID)...
+-- we'll test for the txn side-effects to ensure it didn't run
+SET client_min_messages TO FATAL;
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+
+SET client_min_messages TO DEFAULT;
+
+SELECT citus.mitmproxy('conn.allow()');
+SELECT shardid FROM pg_dist_shard_placement WHERE shardstate = 3;
+SELECT recover_prepared_transactions();
+
+-- shouldn't see any changes after failed PREPARE
+SELECT * FROM dml_test ORDER BY id ASC;
+
 -- fail at COMMIT
 SELECT citus.mitmproxy('conn.onQuery(query="^COMMIT").kill()');
 
@@ -110,6 +177,22 @@ SELECT recover_prepared_transactions();
 -- should see changes, because of txn recovery
 SELECT * FROM dml_test ORDER BY id ASC;
 
+
+-- cancel at COMMITs are ignored by Postgres
+SELECT citus.mitmproxy('conn.onQuery(query="^COMMIT").cancel(' ||  pg_backend_pid() || ')');
+
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+
+-- should see changes, because cancellation is ignored
+SELECT * FROM dml_test ORDER BY id ASC;
+
 -- drop table and recreate with different replication/sharding
 
 DROP TABLE dml_test;
@@ -175,5 +258,20 @@ COMMIT;
 --- shouldn't see any changes after failed COMMIT
 SELECT * FROM dml_test ORDER BY id ASC;
 
+-- cancel at COMMIT (by cancelling on PREPARE)
+SELECT citus.mitmproxy('conn.onQuery(query="^PREPARE").cancel(' ||  pg_backend_pid() || ')');
+
+BEGIN;
+DELETE FROM dml_test WHERE id = 1;
+DELETE FROM dml_test WHERE id = 2;
+INSERT INTO dml_test VALUES (5, 'Epsilon');
+UPDATE dml_test SET name = 'alpha' WHERE id = 1;
+UPDATE dml_test SET name = 'gamma' WHERE id = 3;
+COMMIT;
+
+--- shouldn't see any changes after cancelled PREPARE
+SELECT * FROM dml_test ORDER BY id ASC;
+
+-- allow connection to allow DROP
 SELECT citus.mitmproxy('conn.allow()');
 DROP TABLE dml_test;