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):