Add cancellation tests for router selects

src/test/regress/expected/failure_single_select.out

@@ -89,6 +89,62 @@ WHERE shardid IN (
   SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'select_test'::regclass
 );
 TRUNCATE select_test;
+-- now the same tests with query cancellation
+-- put data in shard for which mitm node is first placement
+INSERT INTO select_test VALUES (2, 'test data');
+SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+SELECT * FROM select_test WHERE key = 2;
+ERROR:  canceling statement due to user request
+SELECT * FROM select_test WHERE key = 2;
+ERROR:  canceling statement due to user request
+-- cancel after first SELECT; txn should fail and nothing should be marked as invalid
+SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+BEGIN;
+INSERT INTO select_test VALUES (2, 'more data');
+SELECT * FROM select_test WHERE key = 2;
+ERROR:  canceling statement due to user request
+COMMIT;
+-- show that all placements are OK
+SELECT DISTINCT shardstate FROM  pg_dist_shard_placement
+WHERE shardid IN (
+  SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'select_test'::regclass
+);
+ shardstate 
+------------
+          1
+(1 row)
+
+TRUNCATE select_test;
+-- cancel the second query
+-- error after second SELECT; txn should fail
+SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").after(1).cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+BEGIN;
+INSERT INTO select_test VALUES (2, 'more data');
+SELECT * FROM select_test WHERE key = 2;
+ key |   value   
+-----+-----------
+   2 | more data
+(1 row)
+
+INSERT INTO select_test VALUES (2, 'even more data');
+SELECT * FROM select_test WHERE key = 2;
+ERROR:  canceling statement due to user request
+COMMIT;
 -- error after second SELECT; txn should work (though placement marked bad)
 SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").after(1).reset()');
  mitmproxy 
@@ -170,5 +226,20 @@ WARNING:  server closed the connection unexpectedly
 	before or while processing the request.
 CONTEXT:  while executing command on localhost:9060
 ERROR:  could not receive query results
+-- now the same test with query cancellation
+SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").after(1).cancel(' ||  pg_backend_pid() || ')');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+SELECT * FROM select_test WHERE key = 1;
+ key |   value   
+-----+-----------
+   1 | test data
+(1 row)
+
+SELECT * FROM select_test WHERE key = 1;
+ERROR:  canceling statement due to user request
 -- ==== Clean up, we're done here ====
 DROP TABLE select_test;

src/test/regress/sql/failure_single_select.sql

@@ -31,6 +31,41 @@ WHERE shardid IN (
 );
 TRUNCATE select_test;
 
+-- now the same tests with query cancellation
+
+-- put data in shard for which mitm node is first placement
+INSERT INTO select_test VALUES (2, 'test data');
+
+SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").cancel(' ||  pg_backend_pid() || ')');
+SELECT * FROM select_test WHERE key = 2;
+SELECT * FROM select_test WHERE key = 2;
+
+-- cancel after first SELECT; txn should fail and nothing should be marked as invalid
+SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").cancel(' ||  pg_backend_pid() || ')');
+
+BEGIN;
+INSERT INTO select_test VALUES (2, 'more data');
+SELECT * FROM select_test WHERE key = 2;
+COMMIT;
+
+-- show that all placements are OK
+SELECT DISTINCT shardstate FROM  pg_dist_shard_placement
+WHERE shardid IN (
+  SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'select_test'::regclass
+);
+TRUNCATE select_test;
+
+-- cancel the second query
+-- error after second SELECT; txn should fail
+SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").after(1).cancel(' ||  pg_backend_pid() || ')');
+
+BEGIN;
+INSERT INTO select_test VALUES (2, 'more data');
+SELECT * FROM select_test WHERE key = 2;
+INSERT INTO select_test VALUES (2, 'even more data');
+SELECT * FROM select_test WHERE key = 2;
+COMMIT;
+
 -- error after second SELECT; txn should work (though placement marked bad)
 SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").after(1).reset()');
 
@@ -58,6 +93,11 @@ SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").after(1).kill()');
 SELECT * FROM select_test WHERE key = 1;
 SELECT * FROM select_test WHERE key = 1;
 
+-- now the same test with query cancellation
+SELECT citus.mitmproxy('conn.onQuery(query="^SELECT").after(1).cancel(' ||  pg_backend_pid() || ')');
+SELECT * FROM select_test WHERE key = 1;
+SELECT * FROM select_test WHERE key = 1;
+
 -- ==== Clean up, we're done here ====
 
 DROP TABLE select_test;