From 7a1f719d5650f3042dfc7aec8b743daff3083a33 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Thu, 15 Oct 2026 13:01:35 +0300 Subject: [PATCH] Add load_balance_hosts to spread connections across multiple hosts With several hosts, connect() always tries them in the listed order, so every connection, including each one a pool opens, lands on the first host that accepts it. load_balance_hosts="random" shuffles the address list on every connect() call, like libpq's option of the same name. The default "disable" keeps the listed order. The value is read from the keyword argument, the DSN, the service file or PGLOADBALANCEHOSTS. A load_balance_hosts DSN option is no longer sent to the server as a setting. --- asyncpg/connect_utils.py | 41 +++++++++++- asyncpg/connection.py | 17 +++++ tests/test_connect.py | 135 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 189 insertions(+), 4 deletions(-) diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py index e215391b..4caccba8 100644 --- a/asyncpg/connect_utils.py +++ b/asyncpg/connect_utils.py @@ -59,6 +59,11 @@ class SSLNegotiation(compat.StrEnum): direct = "direct" +class LoadBalanceHosts(compat.StrEnum): + disable = "disable" + random = "random" + + _ConnectionParameters = collections.namedtuple( 'ConnectionParameters', [ @@ -70,6 +75,7 @@ class SSLNegotiation(compat.StrEnum): 'ssl_negotiation', 'server_settings', 'target_session_attrs', + 'load_balance_hosts', 'krbsrvname', 'gsslib', ]) @@ -283,7 +289,8 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user, password, passfile, database, ssl, service, servicefile, direct_tls, server_settings, - target_session_attrs, krbsrvname, gsslib): + target_session_attrs, load_balance_hosts, + krbsrvname, gsslib): # `auth_hosts` is the version of host information for the purposes # of reading the pgpass file. auth_hosts = None @@ -431,6 +438,11 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user, if target_session_attrs is None: target_session_attrs = dsn_target_session_attrs + if 'load_balance_hosts' in query: + val = query.pop('load_balance_hosts') + if load_balance_hosts is None: + load_balance_hosts = val + if 'krbsrvname' in query: val = query.pop('krbsrvname') if krbsrvname is None: @@ -548,6 +560,11 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user, if target_session_attrs is None: target_session_attrs = dsn_target_session_attrs + if 'load_balance_hosts' in service_params: + val = service_params.pop('load_balance_hosts') + if load_balance_hosts is None: + load_balance_hosts = val + if 'krbsrvname' in service_params: val = service_params.pop('krbsrvname') if krbsrvname is None: @@ -854,6 +871,20 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user, ) ) from None + if load_balance_hosts is None: + load_balance_hosts = os.getenv( + 'PGLOADBALANCEHOSTS', LoadBalanceHosts.disable + ) + try: + load_balance_hosts = LoadBalanceHosts(load_balance_hosts) + except ValueError: + raise exceptions.ClientConfigurationError( + "load_balance_hosts is expected to be one of " + "{!r}, got {!r}".format( + [m.value for m in LoadBalanceHosts], load_balance_hosts + ) + ) from None + if krbsrvname is None: krbsrvname = os.getenv('PGKRBSRVNAME') @@ -871,6 +902,7 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user, sslmode=sslmode, ssl_negotiation=sslneg, server_settings=server_settings, target_session_attrs=target_session_attrs, + load_balance_hosts=load_balance_hosts, krbsrvname=krbsrvname, gsslib=gsslib) return addrs, params @@ -882,8 +914,8 @@ def _parse_connect_arguments(*, dsn, host, port, user, password, passfile, max_cached_statement_lifetime, max_cacheable_statement_size, ssl, direct_tls, server_settings, - target_session_attrs, krbsrvname, gsslib, - service, servicefile): + target_session_attrs, load_balance_hosts, + krbsrvname, gsslib, service, servicefile): local_vars = locals() for var_name in {'max_cacheable_statement_size', 'max_cached_statement_lifetime', @@ -913,6 +945,7 @@ def _parse_connect_arguments(*, dsn, host, port, user, password, passfile, direct_tls=direct_tls, database=database, server_settings=server_settings, target_session_attrs=target_session_attrs, + load_balance_hosts=load_balance_hosts, krbsrvname=krbsrvname, gsslib=gsslib, service=service, servicefile=servicefile) @@ -1268,6 +1301,8 @@ async def _connect(*, loop, connection_class, record_class, **kwargs): addrs, params, config = _parse_connect_arguments(**kwargs) target_attr = params.target_session_attrs + if params.load_balance_hosts is LoadBalanceHosts.random: + addrs = random.sample(addrs, len(addrs)) candidates = [] chosen_connection = None diff --git a/asyncpg/connection.py b/asyncpg/connection.py index c4067fc7..4cf2b9d4 100644 --- a/asyncpg/connection.py +++ b/asyncpg/connection.py @@ -2131,6 +2131,7 @@ async def connect(dsn=None, *, record_class=protocol.Record, server_settings=None, target_session_attrs=None, + load_balance_hosts=None, krbsrvname=None, gsslib=None): r"""A coroutine to establish a connection to a PostgreSQL server. @@ -2382,6 +2383,18 @@ async def connect(dsn=None, *, or the value of the ``PGTARGETSESSIONATTRS`` environment variable, or ``"any"`` if neither is specified. + :param str load_balance_hosts: + The order in which the addresses in *host* are tried. Can be one of: + + - ``"disable"`` - in the order they are listed + - ``"random"`` - in a random order, chosen again on every + connection attempt, which spreads connections (for example, + those of a pool) across the hosts + + If not specified, the value parsed from the *dsn* argument or the + service file is used, or the value of the ``PGLOADBALANCEHOSTS`` + environment variable, or ``"disable"`` if neither is specified. + :param str krbsrvname: Kerberos service name to use when authenticating with GSSAPI. This must match the server configuration. Defaults to 'postgres'. @@ -2472,6 +2485,9 @@ async def connect(dsn=None, *, SSL contexts created by asyncpg now set the ``postgresql`` ALPN protocol, which PostgreSQL 17+ requires for direct SSL connections. + .. versionchanged:: 0.33.0 + Added the *load_balance_hosts* parameter. + .. _SSLContext: https://docs.python.org/3/library/ssl.html#ssl.SSLContext .. _create_default_context: https://docs.python.org/3/library/ssl.html#ssl.create_default_context @@ -2516,6 +2532,7 @@ async def connect(dsn=None, *, max_cached_statement_lifetime=max_cached_statement_lifetime, max_cacheable_statement_size=max_cacheable_statement_size, target_session_attrs=target_session_attrs, + load_balance_hosts=load_balance_hosts, krbsrvname=krbsrvname, gsslib=gsslib, ) diff --git a/tests/test_connect.py b/tests/test_connect.py index 1c25d8d6..7afbe02e 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -751,6 +751,69 @@ class TestConnectParams(tb.TestCase): }) }, + { + 'name': 'load_balance_hosts_default', + 'dsn': 'postgresql://user@host1,host2/db', + 'result': ([('host1', 5432), ('host2', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'load_balance_hosts': 'disable', + }) + }, + + { + 'name': 'load_balance_hosts_dsn', + 'dsn': 'postgresql://user@host1,host2/db' + '?load_balance_hosts=random', + 'env': { + 'PGLOADBALANCEHOSTS': 'disable', + }, + 'result': ([('host1', 5432), ('host2', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'load_balance_hosts': 'random', + }) + }, + + { + 'name': 'load_balance_hosts_kwarg', + 'dsn': 'postgresql://user@host1,host2/db' + '?load_balance_hosts=random', + 'load_balance_hosts': 'disable', + 'result': ([('host1', 5432), ('host2', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'load_balance_hosts': 'disable', + }) + }, + + { + 'name': 'load_balance_hosts_env', + 'dsn': 'postgresql://user@host1,host2/db', + 'env': { + 'PGLOADBALANCEHOSTS': 'random', + }, + 'result': ([('host1', 5432), ('host2', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'load_balance_hosts': 'random', + }) + }, + + { + 'name': 'load_balance_hosts_invalid', + 'dsn': 'postgresql://user@host1,host2/db' + '?load_balance_hosts=round-robin', + 'error': ( + exceptions.ClientConfigurationError, + 'load_balance_hosts is expected to be one of' + ) + }, + { 'name': 'krbsrvname', 'dsn': 'postgresql://user@host/db?krbsrvname=srv_qs', @@ -1144,7 +1207,7 @@ def run_testcase(self, testcase): test_env = {'PGHOST': None, 'PGPORT': None, 'PGUSER': None, 'PGPASSWORD': None, 'PGDATABASE': None, 'PGSSLMODE': None, - 'PGSERVICE': None, } + 'PGSERVICE': None, 'PGLOADBALANCEHOSTS': None, } test_env.update(env) dsn = testcase.get('dsn') @@ -1158,6 +1221,7 @@ def run_testcase(self, testcase): direct_tls = testcase.get('direct_tls') server_settings = testcase.get('server_settings') target_session_attrs = testcase.get('target_session_attrs') + load_balance_hosts = testcase.get('load_balance_hosts') krbsrvname = testcase.get('krbsrvname') gsslib = testcase.get('gsslib') service = testcase.get('service') @@ -1187,6 +1251,7 @@ def run_testcase(self, testcase): direct_tls=direct_tls, server_settings=server_settings, target_session_attrs=target_session_attrs, + load_balance_hosts=load_balance_hosts, krbsrvname=krbsrvname, gsslib=gsslib, service=service, servicefile=servicefile) @@ -1218,6 +1283,8 @@ def run_testcase(self, testcase): # Avoid the hassle of specifying gsslib # unless explicitly tested for params.pop('gsslib', None) + if 'load_balance_hosts' not in expected[1]: + params.pop('load_balance_hosts', None) self.assertEqual(expected, result, 'Testcase: {}'.format(testcase)) @@ -1290,6 +1357,12 @@ def test_connect_connection_service_file(self): target_session_attrs=primary krbsrvname=fakekrbsrvname gsslib=sspi + +[test_service_load_balance_hosts] +host=somehost +dbname=test_dbname +user=admin +load_balance_hosts=random ''')) connection_service_file.close() os.chmod(connection_service_file.name, stat.S_IWUSR | stat.S_IRUSR) @@ -1369,6 +1442,23 @@ def test_connect_connection_service_file(self): } ) }) + # Test load_balance_hosts from the service file + self.run_testcase({ + 'dsn': 'postgresql://?service=test_service_load_balance_hosts', + 'env': { + 'PGSERVICEFILE': connection_service_file.name, + 'PGLOADBALANCEHOSTS': 'disable', + }, + 'result': ( + [('somehost', 5432)], + { + 'user': 'admin', + 'database': 'test_dbname', + 'target_session_attrs': 'any', + 'load_balance_hosts': 'random', + } + ) + }) finally: os.unlink(connection_service_file.name) @@ -1758,6 +1848,49 @@ async def close(): await asyncio.gather(*pending, return_exceptions=True) +class TestLoadBalanceHosts(tb.TestCase): + + async def connect_order(self, **kwargs): + tried = [] + + async def connect_addr(*, addr, **_): + tried.append(addr[0]) + raise ConnectionRefusedError + + with unittest.mock.patch.object( + connect_utils, '_connect_addr', side_effect=connect_addr, + ): + with self.assertRaises(ConnectionRefusedError): + await asyncpg.connect( + host=['first', 'second', 'third'], user='postgres', + database='postgres', ssl=False, **kwargs) + return tried + + async def test_load_balance_hosts_disable_keeps_order(self): + with unittest.mock.patch.object( + connect_utils.random, 'sample', + side_effect=lambda seq, k: list(reversed(seq)), + ): + self.assertEqual(await self.connect_order(), + ['first', 'second', 'third']) + self.assertEqual( + await self.connect_order(load_balance_hosts='disable'), + ['first', 'second', 'third']) + + async def test_load_balance_hosts_random_shuffles_each_connect(self): + orders = iter([[2, 0, 1], [1, 2, 0]]) + with unittest.mock.patch.object( + connect_utils.random, 'sample', + side_effect=lambda seq, k: [seq[i] for i in next(orders)], + ): + self.assertEqual( + await self.connect_order(load_balance_hosts='random'), + ['third', 'first', 'second']) + self.assertEqual( + await self.connect_order(load_balance_hosts='random'), + ['second', 'third', 'first']) + + class TestConnection(tb.ConnectedTestCase): async def test_connection_cancelled_during_transport_setup(self):