codekingpro/portable-devtools

1"""2psycopg server-side cursor objects.3"""4 5# Copyright (C) 2020 The Psycopg Team6 7from typing import Any, AsyncIterator, List, Iterable, Iterator8from typing import Optional, TypeVar, TYPE_CHECKING, overload9from warnings import warn10 11from . import pq12from . import sql13from . import errors as e14from .abc import ConnectionType, Query, Params, PQGen15from .rows import Row, RowFactory, AsyncRowFactory16from .cursor import BaseCursor, Cursor17from .generators import execute18from .cursor_async import AsyncCursor19 20if TYPE_CHECKING:21    from .connection import Connection22    from .connection_async import AsyncConnection23 24DEFAULT_ITERSIZE = 10025 26TEXT = pq.Format.TEXT27BINARY = pq.Format.BINARY28 29COMMAND_OK = pq.ExecStatus.COMMAND_OK30TUPLES_OK = pq.ExecStatus.TUPLES_OK31 32IDLE = pq.TransactionStatus.IDLE33INTRANS = pq.TransactionStatus.INTRANS34 35 36class ServerCursorMixin(BaseCursor[ConnectionType, Row]):37    """Mixin to add ServerCursor behaviour and implementation a BaseCursor."""38 39    __slots__ = "_name _scrollable _withhold _described itersize _format".split()40 41    def __init__(42        self,43        name: str,44        scrollable: Optional[bool],45        withhold: bool,46    ):47        self._name = name48        self._scrollable = scrollable49        self._withhold = withhold50        self._described = False51        self.itersize: int = DEFAULT_ITERSIZE52        self._format = TEXT53 54    def __repr__(self) -> str:55        # Insert the name as the second word56        parts = super().__repr__().split(None, 1)57        parts.insert(1, f"{self._name!r}")58        return " ".join(parts)59 60    @property61    def name(self) -> str:62        """The name of the cursor."""63        return self._name64 65    @property66    def scrollable(self) -> Optional[bool]:67        """68        Whether the cursor is scrollable or not.69 70        If `!None` leave the choice to the server. Use `!True` if you want to71        use `scroll()` on the cursor.72        """73        return self._scrollable74 75    @property76    def withhold(self) -> bool:77        """78        If the cursor can be used after the creating transaction has committed.79        """80        return self._withhold81 82    @property83    def rownumber(self) -> Optional[int]:84        """Index of the next row to fetch in the current result.85 86        `!None` if there is no result to fetch.87        """88        res = self.pgresult89        # command_status is empty if the result comes from90        # describe_portal, which means that we have just executed the DECLARE,91        # so we can assume we are at the first row.92        tuples = res and (res.status == TUPLES_OK or res.command_status == b"")93        return self._pos if tuples else None94 95    def _declare_gen(96        self,97        query: Query,98        params: Optional[Params] = None,99        binary: Optional[bool] = None,100    ) -> PQGen[None]:101        """Generator implementing `ServerCursor.execute()`."""102 103        query = self._make_declare_statement(query)104 105        # If the cursor is being reused, the previous one must be closed.106        if self._described:107            yield from self._close_gen()108            self._described = False109 110        yield from self._start_query(query)111        pgq = self._convert_query(query, params)112        self._execute_send(pgq, force_extended=True)113        results = yield from execute(self._conn.pgconn)114        if results[-1].status != COMMAND_OK:115            self._raise_for_result(results[-1])116 117        # Set the format, which will be used by describe and fetch operations118        if binary is None:119            self._format = self.format120        else:121            self._format = BINARY if binary else TEXT122 123        # The above result only returned COMMAND_OK. Get the cursor shape124        yield from self._describe_gen()125 126    def _describe_gen(self) -> PQGen[None]:127        self._pgconn.send_describe_portal(self._name.encode(self._encoding))128        results = yield from execute(self._pgconn)129        self._check_results(results)130        self._results = results131        self._select_current_result(0, format=self._format)132        self._described = True133 134    def _close_gen(self) -> PQGen[None]:135        ts = self._conn.pgconn.transaction_status136 137        # if the connection is not in a sane state, don't even try138        if ts != IDLE and ts != INTRANS:139            return140 141        # If we are IDLE, a WITHOUT HOLD cursor will surely have gone already.142        if not self._withhold and ts == IDLE:143            return144 145        # if we didn't declare the cursor ourselves we still have to close it146        # but we must make sure it exists.147        if not self._described:148            query = sql.SQL(149                "SELECT 1 FROM pg_catalog.pg_cursors WHERE name = {}"150            ).format(sql.Literal(self._name))151            res = yield from self._conn._exec_command(query)152            # pipeline mode otherwise, unsupported here.153            assert res is not None154            if res.ntuples == 0:155                return156 157        query = sql.SQL("CLOSE {}").format(sql.Identifier(self._name))158        yield from self._conn._exec_command(query)159 160    def _fetch_gen(self, num: Optional[int]) -> PQGen[List[Row]]:161        if self.closed:162            raise e.InterfaceError("the cursor is closed")163        # If we are stealing the cursor, make sure we know its shape164        if not self._described:165            yield from self._start_query()166            yield from self._describe_gen()167 168        query = sql.SQL("FETCH FORWARD {} FROM {}").format(169            sql.SQL("ALL") if num is None else sql.Literal(num),170            sql.Identifier(self._name),171        )172        res = yield from self._conn._exec_command(query, result_format=self._format)173        # pipeline mode otherwise, unsupported here.174        assert res is not None175 176        self.pgresult = res177        self._tx.set_pgresult(res, set_loaders=False)178        return self._tx.load_rows(0, res.ntuples, self._make_row)179 180    def _scroll_gen(self, value: int, mode: str) -> PQGen[None]:181        if mode not in ("relative", "absolute"):182            raise ValueError(f"bad mode: {mode}. It should be 'relative' or 'absolute'")183        query = sql.SQL("MOVE{} {} FROM {}").format(184            sql.SQL(" ABSOLUTE" if mode == "absolute" else ""),185            sql.Literal(value),186            sql.Identifier(self._name),187        )188        yield from self._conn._exec_command(query)189 190    def _make_declare_statement(self, query: Query) -> sql.Composed:191        if isinstance(query, bytes):192            query = query.decode(self._encoding)193        if not isinstance(query, sql.Composable):194            query = sql.SQL(query)195 196        parts = [197            sql.SQL("DECLARE"),198            sql.Identifier(self._name),199        ]200        if self._scrollable is not None:201            parts.append(sql.SQL("SCROLL" if self._scrollable else "NO SCROLL"))202        parts.append(sql.SQL("CURSOR"))203        if self._withhold:204            parts.append(sql.SQL("WITH HOLD"))205        parts.append(sql.SQL("FOR"))206        parts.append(query)207 208        return sql.SQL(" ").join(parts)209 210 211class ServerCursor(ServerCursorMixin["Connection[Any]", Row], Cursor[Row]):212    __module__ = "psycopg"213    __slots__ = ()214    _Self = TypeVar("_Self", bound="ServerCursor[Any]")215 216    @overload217    def __init__(218        self: "ServerCursor[Row]",219        connection: "Connection[Row]",220        name: str,221        *,222        scrollable: Optional[bool] = None,223        withhold: bool = False,224    ):225        ...226 227    @overload228    def __init__(229        self: "ServerCursor[Row]",230        connection: "Connection[Any]",231        name: str,232        *,233        row_factory: RowFactory[Row],234        scrollable: Optional[bool] = None,235        withhold: bool = False,236    ):237        ...238 239    def __init__(240        self,241        connection: "Connection[Any]",242        name: str,243        *,244        row_factory: Optional[RowFactory[Row]] = None,245        scrollable: Optional[bool] = None,246        withhold: bool = False,247    ):248        Cursor.__init__(249            self, connection, row_factory=row_factory or connection.row_factory250        )251        ServerCursorMixin.__init__(self, name, scrollable, withhold)252 253    def __del__(self) -> None:254        if not self.closed:255            warn(256                f"the server-side cursor {self} was deleted while still open."257                " Please use 'with' or '.close()' to close the cursor properly",258                ResourceWarning,259            )260 261    def close(self) -> None:262        """263        Close the current cursor and free associated resources.264        """265        with self._conn.lock:266            if self.closed:267                return268            if not self._conn.closed:269                self._conn.wait(self._close_gen())270            super().close()271 272    def execute(273        self: _Self,274        query: Query,275        params: Optional[Params] = None,276        *,277        binary: Optional[bool] = None,278        **kwargs: Any,279    ) -> _Self:280        """281        Open a cursor to execute a query to the database.282        """283        if kwargs:284            raise TypeError(f"keyword not supported: {list(kwargs)[0]}")285        if self._pgconn.pipeline_status:286            raise e.NotSupportedError(287                "server-side cursors not supported in pipeline mode"288            )289 290        try:291            with self._conn.lock:292                self._conn.wait(self._declare_gen(query, params, binary))293        except e._NO_TRACEBACK as ex:294            raise ex.with_traceback(None)295 296        return self297 298    def executemany(299        self,300        query: Query,301        params_seq: Iterable[Params],302        *,303        returning: bool = True,304    ) -> None:305        """Method not implemented for server-side cursors."""306        raise e.NotSupportedError("executemany not supported on server-side cursors")307 308    def fetchone(self) -> Optional[Row]:309        with self._conn.lock:310            recs = self._conn.wait(self._fetch_gen(1))311        if recs:312            self._pos += 1313            return recs[0]314        else:315            return None316 317    def fetchmany(self, size: int = 0) -> List[Row]:318        if not size:319            size = self.arraysize320        with self._conn.lock:321            recs = self._conn.wait(self._fetch_gen(size))322        self._pos += len(recs)323        return recs324 325    def fetchall(self) -> List[Row]:326        with self._conn.lock:327            recs = self._conn.wait(self._fetch_gen(None))328        self._pos += len(recs)329        return recs330 331    def __iter__(self) -> Iterator[Row]:332        while True:333            with self._conn.lock:334                recs = self._conn.wait(self._fetch_gen(self.itersize))335            for rec in recs:336                self._pos += 1337                yield rec338            if len(recs) < self.itersize:339                break340 341    def scroll(self, value: int, mode: str = "relative") -> None:342        with self._conn.lock:343            self._conn.wait(self._scroll_gen(value, mode))344        # Postgres doesn't have a reliable way to report a cursor out of bound345        if mode == "relative":346            self._pos += value347        else:348            self._pos = value349 350 351class AsyncServerCursor(352    ServerCursorMixin["AsyncConnection[Any]", Row], AsyncCursor[Row]353):354    __module__ = "psycopg"355    __slots__ = ()356    _Self = TypeVar("_Self", bound="AsyncServerCursor[Any]")357 358    @overload359    def __init__(360        self: "AsyncServerCursor[Row]",361        connection: "AsyncConnection[Row]",362        name: str,363        *,364        scrollable: Optional[bool] = None,365        withhold: bool = False,366    ):367        ...368 369    @overload370    def __init__(371        self: "AsyncServerCursor[Row]",372        connection: "AsyncConnection[Any]",373        name: str,374        *,375        row_factory: AsyncRowFactory[Row],376        scrollable: Optional[bool] = None,377        withhold: bool = False,378    ):379        ...380 381    def __init__(382        self,383        connection: "AsyncConnection[Any]",384        name: str,385        *,386        row_factory: Optional[AsyncRowFactory[Row]] = None,387        scrollable: Optional[bool] = None,388        withhold: bool = False,389    ):390        AsyncCursor.__init__(391            self, connection, row_factory=row_factory or connection.row_factory392        )393        ServerCursorMixin.__init__(self, name, scrollable, withhold)394 395    def __del__(self) -> None:396        if not self.closed:397            warn(398                f"the server-side cursor {self} was deleted while still open."399                " Please use 'with' or '.close()' to close the cursor properly",400                ResourceWarning,401            )402 403    async def close(self) -> None:404        async with self._conn.lock:405            if self.closed:406                return407            if not self._conn.closed:408                await self._conn.wait(self._close_gen())409            await super().close()410 411    async def execute(412        self: _Self,413        query: Query,414        params: Optional[Params] = None,415        *,416        binary: Optional[bool] = None,417        **kwargs: Any,418    ) -> _Self:419        if kwargs:420            raise TypeError(f"keyword not supported: {list(kwargs)[0]}")421        if self._pgconn.pipeline_status:422            raise e.NotSupportedError(423                "server-side cursors not supported in pipeline mode"424            )425 426        try:427            async with self._conn.lock:428                await self._conn.wait(self._declare_gen(query, params, binary))429        except e._NO_TRACEBACK as ex:430            raise ex.with_traceback(None)431 432        return self433 434    async def executemany(435        self,436        query: Query,437        params_seq: Iterable[Params],438        *,439        returning: bool = True,440    ) -> None:441        raise e.NotSupportedError("executemany not supported on server-side cursors")442 443    async def fetchone(self) -> Optional[Row]:444        async with self._conn.lock:445            recs = await self._conn.wait(self._fetch_gen(1))446        if recs:447            self._pos += 1448            return recs[0]449        else:450            return None451 452    async def fetchmany(self, size: int = 0) -> List[Row]:453        if not size:454            size = self.arraysize455        async with self._conn.lock:456            recs = await self._conn.wait(self._fetch_gen(size))457        self._pos += len(recs)458        return recs459 460    async def fetchall(self) -> List[Row]:461        async with self._conn.lock:462            recs = await self._conn.wait(self._fetch_gen(None))463        self._pos += len(recs)464        return recs465 466    async def __aiter__(self) -> AsyncIterator[Row]:467        while True:468            async with self._conn.lock:469                recs = await self._conn.wait(self._fetch_gen(self.itersize))470            for rec in recs:471                self._pos += 1472                yield rec473            if len(recs) < self.itersize:474                break475 476    async def scroll(self, value: int, mode: str = "relative") -> None:477        async with self._conn.lock:478            await self._conn.wait(self._scroll_gen(value, mode))479