"""
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
