Skip to content

Add connect_timeout: a per-host deadline for multi-host connect #1393

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lesnik512 wants to merge 1 commit into
base: master
Choose a base branch
from
+317 −21
Open

Add connect_timeout: a per-host deadline for multi-host connect #1393

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 54 additions & 17 deletions asyncpg/connect_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ class SSLNegotiation(compat.StrEnum):
'ssl_negotiation',
'server_settings',
'target_session_attrs',
'connect_timeout',
'krbsrvname',
'gsslib',
])
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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')

Expand All @@ -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
Expand All @@ -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',
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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.
Expand Down
28 changes: 25 additions & 3 deletions asyncpg/connection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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::

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)
Expand Down
Loading