codekingpro/portable-devtools

1"""2psycopg connection objects3"""4 5# Copyright (C) 2020 The Psycopg Team6 7import logging8import threading9from types import TracebackType10from typing import Any, Callable, cast, Dict, Generator, Generic, Iterator11from typing import List, NamedTuple, Optional, Type, TypeVar, Tuple, Union12from typing import overload, TYPE_CHECKING13from weakref import ref, ReferenceType14from warnings import warn15from functools import partial16from contextlib import contextmanager17from typing_extensions import TypeAlias18 19from . import pq20from . import errors as e21from . import waiting22from . import postgres23from .abc import AdaptContext, ConnectionType, Params, Query, RV24from .abc import PQGen, PQGenConn25from .sql import Composable, SQL26from ._tpc import Xid27from .rows import Row, RowFactory, tuple_row, TupleRow, args_row28from .adapt import AdaptersMap29from ._enums import IsolationLevel30from .cursor import Cursor31from ._compat import LiteralString32from .pq.misc import connection_summary33from .conninfo import make_conninfo, conninfo_to_dict, ConnectionInfo34from ._pipeline import BasePipeline, Pipeline35from .generators import notifies, connect, execute36from ._encodings import pgconn_encoding37from ._preparing import PrepareManager38from .transaction import Transaction39from .server_cursor import ServerCursor40 41if TYPE_CHECKING:42    from .pq.abc import PGconn, PGresult43    from psycopg_pool.base import BasePool44 45 46# Row Type variable for Cursor (when it needs to be distinguished from the47# connection's one)48CursorRow = TypeVar("CursorRow")49 50TEXT = pq.Format.TEXT51BINARY = pq.Format.BINARY52 53OK = pq.ConnStatus.OK54BAD = pq.ConnStatus.BAD55 56COMMAND_OK = pq.ExecStatus.COMMAND_OK57TUPLES_OK = pq.ExecStatus.TUPLES_OK58FATAL_ERROR = pq.ExecStatus.FATAL_ERROR59 60IDLE = pq.TransactionStatus.IDLE61INTRANS = pq.TransactionStatus.INTRANS62 63logger = logging.getLogger("psycopg")64 65 66class Notify(NamedTuple):67    """An asynchronous notification received from the database."""68 69    channel: str70    """The name of the channel on which the notification was received."""71 72    payload: str73    """The message attached to the notification."""74 75    pid: int76    """The PID of the backend process which sent the notification."""77 78 79Notify.__module__ = "psycopg"80 81NoticeHandler: TypeAlias = Callable[[e.Diagnostic], None]82NotifyHandler: TypeAlias = Callable[[Notify], None]83 84 85class BaseConnection(Generic[Row]):86    """87    Base class for different types of connections.88 89    Share common functionalities such as access to the wrapped PGconn, but90    allow different interfaces (sync/async).91    """92 93    # DBAPI2 exposed exceptions94    Warning = e.Warning95    Error = e.Error96    InterfaceError = e.InterfaceError97    DatabaseError = e.DatabaseError98    DataError = e.DataError99    OperationalError = e.OperationalError100    IntegrityError = e.IntegrityError101    InternalError = e.InternalError102    ProgrammingError = e.ProgrammingError103    NotSupportedError = e.NotSupportedError104 105    # Enums useful for the connection106    ConnStatus = pq.ConnStatus107    TransactionStatus = pq.TransactionStatus108 109    def __init__(self, pgconn: "PGconn"):110        self.pgconn = pgconn111        self._autocommit = False112 113        # None, but set to a copy of the global adapters map as soon as requested.114        self._adapters: Optional[AdaptersMap] = None115 116        self._notice_handlers: List[NoticeHandler] = []117        self._notify_handlers: List[NotifyHandler] = []118 119        # Number of transaction blocks currently entered120        self._num_transactions = 0121 122        self._closed = False  # closed by an explicit close()123        self._prepared: PrepareManager = PrepareManager()124        self._tpc: Optional[Tuple[Xid, bool]] = None  # xid, prepared125 126        wself = ref(self)127        pgconn.notice_handler = partial(BaseConnection._notice_handler, wself)128        pgconn.notify_handler = partial(BaseConnection._notify_handler, wself)129 130        # Attribute is only set if the connection is from a pool so we can tell131        # apart a connection in the pool too (when _pool = None)132        self._pool: Optional["BasePool[Any]"]133 134        self._pipeline: Optional[BasePipeline] = None135 136        # Time after which the connection should be closed137        self._expire_at: float138 139        self._isolation_level: Optional[IsolationLevel] = None140        self._read_only: Optional[bool] = None141        self._deferrable: Optional[bool] = None142        self._begin_statement = b""143 144    def __del__(self) -> None:145        # If fails on connection we might not have this attribute yet146        if not hasattr(self, "pgconn"):147            return148 149        # Connection correctly closed150        if self.closed:151            return152 153        # Connection in a pool so terminating with the program is normal154        if hasattr(self, "_pool"):155            return156 157        warn(158            f"connection {self} was deleted while still open."159            " Please use 'with' or '.close()' to close the connection",160            ResourceWarning,161        )162 163    def __repr__(self) -> str:164        cls = f"{self.__class__.__module__}.{self.__class__.__qualname__}"165        info = connection_summary(self.pgconn)166        return f"<{cls} {info} at 0x{id(self):x}>"167 168    @property169    def closed(self) -> bool:170        """`!True` if the connection is closed."""171        return self.pgconn.status == BAD172 173    @property174    def broken(self) -> bool:175        """176        `!True` if the connection was interrupted.177 178        A broken connection is always `closed`, but wasn't closed in a clean179        way, such as using `close()` or a `!with` block.180        """181        return self.pgconn.status == BAD and not self._closed182 183    @property184    def autocommit(self) -> bool:185        """The autocommit state of the connection."""186        return self._autocommit187 188    @autocommit.setter189    def autocommit(self, value: bool) -> None:190        self._set_autocommit(value)191 192    def _set_autocommit(self, value: bool) -> None:193        raise NotImplementedError194 195    def _set_autocommit_gen(self, value: bool) -> PQGen[None]:196        yield from self._check_intrans_gen("autocommit")197        self._autocommit = bool(value)198 199    @property200    def isolation_level(self) -> Optional[IsolationLevel]:201        """202        The isolation level of the new transactions started on the connection.203        """204        return self._isolation_level205 206    @isolation_level.setter207    def isolation_level(self, value: Optional[IsolationLevel]) -> None:208        self._set_isolation_level(value)209 210    def _set_isolation_level(self, value: Optional[IsolationLevel]) -> None:211        raise NotImplementedError212 213    def _set_isolation_level_gen(self, value: Optional[IsolationLevel]) -> PQGen[None]:214        yield from self._check_intrans_gen("isolation_level")215        self._isolation_level = IsolationLevel(value) if value is not None else None216        self._begin_statement = b""217 218    @property219    def read_only(self) -> Optional[bool]:220        """221        The read-only state of the new transactions started on the connection.222        """223        return self._read_only224 225    @read_only.setter226    def read_only(self, value: Optional[bool]) -> None:227        self._set_read_only(value)228 229    def _set_read_only(self, value: Optional[bool]) -> None:230        raise NotImplementedError231 232    def _set_read_only_gen(self, value: Optional[bool]) -> PQGen[None]:233        yield from self._check_intrans_gen("read_only")234        self._read_only = bool(value) if value is not None else None235        self._begin_statement = b""236 237    @property238    def deferrable(self) -> Optional[bool]:239        """240        The deferrable state of the new transactions started on the connection.241        """242        return self._deferrable243 244    @deferrable.setter245    def deferrable(self, value: Optional[bool]) -> None:246        self._set_deferrable(value)247 248    def _set_deferrable(self, value: Optional[bool]) -> None:249        raise NotImplementedError250 251    def _set_deferrable_gen(self, value: Optional[bool]) -> PQGen[None]:252        yield from self._check_intrans_gen("deferrable")253        self._deferrable = bool(value) if value is not None else None254        self._begin_statement = b""255 256    def _check_intrans_gen(self, attribute: str) -> PQGen[None]:257        # Raise an exception if we are in a transaction258        status = self.pgconn.transaction_status259        if status == IDLE and self._pipeline:260            yield from self._pipeline._sync_gen()261            status = self.pgconn.transaction_status262        if status != IDLE:263            if self._num_transactions:264                raise e.ProgrammingError(265                    f"can't change {attribute!r} now: "266                    "connection.transaction() context in progress"267                )268            else:269                raise e.ProgrammingError(270                    f"can't change {attribute!r} now: "271                    "connection in transaction status "272                    f"{pq.TransactionStatus(status).name}"273                )274 275    @property276    def info(self) -> ConnectionInfo:277        """A `ConnectionInfo` attribute to inspect connection properties."""278        return ConnectionInfo(self.pgconn)279 280    @property281    def adapters(self) -> AdaptersMap:282        if not self._adapters:283            self._adapters = AdaptersMap(postgres.adapters)284 285        return self._adapters286 287    @property288    def connection(self) -> "BaseConnection[Row]":289        # implement the AdaptContext protocol290        return self291 292    def fileno(self) -> int:293        """Return the file descriptor of the connection.294 295        This function allows to use the connection as file-like object in296        functions waiting for readiness, such as the ones defined in the297        `selectors` module.298        """299        return self.pgconn.socket300 301    def cancel(self) -> None:302        """Cancel the current operation on the connection."""303        # No-op if the connection is closed304        # this allows to use the method as callback handler without caring305        # about its life.306        if self.closed:307            return308 309        if self._tpc and self._tpc[1]:310            raise e.ProgrammingError(311                "cancel() cannot be used with a prepared two-phase transaction"312            )313 314        self._try_cancel(self.pgconn)315 316    @classmethod317    def _try_cancel(cls, pgconn: "PGconn") -> None:318        try:319            # Can fail if the connection is closed320            c = pgconn.get_cancel()321        except Exception as ex:322            logger.warning("couldn't try to cancel query: %s", ex)323        else:324            c.cancel()325 326    def add_notice_handler(self, callback: NoticeHandler) -> None:327        """328        Register a callable to be invoked when a notice message is received.329 330        :param callback: the callback to call upon message received.331        :type callback: Callable[[~psycopg.errors.Diagnostic], None]332        """333        self._notice_handlers.append(callback)334 335    def remove_notice_handler(self, callback: NoticeHandler) -> None:336        """337        Unregister a notice message callable previously registered.338 339        :param callback: the callback to remove.340        :type callback: Callable[[~psycopg.errors.Diagnostic], None]341        """342        self._notice_handlers.remove(callback)343 344    @staticmethod345    def _notice_handler(346        wself: "ReferenceType[BaseConnection[Row]]", res: "PGresult"347    ) -> None:348        self = wself()349        if not (self and self._notice_handlers):350            return351 352        diag = e.Diagnostic(res, pgconn_encoding(self.pgconn))353        for cb in self._notice_handlers:354            try:355                cb(diag)356            except Exception as ex:357                logger.exception("error processing notice callback '%s': %s", cb, ex)358 359    def add_notify_handler(self, callback: NotifyHandler) -> None:360        """361        Register a callable to be invoked whenever a notification is received.362 363        :param callback: the callback to call upon notification received.364        :type callback: Callable[[~psycopg.Notify], None]365        """366        self._notify_handlers.append(callback)367 368    def remove_notify_handler(self, callback: NotifyHandler) -> None:369        """370        Unregister a notification callable previously registered.371 372        :param callback: the callback to remove.373        :type callback: Callable[[~psycopg.Notify], None]374        """375        self._notify_handlers.remove(callback)376 377    @staticmethod378    def _notify_handler(379        wself: "ReferenceType[BaseConnection[Row]]", pgn: pq.PGnotify380    ) -> None:381        self = wself()382        if not (self and self._notify_handlers):383            return384 385        enc = pgconn_encoding(self.pgconn)386        n = Notify(pgn.relname.decode(enc), pgn.extra.decode(enc), pgn.be_pid)387        for cb in self._notify_handlers:388            cb(n)389 390    @property391    def prepare_threshold(self) -> Optional[int]:392        """393        Number of times a query is executed before it is prepared.394 395        - If it is set to 0, every query is prepared the first time it is396          executed.397        - If it is set to `!None`, prepared statements are disabled on the398          connection.399 400        Default value: 5401        """402        return self._prepared.prepare_threshold403 404    @prepare_threshold.setter405    def prepare_threshold(self, value: Optional[int]) -> None:406        self._prepared.prepare_threshold = value407 408    @property409    def prepared_max(self) -> int:410        """411        Maximum number of prepared statements on the connection.412 413        Default value: 100414        """415        return self._prepared.prepared_max416 417    @prepared_max.setter418    def prepared_max(self, value: int) -> None:419        self._prepared.prepared_max = value420 421    # Generators to perform high-level operations on the connection422    #423    # These operations are expressed in terms of non-blocking generators424    # and the task of waiting when needed (when the generators yield) is left425    # to the connections subclass, which might wait either in blocking mode426    # or through asyncio.427    #428    # All these generators assume exclusive access to the connection: subclasses429    # should have a lock and hold it before calling and consuming them.430 431    @classmethod432    def _connect_gen(433        cls: Type[ConnectionType],434        conninfo: str = "",435        *,436        autocommit: bool = False,437    ) -> PQGenConn[ConnectionType]:438        """Generator to connect to the database and create a new instance."""439        pgconn = yield from connect(conninfo)440        conn = cls(pgconn)441        conn._autocommit = bool(autocommit)442        return conn443 444    def _exec_command(445        self, command: Query, result_format: pq.Format = TEXT446    ) -> PQGen[Optional["PGresult"]]:447        """448        Generator to send a command and receive the result to the backend.449 450        Only used to implement internal commands such as "commit", with eventual451        arguments bound client-side. The cursor can do more complex stuff.452        """453        self._check_connection_ok()454 455        if isinstance(command, str):456            command = command.encode(pgconn_encoding(self.pgconn))457        elif isinstance(command, Composable):458            command = command.as_bytes(self)459 460        if self._pipeline:461            cmd = partial(462                self.pgconn.send_query_params,463                command,464                None,465                result_format=result_format,466            )467            self._pipeline.command_queue.append(cmd)468            self._pipeline.result_queue.append(None)469            return None470 471        self.pgconn.send_query_params(command, None, result_format=result_format)472 473        result = (yield from execute(self.pgconn))[-1]474        if result.status != COMMAND_OK and result.status != TUPLES_OK:475            if result.status == FATAL_ERROR:476                raise e.error_from_result(result, encoding=pgconn_encoding(self.pgconn))477            else:478                raise e.InterfaceError(479                    f"unexpected result {pq.ExecStatus(result.status).name}"480                    f" from command {command.decode()!r}"481                )482        return result483 484    def _check_connection_ok(self) -> None:485        if self.pgconn.status == OK:486            return487 488        if self.pgconn.status == BAD:489            raise e.OperationalError("the connection is closed")490        raise e.InterfaceError(491            "cannot execute operations: the connection is"492            f" in status {self.pgconn.status}"493        )494 495    def _start_query(self) -> PQGen[None]:496        """Generator to start a transaction if necessary."""497        if self._autocommit:498            return499 500        if self.pgconn.transaction_status != IDLE:501            return502 503        yield from self._exec_command(self._get_tx_start_command())504        if self._pipeline:505            yield from self._pipeline._sync_gen()506 507    def _get_tx_start_command(self) -> bytes:508        if self._begin_statement:509            return self._begin_statement510 511        parts = [b"BEGIN"]512 513        if self.isolation_level is not None:514            val = IsolationLevel(self.isolation_level)515            parts.append(b"ISOLATION LEVEL")516            parts.append(val.name.replace("_", " ").encode())517 518        if self.read_only is not None:519            parts.append(b"READ ONLY" if self.read_only else b"READ WRITE")520 521        if self.deferrable is not None:522            parts.append(b"DEFERRABLE" if self.deferrable else b"NOT DEFERRABLE")523 524        self._begin_statement = b" ".join(parts)525        return self._begin_statement526 527    def _commit_gen(self) -> PQGen[None]:528        """Generator implementing `Connection.commit()`."""529        if self._num_transactions:530            raise e.ProgrammingError(531                "Explicit commit() forbidden within a Transaction "532                "context. (Transaction will be automatically committed "533                "on successful exit from context.)"534            )535        if self._tpc:536            raise e.ProgrammingError(537                "commit() cannot be used during a two-phase transaction"538            )539        if self.pgconn.transaction_status == IDLE:540            return541 542        yield from self._exec_command(b"COMMIT")543 544        if self._pipeline:545            yield from self._pipeline._sync_gen()546 547    def _rollback_gen(self) -> PQGen[None]:548        """Generator implementing `Connection.rollback()`."""549        if self._num_transactions:550            raise e.ProgrammingError(551                "Explicit rollback() forbidden within a Transaction "552                "context. (Either raise Rollback() or allow "553                "an exception to propagate out of the context.)"554            )555        if self._tpc:556            raise e.ProgrammingError(557                "rollback() cannot be used during a two-phase transaction"558            )559 560        # Get out of a "pipeline aborted" state561        if self._pipeline:562            yield from self._pipeline._sync_gen()563 564        if self.pgconn.transaction_status == IDLE:565            return566 567        yield from self._exec_command(b"ROLLBACK")568        self._prepared.clear()569        for cmd in self._prepared.get_maintenance_commands():570            yield from self._exec_command(cmd)571 572        if self._pipeline:573            yield from self._pipeline._sync_gen()574 575    def xid(self, format_id: int, gtrid: str, bqual: str) -> Xid:576        """577        Returns a `Xid` to pass to the `!tpc_*()` methods of this connection.578 579        The argument types and constraints are explained in580        :ref:`two-phase-commit`.581 582        The values passed to the method will be available on the returned583        object as the members `~Xid.format_id`, `~Xid.gtrid`, `~Xid.bqual`.584        """585        self._check_tpc()586        return Xid.from_parts(format_id, gtrid, bqual)587 588    def _tpc_begin_gen(self, xid: Union[Xid, str]) -> PQGen[None]:589        self._check_tpc()590 591        if not isinstance(xid, Xid):592            xid = Xid.from_string(xid)593 594        if self.pgconn.transaction_status != IDLE:595            raise e.ProgrammingError(596                "can't start two-phase transaction: connection in status"597                f" {pq.TransactionStatus(self.pgconn.transaction_status).name}"598            )599 600        if self._autocommit:601            raise e.ProgrammingError(602                "can't use two-phase transactions in autocommit mode"603            )604 605        self._tpc = (xid, False)606        yield from self._exec_command(self._get_tx_start_command())607 608    def _tpc_prepare_gen(self) -> PQGen[None]:609        if not self._tpc:610            raise e.ProgrammingError(611                "'tpc_prepare()' must be called inside a two-phase transaction"612            )613        if self._tpc[1]:614            raise e.ProgrammingError(615                "'tpc_prepare()' cannot be used during a prepared two-phase transaction"616            )617        xid = self._tpc[0]618        self._tpc = (xid, True)619        yield from self._exec_command(SQL("PREPARE TRANSACTION {}").format(str(xid)))620        if self._pipeline:621            yield from self._pipeline._sync_gen()622 623    def _tpc_finish_gen(624        self, action: LiteralString, xid: Union[Xid, str, None]625    ) -> PQGen[None]:626        fname = f"tpc_{action.lower()}()"627        if xid is None:628            if not self._tpc:629                raise e.ProgrammingError(630                    f"{fname} without xid must must be"631                    " called inside a two-phase transaction"632                )633            xid = self._tpc[0]634        else:635            if self._tpc:636                raise e.ProgrammingError(637                    f"{fname} with xid must must be called"638                    " outside a two-phase transaction"639                )640            if not isinstance(xid, Xid):641                xid = Xid.from_string(xid)642 643        if self._tpc and not self._tpc[1]:644            meth: Callable[[], PQGen[None]]645            meth = getattr(self, f"_{action.lower()}_gen")646            self._tpc = None647            yield from meth()648        else:649            yield from self._exec_command(650                SQL("{} PREPARED {}").format(SQL(action), str(xid))651            )652            self._tpc = None653 654    def _check_tpc(self) -> None:655        """Raise NotSupportedError if TPC is not supported."""656        # TPC supported on every supported PostgreSQL version.657        pass658 659 660class Connection(BaseConnection[Row]):661    """662    Wrapper for a connection to the database.663    """664 665    __module__ = "psycopg"666 667    cursor_factory: Type[Cursor[Row]]668    server_cursor_factory: Type[ServerCursor[Row]]669    row_factory: RowFactory[Row]670    _pipeline: Optional[Pipeline]671    _Self = TypeVar("_Self", bound="Connection[Any]")672 673    def __init__(674        self,675        pgconn: "PGconn",676        row_factory: RowFactory[Row] = cast(RowFactory[Row], tuple_row),677    ):678        super().__init__(pgconn)679        self.row_factory = row_factory680        self.lock = threading.Lock()681        self.cursor_factory = Cursor682        self.server_cursor_factory = ServerCursor683 684    @overload685    @classmethod686    def connect(687        cls,688        conninfo: str = "",689        *,690        autocommit: bool = False,691        row_factory: RowFactory[Row],692        prepare_threshold: Optional[int] = 5,693        cursor_factory: Optional[Type[Cursor[Row]]] = None,694        context: Optional[AdaptContext] = None,695        **kwargs: Union[None, int, str],696    ) -> "Connection[Row]":697        # TODO: returned type should be _Self. See #308.698        ...699 700    @overload701    @classmethod702    def connect(703        cls,704        conninfo: str = "",705        *,706        autocommit: bool = False,707        prepare_threshold: Optional[int] = 5,708        cursor_factory: Optional[Type[Cursor[Any]]] = None,709        context: Optional[AdaptContext] = None,710        **kwargs: Union[None, int, str],711    ) -> "Connection[TupleRow]":712        ...713 714    @classmethod  # type: ignore[misc] # https://github.com/python/mypy/issues/11004715    def connect(716        cls,717        conninfo: str = "",718        *,719        autocommit: bool = False,720        prepare_threshold: Optional[int] = 5,721        row_factory: Optional[RowFactory[Row]] = None,722        cursor_factory: Optional[Type[Cursor[Row]]] = None,723        context: Optional[AdaptContext] = None,724        **kwargs: Any,725    ) -> "Connection[Any]":726        """727        Connect to a database server and return a new `Connection` instance.728        """729        params = cls._get_connection_params(conninfo, **kwargs)730        conninfo = make_conninfo(**params)731 732        try:733            rv = cls._wait_conn(734                cls._connect_gen(conninfo, autocommit=autocommit),735                timeout=params["connect_timeout"],736            )737        except e._NO_TRACEBACK as ex:738            raise ex.with_traceback(None)739 740        if row_factory:741            rv.row_factory = row_factory742        if cursor_factory:743            rv.cursor_factory = cursor_factory744        if context:745            rv._adapters = AdaptersMap(context.adapters)746        rv.prepare_threshold = prepare_threshold747        return rv748 749    def __enter__(self: _Self) -> _Self:750        return self751 752    def __exit__(753        self,754        exc_type: Optional[Type[BaseException]],755        exc_val: Optional[BaseException],756        exc_tb: Optional[TracebackType],757    ) -> None:758        if self.closed:759            return760 761        if exc_type:762            # try to rollback, but if there are problems (connection in a bad763            # state) just warn without clobbering the exception bubbling up.764            try:765                self.rollback()766            except Exception as exc2:767                logger.warning(768                    "error ignored in rollback on %s: %s",769                    self,770                    exc2,771                )772        else:773            self.commit()774 775        # Close the connection only if it doesn't belong to a pool.776        if not getattr(self, "_pool", None):777            self.close()778 779    @classmethod780    def _get_connection_params(cls, conninfo: str, **kwargs: Any) -> Dict[str, Any]:781        """Manipulate connection parameters before connecting.782 783        :param conninfo: Connection string as received by `~Connection.connect()`.784        :param kwargs: Overriding connection arguments as received by `!connect()`.785        :return: Connection arguments merged and eventually modified, in a786            format similar to `~conninfo.conninfo_to_dict()`.787        """788        params = conninfo_to_dict(conninfo, **kwargs)789 790        # Make sure there is an usable connect_timeout791        if "connect_timeout" in params:792            params["connect_timeout"] = int(params["connect_timeout"])793        else:794            params["connect_timeout"] = None795 796        return params797 798    def close(self) -> None:799        """Close the database connection."""800        if self.closed:801            return802        self._closed = True803 804        # TODO: maybe send a cancel on close, if the connection is ACTIVE?805 806        self.pgconn.finish()807 808    @overload809    def cursor(self, *, binary: bool = False) -> Cursor[Row]:810        ...811 812    @overload813    def cursor(814        self, *, binary: bool = False, row_factory: RowFactory[CursorRow]815    ) -> Cursor[CursorRow]:816        ...817 818    @overload819    def cursor(820        self,821        name: str,822        *,823        binary: bool = False,824        scrollable: Optional[bool] = None,825        withhold: bool = False,826    ) -> ServerCursor[Row]:827        ...828 829    @overload830    def cursor(831        self,832        name: str,833        *,834        binary: bool = False,835        row_factory: RowFactory[CursorRow],836        scrollable: Optional[bool] = None,837        withhold: bool = False,838    ) -> ServerCursor[CursorRow]:839        ...840 841    def cursor(842        self,843        name: str = "",844        *,845        binary: bool = False,846        row_factory: Optional[RowFactory[Any]] = None,847        scrollable: Optional[bool] = None,848        withhold: bool = False,849    ) -> Union[Cursor[Any], ServerCursor[Any]]:850        """851        Return a new cursor to send commands and queries to the connection.852        """853        self._check_connection_ok()854 855        if not row_factory:856            row_factory = self.row_factory857 858        cur: Union[Cursor[Any], ServerCursor[Any]]859        if name:860            cur = self.server_cursor_factory(861                self,862                name=name,863                row_factory=row_factory,864                scrollable=scrollable,865                withhold=withhold,866            )867        else:868            cur = self.cursor_factory(self, row_factory=row_factory)869 870        if binary:871            cur.format = BINARY872 873        return cur874 875    def execute(876        self,877        query: Query,878        params: Optional[Params] = None,879        *,880        prepare: Optional[bool] = None,881        binary: bool = False,882    ) -> Cursor[Row]:883        """Execute a query and return a cursor to read its results."""884        try:885            cur = self.cursor()886            if binary:887                cur.format = BINARY888 889            return cur.execute(query, params, prepare=prepare)890 891        except e._NO_TRACEBACK as ex:892            raise ex.with_traceback(None)893 894    def commit(self) -> None:895        """Commit any pending transaction to the database."""896        with self.lock:897            self.wait(self._commit_gen())898 899    def rollback(self) -> None:900        """Roll back to the start of any pending transaction."""901        with self.lock:902            self.wait(self._rollback_gen())903 904    @contextmanager905    def transaction(906        self,907        savepoint_name: Optional[str] = None,908        force_rollback: bool = False,909    ) -> Iterator[Transaction]:910        """911        Start a context block with a new transaction or nested transaction.912 913        :param savepoint_name: Name of the savepoint used to manage a nested914            transaction. If `!None`, one will be chosen automatically.915        :param force_rollback: Roll back the transaction at the end of the916            block even if there were no error (e.g. to try a no-op process).917        :rtype: Transaction918        """919        tx = Transaction(self, savepoint_name, force_rollback)920        if self._pipeline:921            with self.pipeline(), tx, self.pipeline():922                yield tx923        else:924            with tx:925                yield tx926 927    def notifies(self) -> Generator[Notify, None, None]:928        """929        Yield `Notify` objects as soon as they are received from the database.930        """931        while True:932            with self.lock:933                try:934                    ns = self.wait(notifies(self.pgconn))935                except e._NO_TRACEBACK as ex:936                    raise ex.with_traceback(None)937            enc = pgconn_encoding(self.pgconn)938            for pgn in ns:939                n = Notify(pgn.relname.decode(enc), pgn.extra.decode(enc), pgn.be_pid)940                yield n941 942    @contextmanager943    def pipeline(self) -> Iterator[Pipeline]:944        """Switch the connection into pipeline mode."""945        with self.lock:946            self._check_connection_ok()947 948            pipeline = self._pipeline949            if pipeline is None:950                # WARNING: reference loop, broken ahead.951                pipeline = self._pipeline = Pipeline(self)952 953        try:954            with pipeline:955                yield pipeline956        finally:957            if pipeline.level == 0:958                with self.lock:959                    assert pipeline is self._pipeline960                    self._pipeline = None961 962    def wait(self, gen: PQGen[RV], timeout: Optional[float] = 0.1) -> RV:963        """964        Consume a generator operating on the connection.965 966        The function must be used on generators that don't change connection967        fd (i.e. not on connect and reset).968        """969        try:970            return waiting.wait(gen, self.pgconn.socket, timeout=timeout)971        except KeyboardInterrupt:972            # On Ctrl-C, try to cancel the query in the server, otherwise973            # the connection will remain stuck in ACTIVE state.974            self._try_cancel(self.pgconn)975            try:976                waiting.wait(gen, self.pgconn.socket, timeout=timeout)977            except e.QueryCanceled:978                pass  # as expected979            raise980 981    @classmethod982    def _wait_conn(cls, gen: PQGenConn[RV], timeout: Optional[int]) -> RV:983        """Consume a connection generator."""984        return waiting.wait_conn(gen, timeout=timeout)985 986    def _set_autocommit(self, value: bool) -> None:987        with self.lock:988            self.wait(self._set_autocommit_gen(value))989 990    def _set_isolation_level(self, value: Optional[IsolationLevel]) -> None:991        with self.lock:992            self.wait(self._set_isolation_level_gen(value))993 994    def _set_read_only(self, value: Optional[bool]) -> None:995        with self.lock:996            self.wait(self._set_read_only_gen(value))997 998    def _set_deferrable(self, value: Optional[bool]) -> None:999        with self.lock:1000            self.wait(self._set_deferrable_gen(value))1001 1002    def tpc_begin(self, xid: Union[Xid, str]) -> None:1003        """1004        Begin a TPC transaction with the given transaction ID `!xid`.1005        """1006        with self.lock:1007            self.wait(self._tpc_begin_gen(xid))1008 1009    def tpc_prepare(self) -> None:1010        """1011        Perform the first phase of a transaction started with `tpc_begin()`.1012        """1013        try:1014            with self.lock:1015                self.wait(self._tpc_prepare_gen())1016        except e.ObjectNotInPrerequisiteState as ex:1017            raise e.NotSupportedError(str(ex)) from None1018 1019    def tpc_commit(self, xid: Union[Xid, str, None] = None) -> None:1020        """1021        Commit a prepared two-phase transaction.1022        """1023        with self.lock:1024            self.wait(self._tpc_finish_gen("COMMIT", xid))1025 1026    def tpc_rollback(self, xid: Union[Xid, str, None] = None) -> None:1027        """1028        Roll back a prepared two-phase transaction.1029        """1030        with self.lock:1031            self.wait(self._tpc_finish_gen("ROLLBACK", xid))1032 1033    def tpc_recover(self) -> List[Xid]:1034        self._check_tpc()1035        status = self.info.transaction_status1036        with self.cursor(row_factory=args_row(Xid._from_record)) as cur:1037            cur.execute(Xid._get_recover_query())1038            res = cur.fetchall()1039 1040        if status == IDLE and self.info.transaction_status == INTRANS:1041            self.rollback()1042 1043        return res1044