Merge pull request #5403 from citusdata/reuse_to_be_deleted_connection_in_same_transaction

Don't skip connections with forceCloseAtTransactionEnd that Sent Begin in FindAvailableConnection

src/backend/distributed/connection/connection_management.c

@@ -429,12 +429,18 @@ FindAvailableConnection(dlist_head *connections, uint32 flags)
 			continue;
 		}
 
-		if (connection->forceCloseAtTransactionEnd)
+		if (connection->forceCloseAtTransactionEnd &&
+			!connection->remoteTransaction.beginSent)
 		{
 			/*
-			 * This is a connection that should be closed, probabably because
+			 * This is a connection that should be closed, probably because
-			 * of old connection options. So we ignore it. It will
+			 * of old connection options or removing a node. This will
-			 * automatically be closed at the end of the transaction.
+			 * automatically be closed at the end of the transaction. But, if we are still
+			 * inside a transaction, we should keep using this connection as long as a remote
+			 * transaction is in progress over the connection. The main use for this case
+			 * is having some commands inside a transaction block after removing nodes. And, we
+			 * currently allow very limited operations after removing a node inside a
+			 * transaction block (e.g., no placement access can happen).
 			 */
 			continue;
 		}

src/test/regress/expected/node_conninfo_reload.out

@@ -458,5 +458,56 @@ select count(*) from test where a = 0;
      0
 (1 row)
 
+-- Test connecting all the shards
+ALTER SYSTEM SET citus.node_conninfo = 'sslmode=doesnotexist';
+BEGIN;
+ALTER TABLE test ADD COLUMN b INT;
+select pg_reload_conf();
+ pg_reload_conf
+---------------------------------------------------------------------
+ t
+(1 row)
+
+select pg_sleep(0.1); -- wait for config reload to apply
+ pg_sleep
+---------------------------------------------------------------------
+
+(1 row)
+
+show citus.node_conninfo;
+ citus.node_conninfo
+---------------------------------------------------------------------
+ sslmode=doesnotexist
+(1 row)
+
+-- Should work since connections to the same shards that BEGIN is sent
+-- are reused.
+ALTER TABLE test ADD COLUMN c INT;
+COMMIT;
+-- Should fail now, when transaction is finished
+ALTER TABLE test ADD COLUMN d INT;
+ERROR:  connection to the remote node localhost:xxxxx failed with the following error: invalid sslmode value: "doesnotexist"
+-- Reset it again
+ALTER SYSTEM RESET citus.node_conninfo;
+select pg_reload_conf();
+ pg_reload_conf
+---------------------------------------------------------------------
+ t
+(1 row)
+
+select pg_sleep(0.1); -- wait for config reload to apply
+ pg_sleep
+---------------------------------------------------------------------
+
+(1 row)
+
+show citus.node_conninfo;
+ citus.node_conninfo
+---------------------------------------------------------------------
+ sslmode=require
+(1 row)
+
+-- Should work again
+ALTER TABLE test ADD COLUMN e INT;
 DROP SCHEMA node_conninfo_reload CASCADE;
 NOTICE:  drop cascades to table test

src/test/regress/sql/node_conninfo_reload.sql

@@ -171,4 +171,28 @@ show citus.node_conninfo;
 -- Should work
 select count(*) from test where a = 0;
 
+-- Test connecting all the shards
+ALTER SYSTEM SET citus.node_conninfo = 'sslmode=doesnotexist';
+BEGIN;
+ALTER TABLE test ADD COLUMN b INT;
+select pg_reload_conf();
+select pg_sleep(0.1); -- wait for config reload to apply
+show citus.node_conninfo;
+-- Should work since connections to the same shards that BEGIN is sent
+-- are reused.
+ALTER TABLE test ADD COLUMN c INT;
+COMMIT;
+
+-- Should fail now, when transaction is finished
+ALTER TABLE test ADD COLUMN d INT;
+
+-- Reset it again
+ALTER SYSTEM RESET citus.node_conninfo;
+select pg_reload_conf();
+select pg_sleep(0.1); -- wait for config reload to apply
+show citus.node_conninfo;
+
+-- Should work again
+ALTER TABLE test ADD COLUMN e INT;
+
 DROP SCHEMA node_conninfo_reload CASCADE;