diff --git a/src/backend/distributed/commands/multi_copy.c b/src/backend/distributed/commands/multi_copy.c index eb47a7ae2..50ca597e3 100644 --- a/src/backend/distributed/commands/multi_copy.c +++ b/src/backend/distributed/commands/multi_copy.c @@ -3583,6 +3583,12 @@ CopyGetPlacementConnection(HTAB *connectionStateHash, ShardPlacement *placement, nodeName, nodePort); + /* + * Make sure that the connection management remembers that Citus + * accesses this placement over the connection. + */ + AssignPlacementListToConnection(list_make1(placementAccess), connection); + return connection; } @@ -3643,6 +3649,12 @@ CopyGetPlacementConnection(HTAB *connectionStateHash, ShardPlacement *placement, connection = GetLeastUtilisedCopyConnection(copyConnectionStateList, nodeName, nodePort); + + /* + * Make sure that the connection management remembers that Citus + * accesses this placement over the connection. + */ + AssignPlacementListToConnection(list_make1(placementAccess), connection); } else { diff --git a/src/test/regress/expected/single_node.out b/src/test/regress/expected/single_node.out index 77b532c86..16f3c37fd 100644 --- a/src/test/regress/expected/single_node.out +++ b/src/test/regress/expected/single_node.out @@ -1415,6 +1415,18 @@ SELECT * FROM adt_ref ORDER BY 1; 5 (2 rows) +-- make sure that COPY (e.g., INSERT .. SELECT) and +-- alter_distributed_table works in the same TX +BEGIN; +SET LOCAL citus.enable_local_execution=OFF; +INSERT INTO adt_table SELECT x, x+1 FROM generate_series(1, 1000) x; +SELECT alter_distributed_table('adt_table', distribution_column:='a'); + alter_distributed_table +--------------------------------------------------------------------- + +(1 row) + +ROLLBACK; BEGIN; INSERT INTO adt_table SELECT x, x+1 FROM generate_series(1, 1000) x; SELECT alter_distributed_table('adt_table', distribution_column:='a'); diff --git a/src/test/regress/sql/single_node.sql b/src/test/regress/sql/single_node.sql index 568a956cb..a837258fe 100644 --- a/src/test/regress/sql/single_node.sql +++ b/src/test/regress/sql/single_node.sql @@ -737,6 +737,14 @@ SELECT * FROM adt_table ORDER BY 1; SELECT * FROM adt_col ORDER BY 1; SELECT * FROM adt_ref ORDER BY 1; +-- make sure that COPY (e.g., INSERT .. SELECT) and +-- alter_distributed_table works in the same TX +BEGIN; +SET LOCAL citus.enable_local_execution=OFF; +INSERT INTO adt_table SELECT x, x+1 FROM generate_series(1, 1000) x; +SELECT alter_distributed_table('adt_table', distribution_column:='a'); +ROLLBACK; + BEGIN; INSERT INTO adt_table SELECT x, x+1 FROM generate_series(1, 1000) x; SELECT alter_distributed_table('adt_table', distribution_column:='a');