Merge pull request #4588 from citusdata/copy_fix_connection

When reaches to pool size, COPY sets the placement access

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
 		{

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');

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');