codekingpro/portable-devtools

1"""2CockroachDB-specific connections.3"""4 5# Copyright (C) 2022 The Psycopg Team6 7import re8from typing import Any, Optional, Type, Union, overload, TYPE_CHECKING9 10from .. import errors as e11from ..abc import AdaptContext12from ..rows import Row, RowFactory, AsyncRowFactory, TupleRow13from ..conninfo import ConnectionInfo14from ..connection import Connection15from .._adapters_map import AdaptersMap16from ..connection_async import AsyncConnection17from ._types import adapters18 19if TYPE_CHECKING:20    from ..pq.abc import PGconn21    from ..cursor import Cursor22    from ..cursor_async import AsyncCursor23 24 25class _CrdbConnectionMixin:26    _adapters: Optional[AdaptersMap]27    pgconn: "PGconn"28 29    @classmethod30    def is_crdb(31        cls, conn: Union[Connection[Any], AsyncConnection[Any], "PGconn"]32    ) -> bool:33        """34        Return `!True` if the server connected to `!conn` is CockroachDB.35        """36        if isinstance(conn, (Connection, AsyncConnection)):37            conn = conn.pgconn38 39        return bool(conn.parameter_status(b"crdb_version"))40 41    @property42    def adapters(self) -> AdaptersMap:43        if not self._adapters:44            # By default, use CockroachDB adapters map45            self._adapters = AdaptersMap(adapters)46 47        return self._adapters48 49    @property50    def info(self) -> "CrdbConnectionInfo":51        return CrdbConnectionInfo(self.pgconn)52 53    def _check_tpc(self) -> None:54        if self.is_crdb(self.pgconn):55            raise e.NotSupportedError("CockroachDB doesn't support prepared statements")56 57 58class CrdbConnection(_CrdbConnectionMixin, Connection[Row]):59    """60    Wrapper for a connection to a CockroachDB database.61    """62 63    __module__ = "psycopg.crdb"64 65    # TODO: this method shouldn't require re-definition if the base class66    # implements a generic self.67    # https://github.com/psycopg/psycopg/issues/30868    @overload69    @classmethod70    def connect(71        cls,72        conninfo: str = "",73        *,74        autocommit: bool = False,75        row_factory: RowFactory[Row],76        prepare_threshold: Optional[int] = 5,77        cursor_factory: "Optional[Type[Cursor[Row]]]" = None,78        context: Optional[AdaptContext] = None,79        **kwargs: Union[None, int, str],80    ) -> "CrdbConnection[Row]":81        ...82 83    @overload84    @classmethod85    def connect(86        cls,87        conninfo: str = "",88        *,89        autocommit: bool = False,90        prepare_threshold: Optional[int] = 5,91        cursor_factory: "Optional[Type[Cursor[Any]]]" = None,92        context: Optional[AdaptContext] = None,93        **kwargs: Union[None, int, str],94    ) -> "CrdbConnection[TupleRow]":95        ...96 97    @classmethod98    def connect(cls, conninfo: str = "", **kwargs: Any) -> "CrdbConnection[Any]":99        """100        Connect to a database server and return a new `CrdbConnection` instance.101        """102        return super().connect(conninfo, **kwargs)  # type: ignore[return-value]103 104 105class AsyncCrdbConnection(_CrdbConnectionMixin, AsyncConnection[Row]):106    """107    Wrapper for an async connection to a CockroachDB database.108    """109 110    __module__ = "psycopg.crdb"111 112    # TODO: this method shouldn't require re-definition if the base class113    # implements a generic self.114    # https://github.com/psycopg/psycopg/issues/308115    @overload116    @classmethod117    async def connect(118        cls,119        conninfo: str = "",120        *,121        autocommit: bool = False,122        prepare_threshold: Optional[int] = 5,123        row_factory: AsyncRowFactory[Row],124        cursor_factory: "Optional[Type[AsyncCursor[Row]]]" = None,125        context: Optional[AdaptContext] = None,126        **kwargs: Union[None, int, str],127    ) -> "AsyncCrdbConnection[Row]":128        ...129 130    @overload131    @classmethod132    async def connect(133        cls,134        conninfo: str = "",135        *,136        autocommit: bool = False,137        prepare_threshold: Optional[int] = 5,138        cursor_factory: "Optional[Type[AsyncCursor[Any]]]" = None,139        context: Optional[AdaptContext] = None,140        **kwargs: Union[None, int, str],141    ) -> "AsyncCrdbConnection[TupleRow]":142        ...143 144    @classmethod145    async def connect(146        cls, conninfo: str = "", **kwargs: Any147    ) -> "AsyncCrdbConnection[Any]":148        return await super().connect(conninfo, **kwargs)  # type: ignore [no-any-return]149 150 151class CrdbConnectionInfo(ConnectionInfo):152    """153    `~psycopg.ConnectionInfo` subclass to get info about a CockroachDB database.154    """155 156    __module__ = "psycopg.crdb"157 158    @property159    def vendor(self) -> str:160        return "CockroachDB"161 162    @property163    def server_version(self) -> int:164        """165        Return the CockroachDB server version connected.166 167        Return a number in the PostgreSQL format (e.g. 21.2.10 -> 210210).168        """169        sver = self.parameter_status("crdb_version")170        if not sver:171            raise e.InternalError("'crdb_version' parameter status not set")172 173        ver = self.parse_crdb_version(sver)174        if ver is None:175            raise e.InterfaceError(f"couldn't parse CockroachDB version from: {sver!r}")176 177        return ver178 179    @classmethod180    def parse_crdb_version(self, sver: str) -> Optional[int]:181        m = re.search(r"\bv(\d+)\.(\d+)\.(\d+)", sver)182        if not m:183            return None184 185        return int(m.group(1)) * 10000 + int(m.group(2)) * 100 + int(m.group(3))186