Hi,

Attached is the rebased patch. The use case I am targeting is exactly
what Artem describes above. Happy to make any necessary changes.

Thanks,
Andrew Jackson
From b0559224bdb4fc8d72323b25ee065913c0110cfc Mon Sep 17 00:00:00 2001
From: CommanderKeynes <andrewjackson...@gmail.coma>
Date: Sat, 17 May 2025 08:29:01 -0500
Subject: [PATCH] Add option to check all addrs for target_session.

The current behaviour of libpq with regard to searching
for a matching target_session_attrs in a list of addrs is
that after successfully connecting to a server if the servers
session_attr does not match the request target_session_attrs
no futher address is considered. This behaviour is extremely
inconvenient in environments where the user is attempting to
implement a high availability setup without having to modify
DNS records or a proxy server config.

This PR adds a client side option called check_all_addrs.
When set to 1 this option will tell libpq to continue checking
any remaining addresses even if there was a target_session_attrs
mismatch on one of them.

Author: Andrew Jackson
Reviewed-by: Andrey Borodin
Discussion: https://www.postgresql.org/message-id/flat/CAKK5BkESSc69sp2TiTWHvvOHCUey0rDWXSrR9pinyRqyfamUYg%40mail.gmail.com
---
 doc/src/sgml/libpq.sgml                       |  33 +++++
 src/interfaces/libpq/fe-connect.c             |  25 ++--
 src/interfaces/libpq/libpq-int.h              |   1 +
 .../libpq/t/007_target_session_attr_dns.pl    | 129 ++++++++++++++++++
 .../t/008_load_balance_dns_check_all_addrs.pl | 128 +++++++++++++++++
 5 files changed, 306 insertions(+), 10 deletions(-)
 create mode 100644 src/interfaces/libpq/t/007_target_session_attr_dns.pl
 create mode 100644 src/interfaces/libpq/t/008_load_balance_dns_check_all_addrs.pl

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 5bf59a1985..ff3ddddd63 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -2568,6 +2568,39 @@ postgresql://%2Fvar%2Flib%2Fpostgresql/dbname
       </listitem>
      </varlistentry>
 
+     <varlistentry id="libpq-connect-check-all-addrs" xreflabel="check_all_addrs">
+      <term><literal>check_all_addrs</literal></term>
+      <listitem>
+       <para>
+        Controls whether or not all addresses within a hostname are checked when trying to make a connection
+        when attempting to find a connection with a matching <xref linkend="libpq-connect-target-session-attrs"/>.
+
+        There are two modes:
+        <variablelist>
+         <varlistentry>
+          <term><literal>0</literal> (default)</term>
+          <listitem>
+           <para>
+            If a successful connection is made and that connection is found to have a
+            mismatching <xref linkend="libpq-connect-target-session-attrs"/> do not check
+            any additional addresses and move onto the next host if one was provided.
+           </para>
+          </listitem>
+         </varlistentry>
+         <varlistentry>
+          <term><literal>1</literal></term>
+          <listitem>
+           <para>
+            If a successful connection is made and that connection is found to have a
+            mismatching <xref linkend="libpq-connect-target-session-attrs"/> proceed
+            to check any additional addresses.
+           </para>
+          </listitem>
+         </varlistentry>
+        </variablelist>
+       </para>
+      </listitem>
+     </varlistentry>
     </variablelist>
    </para>
   </sect2>
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index a3d12931ff..a007761116 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -393,6 +393,10 @@ static const internalPQconninfoOption PQconninfoOptions[] = {
 
 	{"scram_server_key", NULL, NULL, NULL, "SCRAM-Server-Key", "D", SCRAM_MAX_KEY_LEN * 2,
 	offsetof(struct pg_conn, scram_server_key)},
+	{"check_all_addrs", "PGCHECKALLADDRS",
+		DefaultLoadBalanceHosts, NULL,
+		"Check-All-Addrs", "", 1,
+	offsetof(struct pg_conn, check_all_addrs)},
 
 	/* OAuth v2 */
 	{"oauth_issuer", NULL, NULL, NULL,
@@ -4434,11 +4438,11 @@ keep_going:						/* We will come back to here until there is
 						conn->status = CONNECTION_OK;
 						sendTerminateConn(conn);
 
-						/*
-						 * Try next host if any, but we don't want to consider
-						 * additional addresses for this host.
-						 */
-						conn->try_next_host = true;
+						if (strcmp(conn->check_all_addrs, "1") == 0)
+							conn->try_next_addr = true;
+						else
+							conn->try_next_host = true;
+
 						goto keep_going;
 					}
 				}
