diff --git a/src/backend/distributed/commands/multi_copy.c b/src/backend/distributed/commands/multi_copy.c index 9c439f831..6cc0c5dca 100644 --- a/src/backend/distributed/commands/multi_copy.c +++ b/src/backend/distributed/commands/multi_copy.c @@ -2262,6 +2262,10 @@ CitusCopyDestReceiverReceive(TupleTableSlot *slot, DestReceiver *dest) } PG_CATCH(); { + /* + * We might be able to recover from errors with ROLLBACK TO SAVEPOINT, + * so unclaim the connections before throwing errors. + */ HTAB *shardConnectionHash = copyDest->shardConnectionHash; UnclaimAllShardConnections(shardConnectionHash); diff --git a/src/backend/distributed/connection/remote_commands.c b/src/backend/distributed/connection/remote_commands.c index 8fe174251..21d8bd56e 100644 --- a/src/backend/distributed/connection/remote_commands.c +++ b/src/backend/distributed/connection/remote_commands.c @@ -1038,8 +1038,8 @@ BuildWaitEventSet(MultiConnection **allConnections, int totalConnectionCount, /* - * MultiClientCancel sends a cancelation request on the given connection. Return - * value indicates whether the cancelation request was sent successfully. + * SendCancelationRequest sends a cancelation request on the given connection. + * Return value indicates whether the cancelation request was sent successfully. */ bool SendCancelationRequest(MultiConnection *connection) diff --git a/src/backend/distributed/executor/multi_router_executor.c b/src/backend/distributed/executor/multi_router_executor.c index 8bdb8f4f5..c5259cc6f 100644 --- a/src/backend/distributed/executor/multi_router_executor.c +++ b/src/backend/distributed/executor/multi_router_executor.c @@ -1640,6 +1640,10 @@ ExecuteModifyTasks(List *taskList, bool expectResults, ParamListInfo paramListIn } PG_CATCH(); { + /* + * We might be able to recover from errors with ROLLBACK TO SAVEPOINT, + * so unclaim the connections before throwing errors. + */ UnclaimAllShardConnections(shardConnectionHash); PG_RE_THROW(); } diff --git a/src/test/regress/expected/multi_subtransactions.out b/src/test/regress/expected/multi_subtransactions.out index d8020bc8d..7a6812dce 100644 --- a/src/test/regress/expected/multi_subtransactions.out +++ b/src/test/regress/expected/multi_subtransactions.out @@ -193,6 +193,76 @@ SELECT * FROM artists WHERE id IN (11, 12) ORDER BY id; 12 | Marc Chagall (1 row) +-- Recover from multi-shard CTE modify failures +create table t1(a int, b int); +create table t2(a int, b int CHECK(b > 0)); +ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 1190000; +select create_distributed_table('t1', 'a'), + create_distributed_table('t2', 'a'); + create_distributed_table | create_distributed_table +--------------------------+-------------------------- + | +(1 row) + +begin; +insert into t2 select i, i+1 from generate_series(1, 3) i; +with r AS ( + update t2 set b = b + 1 + returning * +) insert into t1 select * from r; +savepoint s1; +with r AS ( + update t1 set b = b - 10 + returning * +) insert into t2 select * from r; +ERROR: new row for relation "t2_1190005" violates check constraint "t2_b_check" +rollback to savepoint s1; +savepoint s2; +with r AS ( + update t2 set b = b - 10 + returning * +) insert into t1 select * from r; +ERROR: new row for relation "t2_1190004" violates check constraint "t2_b_check" +rollback to savepoint s2; +savepoint s3; +with r AS ( + insert into t2 select i, i+1 from generate_series(-10,-5) i + returning * +) insert into t1 select * from r; +ERROR: new row for relation "t2_1190004" violates check constraint "t2_b_check" +rollback to savepoint s3; +savepoint s4; +with r AS ( + insert into t1 select i, i+1 from generate_series(-10,-5) i + returning * +) insert into t2 select * from r; +ERROR: new row for relation "t2_1190005" violates check constraint "t2_b_check" +rollback to savepoint s4; +with r AS ( + update t2 set b = b + 1 + returning * +) insert into t1 select * from r; +commit; +select * from t2 order by a, b; + a | b +---+--- + 1 | 4 + 2 | 5 + 3 | 6 +(3 rows) + +select * from t1 order by a, b; + a | b +---+--- + 1 | 3 + 1 | 4 + 2 | 4 + 2 | 5 + 3 | 5 + 3 | 6 +(6 rows) + +drop table t1, t2; -- =================================================================== -- Tests for replication factor > 1 -- =================================================================== diff --git a/src/test/regress/expected/multi_subtransactions_0.out b/src/test/regress/expected/multi_subtransactions_0.out index 9caa09e59..3f7f8aa2c 100644 --- a/src/test/regress/expected/multi_subtransactions_0.out +++ b/src/test/regress/expected/multi_subtransactions_0.out @@ -193,6 +193,76 @@ SELECT * FROM artists WHERE id IN (11, 12) ORDER BY id; 12 | Marc Chagall (1 row) +-- Recover from multi-shard CTE modify failures +create table t1(a int, b int); +create table t2(a int, b int CHECK(b > 0)); +ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 1190000; +select create_distributed_table('t1', 'a'), + create_distributed_table('t2', 'a'); + create_distributed_table | create_distributed_table +--------------------------+-------------------------- + | +(1 row) + +begin; +insert into t2 select i, i+1 from generate_series(1, 3) i; +with r AS ( + update t2 set b = b + 1 + returning * +) insert into t1 select * from r; +savepoint s1; +with r AS ( + update t1 set b = b - 10 + returning * +) insert into t2 select * from r; +ERROR: new row for relation "t2_1190005" violates check constraint "t2_b_check" +rollback to savepoint s1; +savepoint s2; +with r AS ( + update t2 set b = b - 10 + returning * +) insert into t1 select * from r; +ERROR: new row for relation "t2_1190004" violates check constraint "t2_b_check" +rollback to savepoint s2; +savepoint s3; +with r AS ( + insert into t2 select i, i+1 from generate_series(-10,-5) i + returning * +) insert into t1 select * from r; +ERROR: new row for relation "t2_1190004" violates check constraint "t2_b_check" +rollback to savepoint s3; +savepoint s4; +with r AS ( + insert into t1 select i, i+1 from generate_series(-10,-5) i + returning * +) insert into t2 select * from r; +ERROR: new row for relation "t2_1190005" violates check constraint "t2_b_check" +rollback to savepoint s4; +with r AS ( + update t2 set b = b + 1 + returning * +) insert into t1 select * from r; +commit; +select * from t2 order by a, b; + a | b +---+--- + 1 | 4 + 2 | 5 + 3 | 6 +(3 rows) + +select * from t1 order by a, b; + a | b +---+--- + 1 | 3 + 1 | 4 + 2 | 4 + 2 | 5 + 3 | 5 + 3 | 6 +(6 rows) + +drop table t1, t2; -- =================================================================== -- Tests for replication factor > 1 -- =================================================================== diff --git a/src/test/regress/sql/multi_subtransactions.sql b/src/test/regress/sql/multi_subtransactions.sql index d0f163244..2dc8fb2f9 100644 --- a/src/test/regress/sql/multi_subtransactions.sql +++ b/src/test/regress/sql/multi_subtransactions.sql @@ -152,6 +152,55 @@ COMMIT; SELECT * FROM artists WHERE id IN (11, 12) ORDER BY id; +-- Recover from multi-shard CTE modify failures +create table t1(a int, b int); +create table t2(a int, b int CHECK(b > 0)); + +ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 1190000; + +select create_distributed_table('t1', 'a'), + create_distributed_table('t2', 'a'); + +begin; +insert into t2 select i, i+1 from generate_series(1, 3) i; +with r AS ( + update t2 set b = b + 1 + returning * +) insert into t1 select * from r; +savepoint s1; +with r AS ( + update t1 set b = b - 10 + returning * +) insert into t2 select * from r; +rollback to savepoint s1; +savepoint s2; +with r AS ( + update t2 set b = b - 10 + returning * +) insert into t1 select * from r; +rollback to savepoint s2; +savepoint s3; +with r AS ( + insert into t2 select i, i+1 from generate_series(-10,-5) i + returning * +) insert into t1 select * from r; +rollback to savepoint s3; +savepoint s4; +with r AS ( + insert into t1 select i, i+1 from generate_series(-10,-5) i + returning * +) insert into t2 select * from r; +rollback to savepoint s4; +with r AS ( + update t2 set b = b + 1 + returning * +) insert into t1 select * from r; +commit; + +select * from t2 order by a, b; +select * from t1 order by a, b; + +drop table t1, t2; -- =================================================================== -- Tests for replication factor > 1