diff --git a/patroni/psycopg.py b/patroni/psycopg.py index 86e58b31..6b9b5629 100644 --- a/patroni/psycopg.py +++ b/patroni/psycopg.py @@ -1,3 +1,11 @@ +"""Abstraction layer for ``psycopg`` module. + +This module is able to handle both ``pyscopg2`` and ``psycopg3``, and it exposes a common interface for both. +``psycopg2`` takes precedence. ``psycopg3`` will only be used if ``psycopg2`` is either absent or older than ``2.5.4``. +""" +from typing import Any, Optional + + __all__ = ['connect', 'quote_ident', 'quote_literal', 'DatabaseError', 'Error', 'OperationalError', 'ProgrammingError'] _legacy = False @@ -14,27 +22,77 @@ try: except ImportError: _legacy = True - def quote_literal(value, conn=None): + def quote_literal(value: Any, conn: Optional[Any] = None) -> str: + """Quote *value* as a SQL literal. + + .. note:: + *value* is quoted through ``psycopg`` adapters. + + :param value: value to be quoted. + :param conn: if a connection is given then :func:`quote_literal` checks if any special handling based on server + parameters needs to be applied to *value* before quoting it as a SQL literal. + + :returns: *value* quoted as a SQL literal. + """ value = adapt(value) if conn: value.prepare(conn) return value.getquoted().decode('utf-8') except ImportError: - from psycopg import connect as __connect, sql, Error, DatabaseError, OperationalError, ProgrammingError + from psycopg import connect as __connect, sql, Error, DatabaseError, OperationalError, ProgrammingError, Connection - def _connect(*args, **kwargs): - ret = __connect(*args, **kwargs) - ret.server_version = ret.pgconn.server_version # compatibility with psycopg2 + def _connect(dsn: str = "", **kwargs: Any) -> Any: + """Call ``psycopg.connect`` with ``dsn`` and ``**kwargs``. + + .. note:: + Will create ``server_version`` attribute in the returning connection, so it keeps compatibility with the + object that would be returned by ``psycopg2.connect``. + + :param dsn: DSN to call ``psycopg.connect`` with. + :param kwargs: keyword arguments to call ``psycopg.connect`` with. + + :returns: a connection to the database. + """ + ret = __connect(dsn, **kwargs) + setattr(ret, 'server_version', ret.pgconn.server_version) # compatibility with psycopg2 return ret - def _quote_ident(value, conn): + def _quote_ident(value: Any, conn: Connection) -> str: + """Quote *value* as a SQL identifier. + + :param value: value to be quoted. + :param conn: connection to evaluate the returning string into. + + :returns: *value* quoted as a SQL identifier. + """ return sql.Identifier(value).as_string(conn) - def quote_literal(value, conn=None): + def quote_literal(value: Any, conn: Optional[Any] = None) -> str: + """Quote *value* as a SQL literal. + + :param value: value to be quoted. + :param conn: connection to evaluate the returning string into. + + :returns: *value* quoted as a SQL literal. + """ return sql.Literal(value).as_string(conn) -def connect(*args, **kwargs): +def connect(*args: Any, **kwargs: Any) -> Any: + """Get a connection to the database. + + .. note:: + The connection will have ``autocommit`` enabled. + + It also enforces ``search_path=pg_catalog`` for non-replication connections to mitigate security issues as + Patroni relies on superuser connections. + + :param args: positional arguments to call ``connect`` function from ``psycopg`` module. + :param kwargs: keyword arguments to call ``connect`` function from ``psycopg`` module. + + :returns: a connection to the database. Can be either a :class:`psycopg.Connection` if using ``psycopg3``, or a + :class:`psycopg2.extensions.connection` if using ``psycopg2``. + """ if kwargs and 'replication' not in kwargs and kwargs.get('fallback_application_name') != 'Patroni ctl': options = [kwargs['options']] if 'options' in kwargs else [] options.append('-c search_path=pg_catalog') @@ -44,7 +102,15 @@ def connect(*args, **kwargs): return ret -def quote_ident(value, conn=None): +def quote_ident(value: Any, conn: Optional[Any] = None) -> str: + """Quote *value* as a SQL identifier. + + :param value: value to be quoted. + :param conn: connection to evaluate the returning string into. Can be either a :class:`psycopg.Connection` if + using ``psycopg3``, or a :class:`psycopg2.extensions.connection` if using ``psycopg2``. + + :returns: *value* quoted as a SQL identifier. + """ if _legacy or conn is None: return '"{0}"'.format(value.replace('"', '""')) return _quote_ident(value, conn)