codekingpro/portable-devtools

1"""2Support for prepared statements3"""4 5# Copyright (C) 2020 The Psycopg Team6 7from enum import IntEnum, auto8from typing import Iterator, Optional, Sequence, Tuple, TYPE_CHECKING9from collections import OrderedDict10from typing_extensions import TypeAlias11 12from . import pq13from ._compat import Deque14from ._queries import PostgresQuery15 16if TYPE_CHECKING:17    from .pq.abc import PGresult18 19Key: TypeAlias = Tuple[bytes, Tuple[int, ...]]20 21COMMAND_OK = pq.ExecStatus.COMMAND_OK22TUPLES_OK = pq.ExecStatus.TUPLES_OK23 24 25class Prepare(IntEnum):26    NO = auto()27    YES = auto()28    SHOULD = auto()29 30 31class PrepareManager:32    # Number of times a query is executed before it is prepared.33    prepare_threshold: Optional[int] = 534 35    # Maximum number of prepared statements on the connection.36    prepared_max: int = 10037 38    def __init__(self) -> None:39        # Map (query, types) to the number of times the query was seen.40        self._counts: OrderedDict[Key, int] = OrderedDict()41 42        # Map (query, types) to the name of the statement if  prepared.43        self._names: OrderedDict[Key, bytes] = OrderedDict()44 45        # Counter to generate prepared statements names46        self._prepared_idx = 047 48        self._maint_commands = Deque[bytes]()49 50    @staticmethod51    def key(query: PostgresQuery) -> Key:52        return (query.query, query.types)53 54    def get(55        self, query: PostgresQuery, prepare: Optional[bool] = None56    ) -> Tuple[Prepare, bytes]:57        """58        Check if a query is prepared, tell back whether to prepare it.59        """60        if prepare is False or self.prepare_threshold is None:61            # The user doesn't want this query to be prepared62            return Prepare.NO, b""63 64        key = self.key(query)65        name = self._names.get(key)66        if name:67            # The query was already prepared in this session68            return Prepare.YES, name69 70        count = self._counts.get(key, 0)71        if count >= self.prepare_threshold or prepare:72            # The query has been executed enough times and needs to be prepared73            name = f"_pg3_{self._prepared_idx}".encode()74            self._prepared_idx += 175            return Prepare.SHOULD, name76        else:77            # The query is not to be prepared yet78            return Prepare.NO, b""79 80    def _should_discard(self, prep: Prepare, results: Sequence["PGresult"]) -> bool:81        """Check if we need to discard our entire state: it should happen on82        rollback or on dropping objects, because the same object may get83        recreated and postgres would fail internal lookups.84        """85        if self._names or prep == Prepare.SHOULD:86            for result in results:87                if result.status != COMMAND_OK:88                    continue89                cmdstat = result.command_status90                if cmdstat and (cmdstat.startswith(b"DROP ") or cmdstat == b"ROLLBACK"):91                    return self.clear()92        return False93 94    @staticmethod95    def _check_results(results: Sequence["PGresult"]) -> bool:96        """Return False if 'results' are invalid for prepared statement cache."""97        if len(results) != 1:98            # We cannot prepare a multiple statement99            return False100 101        status = results[0].status102        if COMMAND_OK != status != TUPLES_OK:103            # We don't prepare failed queries or other weird results104            return False105 106        return True107 108    def _rotate(self) -> None:109        """Evict an old value from the cache.110 111        If it was prepared, deallocate it. Do it only once: if the cache was112        resized, deallocate gradually.113        """114        if len(self._counts) > self.prepared_max:115            self._counts.popitem(last=False)116 117        if len(self._names) > self.prepared_max:118            name = self._names.popitem(last=False)[1]119            self._maint_commands.append(b"DEALLOCATE " + name)120 121    def maybe_add_to_cache(122        self, query: PostgresQuery, prep: Prepare, name: bytes123    ) -> Optional[Key]:124        """Handle 'query' for possible addition to the cache.125 126        If a new entry has been added, return its key. Return None otherwise127        (meaning the query is already in cache or cache is not enabled).128        """129        # don't do anything if prepared statements are disabled130        if self.prepare_threshold is None:131            return None132 133        key = self.key(query)134        if key in self._counts:135            if prep is Prepare.SHOULD:136                del self._counts[key]137                self._names[key] = name138            else:139                self._counts[key] += 1140                self._counts.move_to_end(key)141            return None142 143        elif key in self._names:144            self._names.move_to_end(key)145            return None146 147        else:148            if prep is Prepare.SHOULD:149                self._names[key] = name150            else:151                self._counts[key] = 1152            return key153 154    def validate(155        self,156        key: Key,157        prep: Prepare,158        name: bytes,159        results: Sequence["PGresult"],160    ) -> None:161        """Validate cached entry with 'key' by checking query 'results'.162 163        Possibly record a command to perform maintenance on database side.164        """165        if self._should_discard(prep, results):166            return167 168        if not self._check_results(results):169            self._names.pop(key, None)170            self._counts.pop(key, None)171        else:172            self._rotate()173 174    def clear(self) -> bool:175        """Clear the cache of the maintenance commands.176 177        Clear the internal state and prepare a command to clear the state of178        the server.179        """180        self._counts.clear()181        if self._names:182            self._names.clear()183            self._maint_commands.clear()184            self._maint_commands.append(b"DEALLOCATE ALL")185            return True186        else:187            return False188 189    def get_maintenance_commands(self) -> Iterator[bytes]:190        """191        Iterate over the commands needed to align the server state to our state192        """193        while self._maint_commands:194            yield self._maint_commands.popleft()195