SELECT citus.mitmproxy('conn.allow()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

SET citus.shard_count = 2;
SET citus.shard_replication_factor = 1; -- one shard per worker
SET citus.next_shard_id TO 103400;
ALTER SEQUENCE pg_catalog.pg_dist_placement_placementid_seq RESTART 100;
CREATE TABLE dml_test (id integer, name text);
SELECT create_distributed_table('dml_test', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

COPY dml_test FROM STDIN WITH CSV;
SELECT citus.clear_network_traffic();
 clear_network_traffic
---------------------------------------------------------------------

(1 row)

---- test multiple statements spanning multiple shards,
---- at each significant point. These transactions are 2pc
-- fail at DELETE
SELECT citus.mitmproxy('conn.onQuery(query="DELETE").kill()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

BEGIN;
DELETE FROM dml_test WHERE id = 1;
ERROR:  connection to the remote node postgres@localhost:xxxxx failed with the following error: connection not open
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)

-- 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
---------------------------------------------------------------------

(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:  connection to the remote node postgres@localhost:xxxxx failed with the following error: connection not open
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)

-- 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
---------------------------------------------------------------------

(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:  connection to the remote node postgres@localhost:xxxxx failed with the following error: connection not open
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)

-- 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
---------------------------------------------------------------------

(1 row)

-- this transaction block will be sent to the coordinator as a remote command to hide the
-- error message that is caused during commit.
-- we'll test for the txn side-effects to ensure it didn't run
SELECT master_run_on_worker(
    ARRAY['localhost']::text[],
    ARRAY[:master_port]::int[],
    ARRAY['
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;
    '],
    false
);
WARNING:  connection not open
CONTEXT:  while executing command on localhost:xxxxx
while executing command on localhost:xxxxx
   master_run_on_worker
---------------------------------------------------------------------
 (localhost,57636,t,BEGIN)
(1 row)

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)

-- cancel at PREPARE TRANSACTION
SELECT citus.mitmproxy('conn.onQuery(query="^PREPARE TRANSACTION").cancel(' ||  pg_backend_pid() || ')');
 mitmproxy
---------------------------------------------------------------------

(1 row)

-- we'll test for the txn side-effects to ensure it didn't run
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
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
---------------------------------------------------------------------

(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 ERROR;
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
---------------------------------------------------------------------
                             1
(1 row)

-- should see changes, because of txn recovery
SELECT * FROM dml_test ORDER BY id ASC;
 id |  name
---------------------------------------------------------------------
  3 | gamma
  4 | Delta
  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;
SET citus.shard_replication_factor = 2; -- two placements
CREATE TABLE dml_test (id integer, name text);
SELECT create_distributed_table('dml_test', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

COPY dml_test FROM STDIN WITH CSV;
---- test multiple statements against a single shard, but with two placements
-- fail at PREPARED COMMIT as we use 2PC
SELECT citus.mitmproxy('conn.onQuery(query="^COMMIT").kill()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO ERROR;
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;
RESET client_min_messages;
-- all changes should be committed because we injected
-- the failure on the COMMIT time. And, we should not
-- mark any placements as INVALID
SELECT citus.mitmproxy('conn.allow()');
 mitmproxy
---------------------------------------------------------------------

(1 row)

SELECT recover_prepared_transactions();
 recover_prepared_transactions
---------------------------------------------------------------------
                             1
(1 row)

SELECT shardid FROM pg_dist_shard_placement WHERE shardstate = 3;
 shardid
---------------------------------------------------------------------
(0 rows)

SET citus.task_assignment_policy TO "round-robin";
SELECT * FROM dml_test ORDER BY id ASC;
 id |  name
---------------------------------------------------------------------
  3 | gamma
  4 | Delta
  5 | Epsilon
(3 rows)

SELECT * FROM dml_test ORDER BY id ASC;
 id |  name
---------------------------------------------------------------------
  3 | gamma
  4 | Delta
  5 | Epsilon
(3 rows)

RESET citus.task_assignment_policy;
-- drop table and recreate as reference table
DROP TABLE dml_test;
SET citus.shard_count = 2;
SET citus.shard_replication_factor = 1;
CREATE TABLE dml_test (id integer, name text);
SELECT create_reference_table('dml_test');
 create_reference_table
---------------------------------------------------------------------

(1 row)

COPY dml_test FROM STDIN WITH CSV;
-- fail at COMMIT (by failing to PREPARE)
SELECT citus.mitmproxy('conn.onQuery(query="^PREPARE").kill()');
 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:  connection not open
CONTEXT:  while executing command on localhost:xxxxx
--- shouldn't see any changes after failed COMMIT
SELECT * FROM dml_test ORDER BY id ASC;
 id | name
---------------------------------------------------------------------
  1 | Alpha
  2 | Beta
  3 | Gamma
  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
---------------------------------------------------------------------

(1 row)

DROP TABLE dml_test;
