codekingpro/portable-devtools

1"""2psycopg two-phase commit support3"""4 5# Copyright (C) 2021 The Psycopg Team6 7import re8import datetime as dt9from base64 import b64encode, b64decode10from typing import Optional, Union11from dataclasses import dataclass, replace12 13_re_xid = re.compile(r"^(\d+)_([^_]*)_([^_]*)$")14 15 16@dataclass(frozen=True)17class Xid:18    """A two-phase commit transaction identifier.19 20    The object can also be unpacked as a 3-item tuple (`format_id`, `gtrid`,21    `bqual`).22 23    """24 25    format_id: Optional[int]26    gtrid: str27    bqual: Optional[str]28    prepared: Optional[dt.datetime] = None29    owner: Optional[str] = None30    database: Optional[str] = None31 32    @classmethod33    def from_string(cls, s: str) -> "Xid":34        """Try to parse an XA triple from the string.35 36        This may fail for several reasons. In such case return an unparsed Xid.37        """38        try:39            return cls._parse_string(s)40        except Exception:41            return Xid(None, s, None)42 43    def __str__(self) -> str:44        return self._as_tid()45 46    def __len__(self) -> int:47        return 348 49    def __getitem__(self, index: int) -> Union[int, str, None]:50        return (self.format_id, self.gtrid, self.bqual)[index]51 52    @classmethod53    def _parse_string(cls, s: str) -> "Xid":54        m = _re_xid.match(s)55        if not m:56            raise ValueError("bad Xid format")57 58        format_id = int(m.group(1))59        gtrid = b64decode(m.group(2)).decode()60        bqual = b64decode(m.group(3)).decode()61        return cls.from_parts(format_id, gtrid, bqual)62 63    @classmethod64    def from_parts(65        cls, format_id: Optional[int], gtrid: str, bqual: Optional[str]66    ) -> "Xid":67        if format_id is not None:68            if bqual is None:69                raise TypeError("if format_id is specified, bqual must be too")70            if not 0 <= format_id < 0x80000000:71                raise ValueError("format_id must be a non-negative 32-bit integer")72            if len(bqual) > 64:73                raise ValueError("bqual must be not longer than 64 chars")74            if len(gtrid) > 64:75                raise ValueError("gtrid must be not longer than 64 chars")76 77        elif bqual is None:78            raise TypeError("if format_id is None, bqual must be None too")79 80        return Xid(format_id, gtrid, bqual)81 82    def _as_tid(self) -> str:83        """84        Return the PostgreSQL transaction_id for this XA xid.85 86        PostgreSQL wants just a string, while the DBAPI supports the XA87        standard and thus a triple. We use the same conversion algorithm88        implemented by JDBC in order to allow some form of interoperation.89 90        see also: the pgjdbc implementation91          http://cvs.pgfoundry.org/cgi-bin/cvsweb.cgi/jdbc/pgjdbc/org/92            postgresql/xa/RecoveredXid.java?rev=1.293        """94        if self.format_id is None or self.bqual is None:95            # Unparsed xid: return the gtrid.96            return self.gtrid97 98        # XA xid: mash together the components.99        egtrid = b64encode(self.gtrid.encode()).decode()100        ebqual = b64encode(self.bqual.encode()).decode()101 102        return f"{self.format_id}_{egtrid}_{ebqual}"103 104    @classmethod105    def _get_recover_query(cls) -> str:106        return "SELECT gid, prepared, owner, database FROM pg_prepared_xacts"107 108    @classmethod109    def _from_record(110        cls, gid: str, prepared: dt.datetime, owner: str, database: str111    ) -> "Xid":112        xid = Xid.from_string(gid)113        return replace(xid, prepared=prepared, owner=owner, database=database)114 115 116Xid.__module__ = "psycopg"117