codekingpro/portable-devtools

1"""2Functions to manipulate conninfo strings3"""4 5# Copyright (C) 2020 The Psycopg Team6 7import os8import re9import socket10import asyncio11from typing import Any, Dict, List, Optional12from pathlib import Path13from datetime import tzinfo14from functools import lru_cache15from ipaddress import ip_address16 17from . import pq18from . import errors as e19from ._tz import get_tzinfo20from ._encodings import pgconn_encoding21 22 23def make_conninfo(conninfo: str = "", **kwargs: Any) -> str:24    """25    Merge a string and keyword params into a single conninfo string.26 27    :param conninfo: A `connection string`__ as accepted by PostgreSQL.28    :param kwargs: Parameters overriding the ones specified in `!conninfo`.29    :return: A connection string valid for PostgreSQL, with the `!kwargs`30        parameters merged.31 32    Raise `~psycopg.ProgrammingError` if the input doesn't make a valid33    conninfo string.34 35    .. __: https://www.postgresql.org/docs/current/libpq-connect.html36           #LIBPQ-CONNSTRING37    """38    if not conninfo and not kwargs:39        return ""40 41    # If no kwarg specified don't mung the conninfo but check if it's correct.42    # Make sure to return a string, not a subtype, to avoid making Liskov sad.43    if not kwargs:44        _parse_conninfo(conninfo)45        return str(conninfo)46 47    # Override the conninfo with the parameters48    # Drop the None arguments49    kwargs = {k: v for (k, v) in kwargs.items() if v is not None}50 51    if conninfo:52        tmp = conninfo_to_dict(conninfo)53        tmp.update(kwargs)54        kwargs = tmp55 56    conninfo = " ".join(f"{k}={_param_escape(str(v))}" for (k, v) in kwargs.items())57 58    # Verify the result is valid59    _parse_conninfo(conninfo)60 61    return conninfo62 63 64def conninfo_to_dict(conninfo: str = "", **kwargs: Any) -> Dict[str, Any]:65    """66    Convert the `!conninfo` string into a dictionary of parameters.67 68    :param conninfo: A `connection string`__ as accepted by PostgreSQL.69    :param kwargs: Parameters overriding the ones specified in `!conninfo`.70    :return: Dictionary with the parameters parsed from `!conninfo` and71        `!kwargs`.72 73    Raise `~psycopg.ProgrammingError` if `!conninfo` is not a a valid connection74    string.75 76    .. __: https://www.postgresql.org/docs/current/libpq-connect.html77           #LIBPQ-CONNSTRING78    """79    opts = _parse_conninfo(conninfo)80    rv = {opt.keyword.decode(): opt.val.decode() for opt in opts if opt.val is not None}81    for k, v in kwargs.items():82        if v is not None:83            rv[k] = v84    return rv85 86 87def _parse_conninfo(conninfo: str) -> List[pq.ConninfoOption]:88    """89    Verify that `!conninfo` is a valid connection string.90 91    Raise ProgrammingError if the string is not valid.92 93    Return the result of pq.Conninfo.parse() on success.94    """95    try:96        return pq.Conninfo.parse(conninfo.encode())97    except e.OperationalError as ex:98        raise e.ProgrammingError(str(ex))99 100 101re_escape = re.compile(r"([\\'])")102re_space = re.compile(r"\s")103 104 105def _param_escape(s: str) -> str:106    """107    Apply the escaping rule required by PQconnectdb108    """109    if not s:110        return "''"111 112    s = re_escape.sub(r"\\\1", s)113    if re_space.search(s):114        s = "'" + s + "'"115 116    return s117 118 119class ConnectionInfo:120    """Allow access to information about the connection."""121 122    __module__ = "psycopg"123 124    def __init__(self, pgconn: pq.abc.PGconn):125        self.pgconn = pgconn126 127    @property128    def vendor(self) -> str:129        """A string representing the database vendor connected to."""130        return "PostgreSQL"131 132    @property133    def host(self) -> str:134        """The server host name of the active connection. See :pq:`PQhost()`."""135        return self._get_pgconn_attr("host")136 137    @property138    def hostaddr(self) -> str:139        """The server IP address of the connection. See :pq:`PQhostaddr()`."""140        return self._get_pgconn_attr("hostaddr")141 142    @property143    def port(self) -> int:144        """The port of the active connection. See :pq:`PQport()`."""145        return int(self._get_pgconn_attr("port"))146 147    @property148    def dbname(self) -> str:149        """The database name of the connection. See :pq:`PQdb()`."""150        return self._get_pgconn_attr("db")151 152    @property153    def user(self) -> str:154        """The user name of the connection. See :pq:`PQuser()`."""155        return self._get_pgconn_attr("user")156 157    @property158    def password(self) -> str:159        """The password of the connection. See :pq:`PQpass()`."""160        return self._get_pgconn_attr("password")161 162    @property163    def options(self) -> str:164        """165        The command-line options passed in the connection request.166        See :pq:`PQoptions`.167        """168        return self._get_pgconn_attr("options")169 170    def get_parameters(self) -> Dict[str, str]:171        """Return the connection parameters values.172 173        Return all the parameters set to a non-default value, which might come174        either from the connection string and parameters passed to175        `~Connection.connect()` or from environment variables. The password176        is never returned (you can read it using the `password` attribute).177        """178        pyenc = self.encoding179 180        # Get the known defaults to avoid reporting them181        defaults = {182            i.keyword: i.compiled183            for i in pq.Conninfo.get_defaults()184            if i.compiled is not None185        }186        # Not returned by the libq. Bug? Bet we're using SSH.187        defaults.setdefault(b"channel_binding", b"prefer")188        defaults[b"passfile"] = str(Path.home() / ".pgpass").encode()189 190        return {191            i.keyword.decode(pyenc): i.val.decode(pyenc)192            for i in self.pgconn.info193            if i.val is not None194            and i.keyword != b"password"195            and i.val != defaults.get(i.keyword)196        }197 198    @property199    def dsn(self) -> str:200        """Return the connection string to connect to the database.201 202        The string contains all the parameters set to a non-default value,203        which might come either from the connection string and parameters204        passed to `~Connection.connect()` or from environment variables. The205        password is never returned (you can read it using the `password`206        attribute).207        """208        return make_conninfo(**self.get_parameters())209 210    @property211    def status(self) -> pq.ConnStatus:212        """The status of the connection. See :pq:`PQstatus()`."""213        return pq.ConnStatus(self.pgconn.status)214 215    @property216    def transaction_status(self) -> pq.TransactionStatus:217        """218        The current in-transaction status of the session.219        See :pq:`PQtransactionStatus()`.220        """221        return pq.TransactionStatus(self.pgconn.transaction_status)222 223    @property224    def pipeline_status(self) -> pq.PipelineStatus:225        """226        The current pipeline status of the client.227        See :pq:`PQpipelineStatus()`.228        """229        return pq.PipelineStatus(self.pgconn.pipeline_status)230 231    def parameter_status(self, param_name: str) -> Optional[str]:232        """233        Return a parameter setting of the connection.234 235        Return `None` is the parameter is unknown.236        """237        res = self.pgconn.parameter_status(param_name.encode(self.encoding))238        return res.decode(self.encoding) if res is not None else None239 240    @property241    def server_version(self) -> int:242        """243        An integer representing the server version. See :pq:`PQserverVersion()`.244        """245        return self.pgconn.server_version246 247    @property248    def backend_pid(self) -> int:249        """250        The process ID (PID) of the backend process handling this connection.251        See :pq:`PQbackendPID()`.252        """253        return self.pgconn.backend_pid254 255    @property256    def error_message(self) -> str:257        """258        The error message most recently generated by an operation on the connection.259        See :pq:`PQerrorMessage()`.260        """261        return self._get_pgconn_attr("error_message")262 263    @property264    def timezone(self) -> tzinfo:265        """The Python timezone info of the connection's timezone."""266        return get_tzinfo(self.pgconn)267 268    @property269    def encoding(self) -> str:270        """The Python codec name of the connection's client encoding."""271        return pgconn_encoding(self.pgconn)272 273    def _get_pgconn_attr(self, name: str) -> str:274        value: bytes = getattr(self.pgconn, name)275        return value.decode(self.encoding)276 277 278async def resolve_hostaddr_async(params: Dict[str, Any]) -> Dict[str, Any]:279    """280    Perform async DNS lookup of the hosts and return a new params dict.281 282    :param params: The input parameters, for instance as returned by283        `~psycopg.conninfo.conninfo_to_dict()`.284 285    If a ``host`` param is present but not ``hostname``, resolve the host286    addresses dynamically.287 288    The function may change the input ``host``, ``hostname``, ``port`` to allow289    connecting without further DNS lookups, eventually removing hosts that are290    not resolved, keeping the lists of hosts and ports consistent.291 292    Raise `~psycopg.OperationalError` if connection is not possible (e.g. no293    host resolve, inconsistent lists length).294    """295    hostaddr_arg = params.get("hostaddr", os.environ.get("PGHOSTADDR", ""))296    if hostaddr_arg:297        # Already resolved298        return params299 300    host_arg: str = params.get("host", os.environ.get("PGHOST", ""))301    if not host_arg:302        # Nothing to resolve303        return params304 305    hosts_in = host_arg.split(",")306    port_arg: str = str(params.get("port", os.environ.get("PGPORT", "")))307    ports_in = port_arg.split(",") if port_arg else []308    default_port = "5432"309 310    if len(ports_in) == 1:311        # If only one port is specified, the libpq will apply it to all312        # the hosts, so don't mangle it.313        default_port = ports_in.pop()314 315    elif len(ports_in) > 1:316        if len(ports_in) != len(hosts_in):317            # ProgrammingError would have been more appropriate, but this is318            # what the raise if the libpq fails connect in the same case.319            raise e.OperationalError(320                f"cannot match {len(hosts_in)} hosts with {len(ports_in)} port numbers"321            )322        ports_out = []323 324    hosts_out = []325    hostaddr_out = []326    loop = asyncio.get_running_loop()327    for i, host in enumerate(hosts_in):328        if not host or host.startswith("/") or host[1:2] == ":":329            # Local path330            hosts_out.append(host)331            hostaddr_out.append("")332            if ports_in:333                ports_out.append(ports_in[i])334            continue335 336        # If the host is already an ip address don't try to resolve it337        if is_ip_address(host):338            hosts_out.append(host)339            hostaddr_out.append(host)340            if ports_in:341                ports_out.append(ports_in[i])342            continue343 344        try:345            port = ports_in[i] if ports_in else default_port346            ans = await loop.getaddrinfo(347                host, port, proto=socket.IPPROTO_TCP, type=socket.SOCK_STREAM348            )349        except OSError as ex:350            last_exc = ex351        else:352            for item in ans:353                hosts_out.append(host)354                hostaddr_out.append(item[4][0])355                if ports_in:356                    ports_out.append(ports_in[i])357 358    # Throw an exception if no host could be resolved359    if not hosts_out:360        raise e.OperationalError(str(last_exc))361 362    out = params.copy()363    out["host"] = ",".join(hosts_out)364    out["hostaddr"] = ",".join(hostaddr_out)365    if ports_in:366        out["port"] = ",".join(ports_out)367 368    return out369 370 371@lru_cache()372def is_ip_address(s: str) -> bool:373    """Return True if the string represent a valid ip address."""374    try:375        ip_address(s)376    except ValueError:377        return False378    return True379