@@ -4489,11 +4493,11 @@ keep_going:						/* We will come back to here until there is
 						conn->status = CONNECTION_OK;
 						sendTerminateConn(conn);
 
-						/*
-						 * Try next host if any, but we don't want to consider
-						 * additional addresses for this host.
-						 */
-						conn->try_next_host = true;
+						if (strcmp(conn->check_all_addrs, "1") == 0)
+							conn->try_next_addr = true;
+						else
+							conn->try_next_host = true;
+
 						goto keep_going;
 					}
 				}
@@ -5127,6 +5131,7 @@ freePGconn(PGconn *conn)
 	free(conn->inBuffer);
 	free(conn->outBuffer);
 	free(conn->rowBuf);
+	free(conn->check_all_addrs);
 	termPQExpBuffer(&conn->errorMessage);
 	termPQExpBuffer(&conn->workBuffer);
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index a701c25038..999eb3cf71 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -430,6 +430,7 @@ struct pg_conn
 	char	   *scram_client_key;	/* base64-encoded SCRAM client key */
 	char	   *scram_server_key;	/* base64-encoded SCRAM server key */
 	char	   *sslkeylogfile;	/* where should the client write ssl keylogs */
+	char       *check_all_addrs;  /* whether to check all ips within a host or terminate on failure */
 
 	bool		cancelRequest;	/* true if this connection is used to send a
 								 * cancel request, instead of being a normal
diff --git a/src/interfaces/libpq/t/007_target_session_attr_dns.pl b/src/interfaces/libpq/t/007_target_session_attr_dns.pl
new file mode 100644
index 0000000000..914f6c472f
--- /dev/null
+++ b/src/interfaces/libpq/t/007_target_session_attr_dns.pl
@@ -0,0 +1,129 @@
+
+# Copyright (c) 2023-2025, PostgreSQL Global Development Group
+use strict;
+use warnings FATAL => 'all';
+use Config;
+use PostgreSQL::Test::Utils;
+use PostgreSQL::Test::Cluster;
+use Test::More;
+
+if (!$ENV{PG_TEST_EXTRA} || $ENV{PG_TEST_EXTRA} !~ /\bload_balance\b/)
+{
+	plan skip_all =>
+	  'Potentially unsafe test load_balance not enabled in PG_TEST_EXTRA';
+}
+
+# Cluster setup which is shared for testing both load balancing methods
+my $can_bind_to_127_0_0_2 =
+  $Config{osname} eq 'linux' || $PostgreSQL::Test::Utils::windows_os;
+
+# Checks for the requirements for testing load balancing method 2
+if (!$can_bind_to_127_0_0_2)
+{
+	plan skip_all => 'load_balance test only supported on Linux and Windows';
+}
+
+my $hosts_path;
+if ($windows_os)
+{
+	$hosts_path = 'c:\Windows\System32\Drivers\etc\hosts';
+}
+else
+{
+	$hosts_path = '/etc/hosts';
+}
+
+my $hosts_content = PostgreSQL::Test::Utils::slurp_file($hosts_path);
+
+my $hosts_count = () =
+  $hosts_content =~ /127\.0\.0\.[1-3] pg-loadbalancetest/g;
+if ($hosts_count != 3)
+{
+	# Host file is not prepared for this test
+	plan skip_all => "hosts file was not prepared for DNS load balance test";
+}
+
+$PostgreSQL::Test::Cluster::use_tcp = 1;
+$PostgreSQL::Test::Cluster::test_pghost = '127.0.0.1';
+my $port = PostgreSQL::Test::Cluster::get_free_port();
+
+my $node_primary1 = PostgreSQL::Test::Cluster->new('primary1', port => $port);
+$node_primary1->init(has_archiving => 1, allows_streaming => 1);
+
+# Start it
+$node_primary1->start;
+
+# Take backup from which all operations will be run
+$node_primary1->backup('my_backup');
+
+my $node_standby = PostgreSQL::Test::Cluster->new('standby', port => $port, own_host => 1);
+$node_standby->init_from_backup($node_primary1, 'my_backup',
+	has_restoring => 1);
+$node_standby->start();
+
+my $node_primary2 = PostgreSQL::Test::Cluster->new('node1', port => $port, own_host => 1);
+$node_primary2 ->init();
+$node_primary2 ->start();
+
+# target_session_attrs=primary should always choose the first one.
+$node_primary1->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=primary check_all_addrs=1",
+	"target_session_attrs=primary connects to the first node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+$node_primary1->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=read-write check_all_addrs=1",
+	"target_session_attrs=read-write connects to the first node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+$node_primary1->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=any check_all_addrs=1",
+	"target_session_attrs=any connects to the first node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+$node_standby->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=standby check_all_addrs=1",
+	"target_session_attrs=standby connects to the third node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+$node_standby->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=read-only check_all_addrs=1",
+	"target_session_attrs=read-only connects to the third node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+
+
+$node_primary1->stop();
+
+# target_session_attrs=primary should always choose the first one.
+$node_primary2->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=primary check_all_addrs=1",
+	"target_session_attrs=primary connects to the first node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+$node_primary2->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=read-write check_all_addrs=1",
+	"target_session_attrs=read-write connects to the first node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+$node_standby->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=any check_all_addrs=1",
+	"target_session_attrs=any connects to the first node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+$node_standby->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=standby check_all_addrs=1",
+	"target_session_attrs=standby connects to the third node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+$node_standby->connect_ok(
+	"host=pg-loadbalancetest port=$port target_session_attrs=read-only check_all_addrs=1",
+	"target_session_attrs=read-only connects to the third node",
+	sql => "SELECT 'connect1'",
+	log_like => [qr/statement: SELECT 'connect1'/]);
+
+$node_primary2->stop();
+$node_standby->stop();
+
+
+done_testing();
diff --git a/src/interfaces/libpq/t/008_load_balance_dns_check_all_addrs.pl b/src/interfaces/libpq/t/008_load_balance_dns_check_all_addrs.pl
new file mode 100644
index 0000000000..d3405598e6
--- /dev/null
+++ b/src/interfaces/libpq/t/008_load_balance_dns_check_all_addrs.pl
@@ -0,0 +1,128 @@
+# Copyright (c) 2023-2025, PostgreSQL Global Development Group
+use strict;
+use warnings FATAL => 'all';
+use Config;
+use PostgreSQL::Test::Utils;
+use PostgreSQL::Test::Cluster;
+use Test::More;
+
+if (!$ENV{PG_TEST_EXTRA} || $ENV{PG_TEST_EXTRA} !~ /\bload_balance\b/)
+{
+	plan skip_all =>
+	  'Potentially unsafe test load_balance not enabled in PG_TEST_EXTRA';
+}
+
+my $can_bind_to_127_0_0_2 =
+  $Config{osname} eq 'linux' || $PostgreSQL::Test::Utils::windows_os;
+
+# Checks for the requirements for testing load balancing method 2
+if (!$can_bind_to_127_0_0_2)
+{
+	plan skip_all => 'load_balance test only supported on Linux and Windows';
+}
+
+my $hosts_path;
+if ($windows_os)
+{
+	$hosts_path = 'c:\Windows\System32\Drivers\etc\hosts';
+}
+else
+{
+	$hosts_path = '/etc/hosts';
+}
+
+my $hosts_content = PostgreSQL::Test::Utils::slurp_file($hosts_path);
+
+my $hosts_count = () =
+  $hosts_content =~ /127\.0\.0\.[1-3] pg-loadbalancetest/g;
+if ($hosts_count != 3)
+{
+	# Host file is not prepared for this test
+	plan skip_all => "hosts file was not prepared for DNS load balance test";
+}
+
+$PostgreSQL::Test::Cluster::use_tcp = 1;
+$PostgreSQL::Test::Cluster::test_pghost = '127.0.0.1';
+
+my $port = PostgreSQL::Test::Cluster::get_free_port();
+local $Test::Builder::Level = $Test::Builder::Level + 1;
+my $node_primary1 = PostgreSQL::Test::Cluster->new("primary1", port => $port);
+$node_primary1->init(has_archiving => 1, allows_streaming => 1);
+
+# Start it
+$node_primary1->start();
+
+# Take backup from which all operations will be run
+$node_primary1->backup("my_backup");
+
+my $node_standby = PostgreSQL::Test::Cluster->new("standby", port => $port, own_host => 1);
+$node_standby->init_from_backup($node_primary1, "my_backup",
+	has_restoring => 1);
+$node_standby->start();
+
+my $node_primary2 = PostgreSQL::Test::Cluster->new("node1", port => $port, own_host => 1);
+$node_primary2->init();
+$node_primary2->start();
+sub test_target_session_attr {
+	my $target_session_attrs = shift;
+	my $test_num = shift;
+	my $primary1_expect_traffic = shift;
+	my $standby_expeect_traffic = shift;
+	my $primary2_expect_traffic = shift;
+	# Statistically the following loop with load_balance_hosts=random will almost
+	# certainly connect at least once to each of the nodes. The chance of that not
+	# happening is so small that it's negligible: (2/3)^50 = 1.56832855e-9
+	foreach my $i (1 .. 50)
+	{
+		$node_primary1->connect_ok(
+			"host=pg-loadbalancetest port=$port load_balance_hosts=random target_session_attrs=${target_session_attrs} check_all_addrs=1",
+			"repeated connections with random load balancing",
+			sql => "SELECT 'connect${test_num}'");
+	}
+	my $node_primary1_occurrences = () =
+	  $node_primary1->log_content() =~ /statement: SELECT 'connect${test_num}'/g;
+	my $node_standby_occurrences = () =
+	  $node_standby->log_content() =~ /statement: SELECT 'connect${test_num}'/g;
+	my $node_primary2_occurrences = () =
+	  $node_primary2->log_content() =~ /statement: SELECT 'connect${test_num}'/g;
+
+	my $total_occurrences =
+	  $node_primary1_occurrences + $node_standby_occurrences + $node_primary2_occurrences;
+
+	if ($primary1_expect_traffic == 1) {
+		ok($node_primary1_occurrences > 0, "received at least one connection on node1");
+	}else{
+		ok($node_primary1_occurrences == 0, "received at least one connection on node1");
+	}
+	if ($standby_expeect_traffic == 1) {
+		ok($node_standby_occurrences > 0, "received at least one connection on node1");
+	}else{
+		ok($node_standby_occurrences == 0, "received at least one connection on node1");
+	}
+
+	if ($primary2_expect_traffic == 1) {
+		ok($node_primary2_occurrences > 0, "received at least one connection on node1");
+	}else{
+		ok($node_primary2_occurrences == 0, "received at least one connection on node1");
+	}
+
+	ok($total_occurrences == 50, "received 50 connections across all nodes");
+}
+
+test_target_session_attr('any',
+	1, 1, 1, 1);
+test_target_session_attr('read-only',
+	2, 0, 1, 0);
+test_target_session_attr('read-write',
+	3, 1, 0, 1);
+test_target_session_attr('primary',
+	4, 1, 0, 1);
+test_target_session_attr('standby',
+	5, 0, 1, 0);
+
+
+$node_primary1->stop();
+$node_primary2->stop();
+$node_standby->stop();
+
+done_testing();
-- 
2.49.0

Reply via email to