From 8d7cb38db09dd58b81c5a1edfdb02cea2f4495bb Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Thu, 15 Oct 2026 12:29:41 +0300 Subject: [PATCH] Add connect_timeout: a per-host deadline for multi-host connect With several hosts, a host that hangs instead of refusing used up the whole connect timeout, so later hosts were never tried. connect_timeout limits each address attempt (TCP, TLS, startup/auth and the target_session_attrs check) and moves on to the next address when it expires, like libpq's connect_timeout. timeout still caps the whole call. The value is read from the keyword argument, the DSN, the service file or PGCONNECT_TIMEOUT. A connect_timeout DSN option is no longer sent to the server as a setting. Fixes #1155. Fixes #516. --- asyncpg/connect_utils.py | 71 +++++++++--- asyncpg/connection.py | 28 ++++- tests/test_connect.py | 239 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 317 insertions(+), 21 deletions(-) diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py index e215391b..f7b1b52b 100644 --- a/asyncpg/connect_utils.py +++ b/asyncpg/connect_utils.py @@ -70,6 +70,7 @@ class SSLNegotiation(compat.StrEnum): 'ssl_negotiation', 'server_settings', 'target_session_attrs', + 'connect_timeout', 'krbsrvname', 'gsslib', ]) @@ -283,7 +284,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, connect_timeout, + krbsrvname, gsslib): # `auth_hosts` is the version of host information for the purposes # of reading the pgpass file. auth_hosts = None @@ -431,6 +433,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 'connect_timeout' in query: + val = query.pop('connect_timeout') + if connect_timeout is None: + connect_timeout = val + if 'krbsrvname' in query: val = query.pop('krbsrvname') if krbsrvname is None: @@ -548,6 +555,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 'connect_timeout' in service_params: + val = service_params.pop('connect_timeout') + if connect_timeout is None: + connect_timeout = val + if 'krbsrvname' in service_params: val = service_params.pop('krbsrvname') if krbsrvname is None: @@ -854,6 +866,21 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user, ) ) from None + if connect_timeout is None: + connect_timeout = os.getenv('PGCONNECT_TIMEOUT') + if connect_timeout is not None: + try: + if isinstance(connect_timeout, bool): + raise ValueError + connect_timeout = float(connect_timeout) + except (TypeError, ValueError): + raise exceptions.ClientConfigurationError( + 'invalid connect_timeout value: ' + 'expected a number of seconds, got {!r}'.format( + connect_timeout)) from None + if not connect_timeout > 0: + connect_timeout = None + if krbsrvname is None: krbsrvname = os.getenv('PGKRBSRVNAME') @@ -871,6 +898,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, + connect_timeout=connect_timeout, krbsrvname=krbsrvname, gsslib=gsslib) return addrs, params @@ -882,8 +910,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, connect_timeout, + krbsrvname, gsslib, service, servicefile): local_vars = locals() for var_name in {'max_cacheable_statement_size', 'max_cached_statement_lifetime', @@ -913,6 +941,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, + connect_timeout=connect_timeout, krbsrvname=krbsrvname, gsslib=gsslib, service=service, servicefile=servicefile) @@ -1270,30 +1299,38 @@ async def _connect(*, loop, connection_class, record_class, **kwargs): target_attr = params.target_session_attrs candidates = [] + unmatched = [] chosen_connection = None last_error = None try: for addr in addrs: try: - conn = await _connect_addr( - addr=addr, - loop=loop, - params=params, - config=config, - connection_class=connection_class, - record_class=record_class, - ) - candidates.append(conn) - if await _can_use_connection(conn, target_attr): - chosen_connection = conn - break + async with compat.timeout(params.connect_timeout): + conn = await _connect_addr( + addr=addr, + loop=loop, + params=params, + config=config, + connection_class=connection_class, + record_class=record_class, + ) + candidates.append(conn) + if await _can_use_connection(conn, target_attr): + chosen_connection = conn + break + unmatched.append(conn) except OSError as ex: last_error = ex + except asyncio.TimeoutError as ex: + # Python < 3.11, where TimeoutError is not an OSError. + if params.connect_timeout is None: + raise + last_error = ex except _NextHostSignal as ex: last_error = ex.__cause__ else: - if target_attr == SessionAttribute.prefer_standby and candidates: - chosen_connection = random.choice(candidates) + if target_attr == SessionAttribute.prefer_standby and unmatched: + chosen_connection = random.choice(unmatched) finally: # Do not wait for rejected hosts to disconnect: their cleanup must # neither delay a successful connection nor outlive the event loop. diff --git a/asyncpg/connection.py b/asyncpg/connection.py index c4067fc7..738ced9a 100644 --- a/asyncpg/connection.py +++ b/asyncpg/connection.py @@ -2121,6 +2121,7 @@ async def connect(dsn=None, *, database=None, loop=None, timeout=60, + connect_timeout=None, statement_cache_size=100, max_cached_statement_lifetime=300, max_cacheable_statement_size=1024 * 15, @@ -2151,8 +2152,9 @@ async def connect(dsn=None, *, The following options are recognized by asyncpg: ``host``, ``port``, ``user``, ``database`` (or ``dbname``), ``password``, ``passfile``, ``sslmode``, ``sslcert``, ``sslkey``, ``sslrootcert``, - and ``sslcrl``. Unlike libpq, asyncpg will treat unrecognized - options as `server settings`_ to be used for the connection. + ``sslcrl``, and ``connect_timeout``. Unlike libpq, asyncpg will + treat unrecognized options as `server settings`_ to be used for the + connection. .. note:: @@ -2240,7 +2242,23 @@ async def connect(dsn=None, *, event loop will be used. :param float timeout: - Connection timeout in seconds. + Connection timeout in seconds. This is the limit for the whole + call, including attempts on every host in *host*. + + :param float connect_timeout: + Maximum time in seconds to wait for each host address to accept + the connection, finish TLS and authentication, and pass the + *target_session_attrs* check. When it expires, the attempt is + abandoned and the next address is tried, as when the host refuses + the connection. If every address times out, + :exc:`TimeoutError` is raised. *timeout* still caps the whole + call, so the worst case is the smaller of *timeout* and the number + of addresses times *connect_timeout*. Zero, a negative value or + ``None`` disables the per-host limit. + + If not specified, the ``connect_timeout`` option from *dsn* or the + service file is used, or the value of the ``PGCONNECT_TIMEOUT`` + environment variable. :param int statement_cache_size: The size of prepared statement LRU cache. Pass ``0`` to @@ -2472,6 +2490,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 *connect_timeout* 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 +2537,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, + connect_timeout=connect_timeout, krbsrvname=krbsrvname, gsslib=gsslib, ) diff --git a/tests/test_connect.py b/tests/test_connect.py index 1c25d8d6..2da3de35 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -751,6 +751,91 @@ class TestConnectParams(tb.TestCase): }) }, + { + 'name': 'connect_timeout_dsn', + 'dsn': 'postgresql://user@host1,host2/db?connect_timeout=5', + 'env': { + 'PGCONNECT_TIMEOUT': '7', + }, + 'result': ([('host1', 5432), ('host2', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'connect_timeout': 5.0, + }) + }, + + { + 'name': 'connect_timeout_kwarg', + 'dsn': 'postgresql://user@host/db?connect_timeout=5', + 'connect_timeout': 2.5, + 'result': ([('host', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'connect_timeout': 2.5, + }) + }, + + { + 'name': 'connect_timeout_env', + 'dsn': 'postgresql://user@host/db', + 'env': { + 'PGCONNECT_TIMEOUT': '7', + }, + 'result': ([('host', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'connect_timeout': 7.0, + }) + }, + + { + 'name': 'connect_timeout_zero', + 'dsn': 'postgresql://user@host/db?connect_timeout=0', + 'env': { + 'PGCONNECT_TIMEOUT': '7', + }, + 'result': ([('host', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'connect_timeout': None, + }) + }, + + { + 'name': 'connect_timeout_negative_kwarg', + 'dsn': 'postgresql://user@host/db?connect_timeout=5', + 'connect_timeout': -1, + 'result': ([('host', 5432)], { + 'database': 'db', + 'user': 'user', + 'target_session_attrs': 'any', + 'connect_timeout': None, + }) + }, + + { + 'name': 'connect_timeout_invalid', + 'dsn': 'postgresql://user@host/db?connect_timeout=soon', + 'error': ( + exceptions.ClientConfigurationError, + 'invalid connect_timeout value' + ) + }, + + { + 'name': 'connect_timeout_invalid_kwarg', + 'dsn': 'postgresql://user@host/db', + 'connect_timeout': True, + 'error': ( + exceptions.ClientConfigurationError, + 'invalid connect_timeout value' + ) + }, + { 'name': 'krbsrvname', 'dsn': 'postgresql://user@host/db?krbsrvname=srv_qs', @@ -1144,7 +1229,7 @@ def run_testcase(self, testcase): test_env = {'PGHOST': None, 'PGPORT': None, 'PGUSER': None, 'PGPASSWORD': None, 'PGDATABASE': None, 'PGSSLMODE': None, - 'PGSERVICE': None, } + 'PGSERVICE': None, 'PGCONNECT_TIMEOUT': None, } test_env.update(env) dsn = testcase.get('dsn') @@ -1158,6 +1243,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') + connect_timeout = testcase.get('connect_timeout') krbsrvname = testcase.get('krbsrvname') gsslib = testcase.get('gsslib') service = testcase.get('service') @@ -1187,6 +1273,7 @@ def run_testcase(self, testcase): direct_tls=direct_tls, server_settings=server_settings, target_session_attrs=target_session_attrs, + connect_timeout=connect_timeout, krbsrvname=krbsrvname, gsslib=gsslib, service=service, servicefile=servicefile) @@ -1290,6 +1377,12 @@ def test_connect_connection_service_file(self): target_session_attrs=primary krbsrvname=fakekrbsrvname gsslib=sspi + +[test_service_connect_timeout] +host=somehost +dbname=test_dbname +user=admin +connect_timeout=3 ''')) connection_service_file.close() os.chmod(connection_service_file.name, stat.S_IWUSR | stat.S_IRUSR) @@ -1369,6 +1462,40 @@ def test_connect_connection_service_file(self): } ) }) + # Test connect_timeout from the service file + self.run_testcase({ + 'dsn': 'postgresql://?service=test_service_connect_timeout', + 'env': { + 'PGSERVICEFILE': connection_service_file.name, + 'PGCONNECT_TIMEOUT': '7', + }, + 'result': ( + [('somehost', 5432)], + { + 'user': 'admin', + 'database': 'test_dbname', + 'target_session_attrs': 'any', + 'connect_timeout': 3.0, + } + ) + }) + # Test that dsn connect_timeout overrides the service file + self.run_testcase({ + 'dsn': 'postgresql://?service=test_service_connect_timeout' + '&connect_timeout=4', + 'env': { + 'PGSERVICEFILE': connection_service_file.name + }, + 'result': ( + [('somehost', 5432)], + { + 'user': 'admin', + 'database': 'test_dbname', + 'target_session_attrs': 'any', + 'connect_timeout': 4.0, + } + ) + }) finally: os.unlink(connection_service_file.name) @@ -1758,6 +1885,116 @@ async def close(): await asyncio.gather(*pending, return_exceptions=True) +class TestConnectTimeout(tb.ClusterTestCase): + + @contextlib.asynccontextmanager + async def silent_server(self): + stop = asyncio.Event() + + async def handle(reader, writer): + await stop.wait() + writer.close() + + server = await asyncio.start_server(handle, '127.0.0.1', 0) + try: + yield server.sockets[0].getsockname()[1] + finally: + stop.set() + server.close() + await server.wait_closed() + + def connect_two_hosts(self, silent_port, **kwargs): + spec = self.get_connection_spec() + return pg_connection.connect( + host=['127.0.0.1', spec['host']], + port=[silent_port, spec['port']], + user=spec['user'], database=spec['database'], + loop=self.loop, **kwargs) + + async def test_connect_timeout_tries_next_host(self): + async with self.silent_server() as silent_port: + with self.assertRunUnder(5): + con = await self.connect_two_hosts( + silent_port, connect_timeout=0.5, timeout=30) + try: + self.assertEqual(await con.fetchval('SELECT 1'), 1) + self.assertFalse(_get_connected_host(con).endswith( + str(silent_port))) + finally: + await con.close() + + async def test_connect_timeout_unset_uses_whole_timeout(self): + async with self.silent_server() as silent_port: + with self.assertRaises(asyncio.TimeoutError): + await self.connect_two_hosts(silent_port, timeout=1) + + async def test_connect_timeout_all_hosts_time_out(self): + async with self.silent_server() as port1, \ + self.silent_server() as port2: + with self.assertRunUnder(5), \ + self.assertRaises(asyncio.TimeoutError): + await pg_connection.connect( + host=['127.0.0.1', '127.0.0.1'], port=[port1, port2], + user='postgres', connect_timeout=0.5, timeout=30, + loop=self.loop) + + async def test_connect_timeout_in_dsn_is_not_a_server_setting(self): + con = await self.connect( + dsn='postgresql:///postgres?connect_timeout=5', + host=self.get_connection_spec()['host'], user='postgres') + await con.close() + + async def test_connect_timeout_covers_target_session_attrs_check(self): + hung = unittest.mock.Mock(spec=pg_connection.Connection) + selected = unittest.mock.Mock(spec=pg_connection.Connection) + + async def can_use(con, attr): + if con is hung: + await asyncio.Future() + return True + + with unittest.mock.patch.object( + connect_utils, '_connect_addr', side_effect=[hung, selected], + ), unittest.mock.patch.object( + connect_utils, '_can_use_connection', side_effect=can_use, + ): + with self.assertRunUnder(5): + con = await asyncpg.connect( + host=['first', 'second'], user='postgres', + database='postgres', ssl=False, + target_session_attrs='primary', + connect_timeout=0.5, timeout=30, + ) + self.assertIs(con, selected) + hung.terminate.assert_called_once_with() + selected.terminate.assert_not_called() + + async def test_connect_timeout_prefer_standby_skips_timed_out_host(self): + hung = unittest.mock.Mock(spec=pg_connection.Connection) + primary = unittest.mock.Mock(spec=pg_connection.Connection) + + async def can_use(con, attr): + if con is hung: + await asyncio.Future() + return False + + with unittest.mock.patch.object( + connect_utils, '_connect_addr', side_effect=[hung, primary], + ), unittest.mock.patch.object( + connect_utils, '_can_use_connection', side_effect=can_use, + ), unittest.mock.patch.object( + connect_utils.random, 'choice', side_effect=lambda seq: seq[0], + ): + con = await asyncpg.connect( + host=['first', 'second'], user='postgres', + database='postgres', ssl=False, + target_session_attrs='prefer-standby', + connect_timeout=0.5, timeout=30, + ) + self.assertIs(con, primary) + hung.terminate.assert_called_once_with() + + class TestConnection(tb.ConnectedTestCase): async def test_connection_cancelled_during_transport_setup(self):