codekingpro/portable-devtools

1"""2commands pipeline management3"""4 5# Copyright (C) 2021 The Psycopg Team6 7import logging8from types import TracebackType9from typing import Any, List, Optional, Union, Tuple, Type, TypeVar, TYPE_CHECKING10from typing_extensions import TypeAlias11 12from . import pq13from . import errors as e14from .abc import PipelineCommand, PQGen15from ._compat import Deque16from .pq.misc import connection_summary17from ._encodings import pgconn_encoding18from ._preparing import Key, Prepare19from .generators import pipeline_communicate, fetch_many, send20 21if TYPE_CHECKING:22    from .pq.abc import PGresult23    from .cursor import BaseCursor24    from .connection import BaseConnection, Connection25    from .connection_async import AsyncConnection26 27 28PendingResult: TypeAlias = Union[29    None, Tuple["BaseCursor[Any, Any]", Optional[Tuple[Key, Prepare, bytes]]]30]31 32FATAL_ERROR = pq.ExecStatus.FATAL_ERROR33PIPELINE_ABORTED = pq.ExecStatus.PIPELINE_ABORTED34BAD = pq.ConnStatus.BAD35 36ACTIVE = pq.TransactionStatus.ACTIVE37 38logger = logging.getLogger("psycopg")39 40 41class BasePipeline:42    command_queue: Deque[PipelineCommand]43    result_queue: Deque[PendingResult]44    _is_supported: Optional[bool] = None45 46    def __init__(self, conn: "BaseConnection[Any]") -> None:47        self._conn = conn48        self.pgconn = conn.pgconn49        self.command_queue = Deque[PipelineCommand]()50        self.result_queue = Deque[PendingResult]()51        self.level = 052 53    def __repr__(self) -> str:54        cls = f"{self.__class__.__module__}.{self.__class__.__qualname__}"55        info = connection_summary(self._conn.pgconn)56        return f"<{cls} {info} at 0x{id(self):x}>"57 58    @property59    def status(self) -> pq.PipelineStatus:60        return pq.PipelineStatus(self.pgconn.pipeline_status)61 62    @classmethod63    def is_supported(cls) -> bool:64        """Return `!True` if the psycopg libpq wrapper supports pipeline mode."""65        if BasePipeline._is_supported is None:66            BasePipeline._is_supported = not cls._not_supported_reason()67        return BasePipeline._is_supported68 69    @classmethod70    def _not_supported_reason(cls) -> str:71        """Return the reason why the pipeline mode is not supported.72 73        Return an empty string if pipeline mode is supported.74        """75        # Support only depends on the libpq functions available in the pq76        # wrapper, not on the database version.77        if pq.version() < 140000:78            return (79                f"libpq too old {pq.version()};"80                " v14 or greater required for pipeline mode"81            )82 83        if pq.__build_version__ < 140000:84            return (85                f"libpq too old: module built for {pq.__build_version__};"86                " v14 or greater required for pipeline mode"87            )88 89        return ""90 91    def _enter_gen(self) -> PQGen[None]:92        if not self.is_supported():93            raise e.NotSupportedError(94                f"pipeline mode not supported: {self._not_supported_reason()}"95            )96        if self.level == 0:97            self.pgconn.enter_pipeline_mode()98        elif self.command_queue or self.pgconn.transaction_status == ACTIVE:99            # Nested pipeline case.100            #  Transaction might be ACTIVE when the pipeline uses an "implicit101            #  transaction", typically in autocommit mode. But when entering a102            #  Psycopg transaction(), we expect the IDLE state. By sync()-ing,103            #  we make sure all previous commands are completed and the104            #  transaction gets back to IDLE.105            yield from self._sync_gen()106        self.level += 1107 108    def _exit(self, exc: Optional[BaseException]) -> None:109        self.level -= 1110        if self.level == 0 and self.pgconn.status != BAD:111            try:112                self.pgconn.exit_pipeline_mode()113            except e.OperationalError as exc2:114                # Notice that this error might be pretty irrecoverable. It115                # happens on COPY, for instance: even if sync succeeds, exiting116                # fails with "cannot exit pipeline mode with uncollected results"117                if exc:118                    logger.warning("error ignored exiting %r: %s", self, exc2)119                else:120                    raise exc2.with_traceback(None)121 122    def _sync_gen(self) -> PQGen[None]:123        self._enqueue_sync()124        yield from self._communicate_gen()125        yield from self._fetch_gen(flush=False)126 127    def _exit_gen(self) -> PQGen[None]:128        """129        Exit current pipeline by sending a Sync and fetch back all remaining results.130        """131        try:132            self._enqueue_sync()133            yield from self._communicate_gen()134        finally:135            # No need to force flush since we emitted a sync just before.136            yield from self._fetch_gen(flush=False)137 138    def _communicate_gen(self) -> PQGen[None]:139        """Communicate with pipeline to send commands and possibly fetch140        results, which are then processed.141        """142        fetched = yield from pipeline_communicate(self.pgconn, self.command_queue)143        exception = None144        for results in fetched:145            queued = self.result_queue.popleft()146            try:147                self._process_results(queued, results)148            except e.Error as exc:149                if exception is None:150                    exception = exc151        if exception is not None:152            raise exception153 154    def _fetch_gen(self, *, flush: bool) -> PQGen[None]:155        """Fetch available results from the connection and process them with156        pipeline queued items.157 158        If 'flush' is True, a PQsendFlushRequest() is issued in order to make159        sure results can be fetched. Otherwise, the caller may emit a160        PQpipelineSync() call to ensure the output buffer gets flushed before161        fetching.162        """163        if not self.result_queue:164            return165 166        if flush:167            self.pgconn.send_flush_request()168            yield from send(self.pgconn)169 170        exception = None171        while self.result_queue:172            results = yield from fetch_many(self.pgconn)173            if not results:174                # No more results to fetch, but there may still be pending175                # commands.176                break177            queued = self.result_queue.popleft()178            try:179                self._process_results(queued, results)180            except e.Error as exc:181                if exception is None:182                    exception = exc183        if exception is not None:184            raise exception185 186    def _process_results(187        self, queued: PendingResult, results: List["PGresult"]188    ) -> None:189        """Process a results set fetched from the current pipeline.190 191        This matches 'results' with its respective element in the pipeline192        queue. For commands (None value in the pipeline queue), results are193        checked directly. For prepare statement creation requests, update the194        cache. Otherwise, results are attached to their respective cursor.195        """196        if queued is None:197            (result,) = results198            if result.status == FATAL_ERROR:199                raise e.error_from_result(result, encoding=pgconn_encoding(self.pgconn))200            elif result.status == PIPELINE_ABORTED:201                raise e.PipelineAborted("pipeline aborted")202        else:203            cursor, prepinfo = queued204            if prepinfo:205                key, prep, name = prepinfo206                # Update the prepare state of the query.207                cursor._conn._prepared.validate(key, prep, name, results)208            cursor._set_results_from_pipeline(results)209 210    def _enqueue_sync(self) -> None:211        """Enqueue a PQpipelineSync() command."""212        self.command_queue.append(self.pgconn.pipeline_sync)213        self.result_queue.append(None)214 215 216class Pipeline(BasePipeline):217    """Handler for connection in pipeline mode."""218 219    __module__ = "psycopg"220    _conn: "Connection[Any]"221    _Self = TypeVar("_Self", bound="Pipeline")222 223    def __init__(self, conn: "Connection[Any]") -> None:224        super().__init__(conn)225 226    def sync(self) -> None:227        """Sync the pipeline, send any pending command and receive and process228        all available results.229        """230        try:231            with self._conn.lock:232                self._conn.wait(self._sync_gen())233        except e._NO_TRACEBACK as ex:234            raise ex.with_traceback(None)235 236    def __enter__(self: _Self) -> _Self:237        with self._conn.lock:238            self._conn.wait(self._enter_gen())239        return self240 241    def __exit__(242        self,243        exc_type: Optional[Type[BaseException]],244        exc_val: Optional[BaseException],245        exc_tb: Optional[TracebackType],246    ) -> None:247        try:248            with self._conn.lock:249                self._conn.wait(self._exit_gen())250        except Exception as exc2:251            # Don't clobber an exception raised in the block with this one252            if exc_val:253                logger.warning("error ignored terminating %r: %s", self, exc2)254            else:255                raise exc2.with_traceback(None)256        finally:257            self._exit(exc_val)258 259 260class AsyncPipeline(BasePipeline):261    """Handler for async connection in pipeline mode."""262 263    __module__ = "psycopg"264    _conn: "AsyncConnection[Any]"265    _Self = TypeVar("_Self", bound="AsyncPipeline")266 267    def __init__(self, conn: "AsyncConnection[Any]") -> None:268        super().__init__(conn)269 270    async def sync(self) -> None:271        try:272            async with self._conn.lock:273                await self._conn.wait(self._sync_gen())274        except e._NO_TRACEBACK as ex:275            raise ex.with_traceback(None)276 277    async def __aenter__(self: _Self) -> _Self:278        async with self._conn.lock:279            await self._conn.wait(self._enter_gen())280        return self281 282    async def __aexit__(283        self,284        exc_type: Optional[Type[BaseException]],285        exc_val: Optional[BaseException],286        exc_tb: Optional[TracebackType],287    ) -> None:288        try:289            async with self._conn.lock:290                await self._conn.wait(self._exit_gen())291        except Exception as exc2:292            # Don't clobber an exception raised in the block with this one293            if exc_val:294                logger.warning("error ignored terminating %r: %s", self, exc2)295            else:296                raise exc2.with_traceback(None)297        finally:298            self._exit(exc_val)299