You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
91 lines
2.8 KiB
91 lines
2.8 KiB
"""
|
|
Separate connection attempts from a connection string.
|
|
"""
|
|
|
|
# Copyright (C) 2024 The Psycopg Team
|
|
|
|
from __future__ import annotations
|
|
|
|
import socket
|
|
import logging
|
|
from random import shuffle
|
|
|
|
from . import errors as e
|
|
from ._conninfo_utils import ConnDict, get_param, is_ip_address, get_param_def
|
|
from ._conninfo_utils import split_attempts
|
|
|
|
logger = logging.getLogger("psycopg")
|
|
|
|
|
|
def conninfo_attempts(params: ConnDict) -> list[ConnDict]:
|
|
"""Split a set of connection params on the single attempts to perform.
|
|
|
|
A connection param can perform more than one attempt more than one ``host``
|
|
is provided.
|
|
|
|
Also perform async resolution of the hostname into hostaddr. Because a host
|
|
can resolve to more than one address, this can lead to yield more attempts
|
|
too. Raise `OperationalError` if no host could be resolved.
|
|
|
|
Because the libpq async function doesn't honour the timeout, we need to
|
|
reimplement the repeated attempts.
|
|
"""
|
|
last_exc = None
|
|
attempts = []
|
|
for attempt in split_attempts(params):
|
|
try:
|
|
attempts.extend(_resolve_hostnames(attempt))
|
|
except OSError as ex:
|
|
logger.debug("failed to resolve host %r: %s", attempt.get("host"), str(ex))
|
|
last_exc = ex
|
|
|
|
if not attempts:
|
|
assert last_exc
|
|
# We couldn't resolve anything
|
|
raise e.OperationalError(str(last_exc))
|
|
|
|
if get_param(params, "load_balance_hosts") == "random":
|
|
shuffle(attempts)
|
|
|
|
return attempts
|
|
|
|
|
|
def _resolve_hostnames(params: ConnDict) -> list[ConnDict]:
|
|
"""
|
|
Perform DNS lookup of the hosts and return a list of connection attempts.
|
|
|
|
If a ``host`` param is present but not ``hostname``, resolve the host
|
|
addresses asynchronously.
|
|
|
|
:param params: The input parameters, for instance as returned by
|
|
`~psycopg.conninfo.conninfo_to_dict()`. The function expects at most
|
|
a single entry for host, hostaddr because it is designed to further
|
|
process the input of split_attempts().
|
|
|
|
:return: A list of attempts to make (to include the case of a hostname
|
|
resolving to more than one IP).
|
|
"""
|
|
host = get_param(params, "host")
|
|
if not host or host.startswith("/") or host[1:2] == ":":
|
|
# Local path, or no host to resolve
|
|
return [params]
|
|
|
|
hostaddr = get_param(params, "hostaddr")
|
|
if hostaddr:
|
|
# Already resolved
|
|
return [params]
|
|
|
|
if is_ip_address(host):
|
|
# If the host is already an ip address don't try to resolve it
|
|
return [{**params, "hostaddr": host}]
|
|
|
|
port = get_param(params, "port")
|
|
if not port:
|
|
port_def = get_param_def("port")
|
|
port = port_def and port_def.compiled or "5432"
|
|
|
|
ans = socket.getaddrinfo(
|
|
host, port, proto=socket.IPPROTO_TCP, type=socket.SOCK_STREAM
|
|
)
|
|
return [{**params, "hostaddr": item[4][0]} for item in ans]
|