""" psycopg connection objects """ # Copyright (C) 2020 The Psycopg Team import logging import threading from types import TracebackType from typing import Any, Callable, cast, Generator, Generic, Iterator from typing import List, NamedTuple, Optional, Type, TypeVar, Tuple, Union from typing import overload, TYPE_CHECKING from weakref import ref, ReferenceType from warnings import warn from functools import partial from contextlib import contextmanager from typing_extensions import TypeAlias from . import pq from . import errors as e from . import waiting from . import postgres from .abc import AdaptContext, ConnectionType, Params, Query, RV from .abc import PQGen, PQGenConn from .sql import Composable, SQL from ._tpc import Xid from .rows import Row, RowFactory, tuple_row, TupleRow, args_row from .adapt import AdaptersMap from ._enums import IsolationLevel from .cursor import Cursor from ._compat import LiteralString from .pq.misc import connection_summary from .conninfo import make_conninfo, conninfo_to_dict, ConnectionInfo from .conninfo import conninfo_attempts, ConnDict from ._pipeline import BasePipeline, Pipeline from .generators import notifies, connect, execute from ._encodings import pgconn_encoding from ._preparing import PrepareManager from .transaction import Transaction from .server_cursor import ServerCursor if TYPE_CHECKING: from .pq.abc import PGconn, PGresult from psycopg_pool.base import BasePool # Row Type variable for Cursor (when it needs to be distinguished from the # connection's one) CursorRow = TypeVar("CursorRow") TEXT = pq.Format.TEXT BINARY = pq.Format.BINARY OK = pq.ConnStatus.OK BAD = pq.ConnStatus.BAD COMMAND_OK = pq.ExecStatus.COMMAND_OK TUPLES_OK = pq.ExecStatus.TUPLES_OK FATAL_ERROR = pq.ExecStatus.FATAL_ERROR IDLE = pq.TransactionStatus.IDLE INTRANS = pq.TransactionStatus.INTRANS logger = logging.getLogger("psycopg") class Notify(NamedTuple): """An asynchronous notification received from the database.""" channel: str """The name of the channel on which the notification was received.""" payload: str """The message attached to the notification.""" pid: int """The PID of the backend process which sent the notification.""" Notify.__module__ = "psycopg" NoticeHandler: TypeAlias = Callable[[e.Diagnostic], None] NotifyHandler: TypeAlias = Callable[[Notify], None] class BaseConnection(Generic[Row]): """ Base class for different types of connections. Share common functionalities such as access to the wrapped PGconn, but allow different interfaces (sync/async). """ # DBAPI2 exposed exceptions Warning = e.Warning Error = e.Error InterfaceError = e.InterfaceError DatabaseError = e.DatabaseError DataError = e.DataError OperationalError = e.OperationalError IntegrityError = e.IntegrityError InternalError = e.InternalError ProgrammingError = e.ProgrammingError NotSupportedError = e.NotSupportedError # Enums useful for the connection ConnStatus = pq.ConnStatus TransactionStatus = pq.TransactionStatus # Default timeout for connection a attempt. # Arbitrary timeout, what applied by the libpq on my computer. # Your mileage won't vary. _DEFAULT_CONNECT_TIMEOUT = 130 def __init__(self, pgconn: "PGconn"): self.pgconn = pgconn self._autocommit = False # None, but set to a copy of the global adapters map as soon as requested. self._adapters: Optional[AdaptersMap] = None self._notice_handlers: List[NoticeHandler] = [] self._notify_handlers: List[NotifyHandler] = [] # Number of transaction blocks currently entered self._num_transactions = 0 self._closed = False # closed by an explicit close() self._prepared: PrepareManager = PrepareManager() self._tpc: Optional[Tuple[Xid, bool]] = None # xid, prepared wself = ref(self) pgconn.notice_handler = partial(BaseConnection._notice_handler, wself) pgconn.notify_handler = partial(BaseConnection._notify_handler, wself) # Attribute is only set if the connection is from a pool so we can tell # apart a connection in the pool too (when _pool = None) self._pool: Optional["BasePool[Any]"] self._pipeline: Optional[BasePipeline] = None # Time after which the connection should be closed self._expire_at: float self._isolation_level: Optional[IsolationLevel] = None self._read_only: Optional[bool] = None self._deferrable: Optional[bool] = None self._begin_statement = b"" def __del__(self) -> None: # If fails on connection we might not have this attribute yet if not hasattr(self, "pgconn"): return # Connection correctly closed if self.closed: return # Connection in a pool so terminating with the program is normal if hasattr(self, "_pool"): return warn( f"connection {self} was deleted while still open." " Please use 'with' or '.close()' to close the connection", ResourceWarning, ) def __repr__(self) -> str: cls = f"{self.__class__.__module__}.{self.__class__.__qualname__}" info = connection_summary(self.pgconn) return f"<{cls} {info} at 0x{id(self):x}>" @property def closed(self) -> bool: """`!True` if the connection is closed.""" return self.pgconn.status == BAD @property def broken(self) -> bool: """ `!True` if the connection was interrupted. A broken connection is always `closed`, but wasn't closed in a clean way, such as using `close()` or a `!with` block. """ return self.pgconn.status == BAD and not self._closed @property def autocommit(self) -> bool: """The autocommit state of the connection.""" return self._autocommit @autocommit.setter def autocommit(self, value: bool) -> None: self._set_autocommit(value) def _set_autocommit(self, value: bool) -> None: raise NotImplementedError def _set_autocommit_gen(self, value: bool) -> PQGen[None]: yield from self._check_intrans_gen("autocommit") self._autocommit = bool(value) @property def isolation_level(self) -> Optional[IsolationLevel]: """ The isolation level of the new transactions started on the connection. """ return self._isolation_level @isolation_level.setter def isolation_level(self, value: Optional[IsolationLevel]) -> None: self._set_isolation_level(value) def _set_isolation_level(self, value: Optional[IsolationLevel]) -> None: raise NotImplementedError def _set_isolation_level_gen(self, value: Optional[IsolationLevel]) -> PQGen[None]: yield from self._check_intrans_gen("isolation_level") self._isolation_level = IsolationLevel(value) if value is not None else None self._begin_statement = b"" @property def read_only(self) -> Optional[bool]: """ The read-only state of the new transactions started on the connection. """ return self._read_only @read_only.setter def read_only(self, value: Optional[bool]) -> None: self._set_read_only(value) def _set_read_only(self, value: Optional[bool]) -> None: raise NotImplementedError def _set_read_only_gen(self, value: Optional[bool]) -> PQGen[None]: yield from self._check_intrans_gen("read_only") self._read_only = bool(value) if value is not None else None self._begin_statement = b"" @property def deferrable(self) -> Optional[bool]: """ The deferrable state of the new transactions started on the connection. """ return self._deferrable @deferrable.setter def deferrable(self, value: Optional[bool]) -> None: self._set_deferrable(value) def _set_deferrable(self, value: Optional[bool]) -> None: raise NotImplementedError def _set_deferrable_gen(self, value: Optional[bool]) -> PQGen[None]: yield from self._check_intrans_gen("deferrable") self._deferrable = bool(value) if value is not None else None self._begin_statement = b"" def _check_intrans_gen(self, attribute: str) -> PQGen[None]: # Raise an exception if we are in a transaction status = self.pgconn.transaction_status if status == IDLE and self._pipeline: yield from self._pipeline._sync_gen() status = self.pgconn.transaction_status if status != IDLE: if self._num_transactions: raise e.ProgrammingError( f"can't change {attribute!r} now: " "connection.transaction() context in progress" ) else: raise e.ProgrammingError( f"can't change {attribute!r} now: " "connection in transaction status " f"{pq.TransactionStatus(status).name}" ) @property def info(self) -> ConnectionInfo: """A `ConnectionInfo` attribute to inspect connection properties.""" return ConnectionInfo(self.pgconn) @property def adapters(self) -> AdaptersMap: if not self._adapters: self._adapters = AdaptersMap(postgres.adapters) return self._adapters @property def connection(self) -> "BaseConnection[Row]": # implement the AdaptContext protocol return self def fileno(self) -> int: """Return the file descriptor of the connection. This function allows to use the connection as file-like object in functions waiting for readiness, such as the ones defined in the `selectors` module. """ return self.pgconn.socket def cancel(self) -> None: """Cancel the current operation on the connection.""" # No-op if the connection is closed # this allows to use the method as callback handler without caring # about its life. if self.closed: return if self._tpc and self._tpc[1]: raise e.ProgrammingError( "cancel() cannot be used with a prepared two-phase transaction" ) self._try_cancel(self.pgconn) @classmethod def _try_cancel(cls, pgconn: "PGconn") -> None: try: # Can fail if the connection is closed c = pgconn.get_cancel() except Exception as ex: logger.warning("couldn't try to cancel query: %s", ex) else: c.cancel() def add_notice_handler(self, callback: NoticeHandler) -> None: """ Register a callable to be invoked when a notice message is received. :param callback: the callback to call upon message received. :type callback: Callable[[~psycopg.errors.Diagnostic], None] """ self._notice_handlers.append(callback) def remove_notice_handler(self, callback: NoticeHandler) -> None: """ Unregister a notice message callable previously registered. :param callback: the callback to remove. :type callback: Callable[[~psycopg.errors.Diagnostic], None] """ self._notice_handlers.remove(callback) @staticmethod def _notice_handler( wself: "ReferenceType[BaseConnection[Row]]", res: "PGresult" ) -> None: self = wself() if not (self and self._notice_handlers): return diag = e.Diagnostic(res, pgconn_encoding(self.pgconn)) for cb in self._notice_handlers: try: cb(diag) except Exception as ex: logger.exception("error processing notice callback '%s': %s", cb, ex) def add_notify_handler(self, callback: NotifyHandler) -> None: """ Register a callable to be invoked whenever a notification is received. :param callback: the callback to call upon notification received. :type callback: Callable[[~psycopg.Notify], None] """ self._notify_handlers.append(callback) def remove_notify_handler(self, callback: NotifyHandler) -> None: """ Unregister a notification callable previously registered. :param callback: the callback to remove. :type callback: Callable[[~psycopg.Notify], None] """ self._notify_handlers.remove(callback) @staticmethod def _notify_handler( wself: "ReferenceType[BaseConnection[Row]]", pgn: pq.PGnotify ) -> None: self = wself() if not (self and self._notify_handlers): return enc = pgconn_encoding(self.pgconn) n = Notify(pgn.relname.decode(enc), pgn.extra.decode(enc), pgn.be_pid) for cb in self._notify_handlers: cb(n) @property def prepare_threshold(self) -> Optional[int]: """ Number of times a query is executed before it is prepared. - If it is set to 0, every query is prepared the first time it is executed. - If it is set to `!None`, prepared statements are disabled on the connection. Default value: 5 """ return self._prepared.prepare_threshold @prepare_threshold.setter def prepare_threshold(self, value: Optional[int]) -> None: self._prepared.prepare_threshold = value @property def prepared_max(self) -> int: """ Maximum number of prepared statements on the connection. Default value: 100 """ return self._prepared.prepared_max @prepared_max.setter def prepared_max(self, value: int) -> None: self._prepared.prepared_max = value # Generators to perform high-level operations on the connection # # These operations are expressed in terms of non-blocking generators # and the task of waiting when needed (when the generators yield) is left # to the connections subclass, which might wait either in blocking mode # or through asyncio. # # All these generators assume exclusive access to the connection: subclasses # should have a lock and hold it before calling and consuming them. @classmethod def _connect_gen( cls: Type[ConnectionType], conninfo: str = "", ) -> PQGenConn[ConnectionType]: """Generator to connect to the database and create a new instance.""" pgconn = yield from connect(conninfo) conn = cls(pgconn) return conn def _exec_command( self, command: Query, result_format: pq.Format = TEXT ) -> PQGen[Optional["PGresult"]]: """ Generator to send a command and receive the result to the backend. Only used to implement internal commands such as "commit", with eventual arguments bound client-side. The cursor can do more complex stuff. """ self._check_connection_ok() if isinstance(command, str): command = command.encode(pgconn_encoding(self.pgconn)) elif isinstance(command, Composable): command = command.as_bytes(self) if self._pipeline: cmd = partial( self.pgconn.send_query_params, command, None, result_format=result_format, ) self._pipeline.command_queue.append(cmd) self._pipeline.result_queue.append(None) return None self.pgconn.send_query_params(command, None, result_format=result_format) result = (yield from execute(self.pgconn))[-1] if result.status != COMMAND_OK and result.status != TUPLES_OK: if result.status == FATAL_ERROR: raise e.error_from_result(result, encoding=pgconn_encoding(self.pgconn)) else: raise e.InterfaceError( f"unexpected result {pq.ExecStatus(result.status).name}" f" from command {command.decode()!r}" ) return result def _check_connection_ok(self) -> None: if self.pgconn.status == OK: return if self.pgconn.status == BAD: raise e.OperationalError("the connection is closed") raise e.InterfaceError( "cannot execute operations: the connection is" f" in status {self.pgconn.status}" ) def _start_query(self) -> PQGen[None]: """Generator to start a transaction if necessary.""" if self._autocommit: return if self.pgconn.transaction_status != IDLE: return yield from self._exec_command(self._get_tx_start_command()) if self._pipeline: yield from self._pipeline._sync_gen() def _get_tx_start_command(self) -> bytes: if self._begin_statement: return self._begin_statement parts = [b"BEGIN"] if self.isolation_level is not None: val = IsolationLevel(self.isolation_level) parts.append(b"ISOLATION LEVEL") parts.append(val.name.replace("_", " ").encode()) if self.read_only is not None: parts.append(b"READ ONLY" if self.read_only else b"READ WRITE") if self.deferrable is not None: parts.append(b"DEFERRABLE" if self.deferrable else b"NOT DEFERRABLE") self._begin_statement = b" ".join(parts) return self._begin_statement def _commit_gen(self) -> PQGen[None]: """Generator implementing `Connection.commit()`.""" if self._num_transactions: raise e.ProgrammingError( "Explicit commit() forbidden within a Transaction " "context. (Transaction will be automatically committed " "on successful exit from context.)" ) if self._tpc: raise e.ProgrammingError( "commit() cannot be used during a two-phase transaction" ) if self.pgconn.transaction_status == IDLE: return yield from self._exec_command(b"COMMIT") if self._pipeline: yield from self._pipeline._sync_gen() def _rollback_gen(self) -> PQGen[None]: """Generator implementing `Connection.rollback()`.""" if self._num_transactions: raise e.ProgrammingError( "Explicit rollback() forbidden within a Transaction " "context. (Either raise Rollback() or allow " "an exception to propagate out of the context.)" ) if self._tpc: raise e.ProgrammingError( "rollback() cannot be used during a two-phase transaction" ) # Get out of a "pipeline aborted" state if self._pipeline: yield from self._pipeline._sync_gen() if self.pgconn.transaction_status == IDLE: return yield from self._exec_command(b"ROLLBACK") self._prepared.clear() for cmd in self._prepared.get_maintenance_commands(): yield from self._exec_command(cmd) if self._pipeline: yield from self._pipeline._sync_gen() def xid(self, format_id: int, gtrid: str, bqual: str) -> Xid: """ Returns a `Xid` to pass to the `!tpc_*()` methods of this connection. The argument types and constraints are explained in :ref:`two-phase-commit`. The values passed to the method will be available on the returned object as the members `~Xid.format_id`, `~Xid.gtrid`, `~Xid.bqual`. """ self._check_tpc() return Xid.from_parts(format_id, gtrid, bqual) def _tpc_begin_gen(self, xid: Union[Xid, str]) -> PQGen[None]: self._check_tpc() if not isinstance(xid, Xid): xid = Xid.from_string(xid) if self.pgconn.transaction_status != IDLE: raise e.ProgrammingError( "can't start two-phase transaction: connection in status" f" {pq.TransactionStatus(self.pgconn.transaction_status).name}" ) if self._autocommit: raise e.ProgrammingError( "can't use two-phase transactions in autocommit mode" ) self._tpc = (xid, False) yield from self._exec_command(self._get_tx_start_command()) def _tpc_prepare_gen(self) -> PQGen[None]: if not self._tpc: raise e.ProgrammingError( "'tpc_prepare()' must be called inside a two-phase transaction" ) if self._tpc[1]: raise e.ProgrammingError( "'tpc_prepare()' cannot be used during a prepared two-phase transaction" ) xid = self._tpc[0] self._tpc = (xid, True) yield from self._exec_command(SQL("PREPARE TRANSACTION {}").format(str(xid))) if self._pipeline: yield from self._pipeline._sync_gen() def _tpc_finish_gen( self, action: LiteralString, xid: Union[Xid, str, None] ) -> PQGen[None]: fname = f"tpc_{action.lower()}()" if xid is None: if not self._tpc: raise e.ProgrammingError( f"{fname} without xid must must be" " called inside a two-phase transaction" ) xid = self._tpc[0] else: if self._tpc: raise e.ProgrammingError( f"{fname} with xid must must be called" " outside a two-phase transaction" ) if not isinstance(xid, Xid): xid = Xid.from_string(xid) if self._tpc and not self._tpc[1]: meth: Callable[[], PQGen[None]] meth = getattr(self, f"_{action.lower()}_gen") self._tpc = None yield from meth() else: yield from self._exec_command( SQL("{} PREPARED {}").format(SQL(action), str(xid)) ) self._tpc = None def _check_tpc(self) -> None: """Raise NotSupportedError if TPC is not supported.""" # TPC supported on every supported PostgreSQL version. pass class Connection(BaseConnection[Row]): """ Wrapper for a connection to the database. """ __module__ = "psycopg" cursor_factory: Type[Cursor[Row]] server_cursor_factory: Type[ServerCursor[Row]] row_factory: RowFactory[Row] _pipeline: Optional[Pipeline] _Self = TypeVar("_Self", bound="Connection[Any]") def __init__( self, pgconn: "PGconn", row_factory: RowFactory[Row] = cast(RowFactory[Row], tuple_row), ): super().__init__(pgconn) self.row_factory = row_factory self.lock = threading.Lock() self.cursor_factory = Cursor self.server_cursor_factory = ServerCursor @overload @classmethod def connect( cls, conninfo: str = "", *, autocommit: bool = False, row_factory: RowFactory[Row], prepare_threshold: Optional[int] = 5, cursor_factory: Optional[Type[Cursor[Row]]] = None, context: Optional[AdaptContext] = None, **kwargs: Union[None, int, str], ) -> "Connection[Row]": # TODO: returned type should be _Self. See #308. ... @overload @classmethod def connect( cls, conninfo: str = "", *, autocommit: bool = False, prepare_threshold: Optional[int] = 5, cursor_factory: Optional[Type[Cursor[Any]]] = None, context: Optional[AdaptContext] = None, **kwargs: Union[None, int, str], ) -> "Connection[TupleRow]": ... @classmethod # type: ignore[misc] # https://github.com/python/mypy/issues/11004 def connect( cls, conninfo: str = "", *, autocommit: bool = False, prepare_threshold: Optional[int] = 5, row_factory: Optional[RowFactory[Row]] = None, cursor_factory: Optional[Type[Cursor[Row]]] = None, context: Optional[AdaptContext] = None, **kwargs: Any, ) -> "Connection[Any]": """ Connect to a database server and return a new `Connection` instance. """ params = cls._get_connection_params(conninfo, **kwargs) timeout = int(params["connect_timeout"]) rv = None for attempt in conninfo_attempts(params): try: conninfo = make_conninfo(**attempt) rv = cls._wait_conn(cls._connect_gen(conninfo), timeout=timeout) break except e._NO_TRACEBACK as ex: last_ex = ex if not rv: assert last_ex raise last_ex.with_traceback(None) rv._autocommit = bool(autocommit) if row_factory: rv.row_factory = row_factory if cursor_factory: rv.cursor_factory = cursor_factory if context: rv._adapters = AdaptersMap(context.adapters) rv.prepare_threshold = prepare_threshold return rv def __enter__(self: _Self) -> _Self: return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: if self.closed: return if exc_type: # try to rollback, but if there are problems (connection in a bad # state) just warn without clobbering the exception bubbling up. try: self.rollback() except Exception as exc2: logger.warning( "error ignored in rollback on %s: %s", self, exc2, ) else: self.commit() # Close the connection only if it doesn't belong to a pool. if not getattr(self, "_pool", None): self.close() @classmethod def _get_connection_params(cls, conninfo: str, **kwargs: Any) -> ConnDict: """Manipulate connection parameters before connecting. :param conninfo: Connection string as received by `~Connection.connect()`. :param kwargs: Overriding connection arguments as received by `!connect()`. :return: Connection arguments merged and eventually modified, in a format similar to `~conninfo.conninfo_to_dict()`. """ params = conninfo_to_dict(conninfo, **kwargs) # Make sure there is an usable connect_timeout if "connect_timeout" in params: params["connect_timeout"] = int(params["connect_timeout"]) else: # The sync connect function will stop on the default socket timeout # Because in async connection mode we need to enforce the timeout # ourselves, we need a finite value. params["connect_timeout"] = cls._DEFAULT_CONNECT_TIMEOUT return params def close(self) -> None: """Close the database connection.""" if self.closed: return self._closed = True # TODO: maybe send a cancel on close, if the connection is ACTIVE? self.pgconn.finish() @overload def cursor(self, *, binary: bool = False) -> Cursor[Row]: ... @overload def cursor( self, *, binary: bool = False, row_factory: RowFactory[CursorRow] ) -> Cursor[CursorRow]: ... @overload def cursor( self, name: str, *, binary: bool = False, scrollable: Optional[bool] = None, withhold: bool = False, ) -> ServerCursor[Row]: ... @overload def cursor( self, name: str, *, binary: bool = False, row_factory: RowFactory[CursorRow], scrollable: Optional[bool] = None, withhold: bool = False, ) -> ServerCursor[CursorRow]: ... def cursor( self, name: str = "", *, binary: bool = False, row_factory: Optional[RowFactory[Any]] = None, scrollable: Optional[bool] = None, withhold: bool = False, ) -> Union[Cursor[Any], ServerCursor[Any]]: """ Return a new cursor to send commands and queries to the connection. """ self._check_connection_ok() if not row_factory: row_factory = self.row_factory cur: Union[Cursor[Any], ServerCursor[Any]] if name: cur = self.server_cursor_factory( self, name=name, row_factory=row_factory, scrollable=scrollable, withhold=withhold, ) else: cur = self.cursor_factory(self, row_factory=row_factory) if binary: cur.format = BINARY return cur def execute( self, query: Query, params: Optional[Params] = None, *, prepare: Optional[bool] = None, binary: bool = False, ) -> Cursor[Row]: """Execute a query and return a cursor to read its results.""" try: cur = self.cursor() if binary: cur.format = BINARY return cur.execute(query, params, prepare=prepare) except e._NO_TRACEBACK as ex: raise ex.with_traceback(None) def commit(self) -> None: """Commit any pending transaction to the database.""" with self.lock: self.wait(self._commit_gen()) def rollback(self) -> None: """Roll back to the start of any pending transaction.""" with self.lock: self.wait(self._rollback_gen()) @contextmanager def transaction( self, savepoint_name: Optional[str] = None, force_rollback: bool = False, ) -> Iterator[Transaction]: """ Start a context block with a new transaction or nested transaction. :param savepoint_name: Name of the savepoint used to manage a nested transaction. If `!None`, one will be chosen automatically. :param force_rollback: Roll back the transaction at the end of the block even if there were no error (e.g. to try a no-op process). :rtype: Transaction """ tx = Transaction(self, savepoint_name, force_rollback) if self._pipeline: with self.pipeline(), tx, self.pipeline(): yield tx else: with tx: yield tx def notifies(self) -> Generator[Notify, None, None]: """ Yield `Notify` objects as soon as they are received from the database. """ while True: with self.lock: try: ns = self.wait(notifies(self.pgconn)) except e._NO_TRACEBACK as ex: raise ex.with_traceback(None) enc = pgconn_encoding(self.pgconn) for pgn in ns: n = Notify(pgn.relname.decode(enc), pgn.extra.decode(enc), pgn.be_pid) yield n @contextmanager def pipeline(self) -> Iterator[Pipeline]: """Switch the connection into pipeline mode.""" with self.lock: self._check_connection_ok() pipeline = self._pipeline if pipeline is None: # WARNING: reference loop, broken ahead. pipeline = self._pipeline = Pipeline(self) try: with pipeline: yield pipeline finally: if pipeline.level == 0: with self.lock: assert pipeline is self._pipeline self._pipeline = None def wait(self, gen: PQGen[RV], timeout: Optional[float] = 0.1) -> RV: """ Consume a generator operating on the connection. The function must be used on generators that don't change connection fd (i.e. not on connect and reset). """ try: return waiting.wait(gen, self.pgconn.socket, timeout=timeout) except KeyboardInterrupt: # On Ctrl-C, try to cancel the query in the server, otherwise # the connection will remain stuck in ACTIVE state. self._try_cancel(self.pgconn) try: waiting.wait(gen, self.pgconn.socket, timeout=timeout) except e.QueryCanceled: pass # as expected raise @classmethod def _wait_conn(cls, gen: PQGenConn[RV], timeout: Optional[int]) -> RV: """Consume a connection generator.""" return waiting.wait_conn(gen, timeout=timeout) def _set_autocommit(self, value: bool) -> None: with self.lock: self.wait(self._set_autocommit_gen(value)) def _set_isolation_level(self, value: Optional[IsolationLevel]) -> None: with self.lock: self.wait(self._set_isolation_level_gen(value)) def _set_read_only(self, value: Optional[bool]) -> None: with self.lock: self.wait(self._set_read_only_gen(value)) def _set_deferrable(self, value: Optional[bool]) -> None: with self.lock: self.wait(self._set_deferrable_gen(value)) def tpc_begin(self, xid: Union[Xid, str]) -> None: """ Begin a TPC transaction with the given transaction ID `!xid`. """ with self.lock: self.wait(self._tpc_begin_gen(xid)) def tpc_prepare(self) -> None: """ Perform the first phase of a transaction started with `tpc_begin()`. """ try: with self.lock: self.wait(self._tpc_prepare_gen()) except e.ObjectNotInPrerequisiteState as ex: raise e.NotSupportedError(str(ex)) from None def tpc_commit(self, xid: Union[Xid, str, None] = None) -> None: """ Commit a prepared two-phase transaction. """ with self.lock: self.wait(self._tpc_finish_gen("COMMIT", xid)) def tpc_rollback(self, xid: Union[Xid, str, None] = None) -> None: """ Roll back a prepared two-phase transaction. """ with self.lock: self.wait(self._tpc_finish_gen("ROLLBACK", xid)) def tpc_recover(self) -> List[Xid]: self._check_tpc() status = self.info.transaction_status with self.cursor(row_factory=args_row(Xid._from_record)) as cur: cur.execute(Xid._get_recover_query()) res = cur.fetchall() if status == IDLE and self.info.transaction_status == INTRANS: self.rollback() return res