pgtern 0.3.0__py3-none-any.whl

pgtern/errors.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DB-API 2.0 exceptions, plus one subclass per common SQLSTATE.
+
+Server errors are raised as the class registered for their SQLSTATE
+(``errors.UniqueViolation`` for ``23505``, ...); unknown codes fall back to the DB-API
+class for their first two characters (``DatabaseError`` if even that is unknown).
+``str(exc)`` starts with the server's primary message; DETAIL, HINT and CONTEXT follow
+on later lines. The class layout follows the DB-API conventions.
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable
+
+
+class Warning(Exception):  # noqa: A001 - the name is required by DB-API 2.0
+    pass
+
+
+_DIAG_FIELDS = {
+    "severity": "S", "severity_nonlocalized": "V", "sqlstate": "C", "message_primary": "M",
+    "message_detail": "D", "message_hint": "H", "statement_position": "P",
+    "internal_position": "p", "internal_query": "q", "context": "W", "schema_name": "s",
+    "table_name": "t", "column_name": "c", "datatype_name": "d", "constraint_name": "n",
+    "source_file": "F", "source_line": "L", "source_function": "R",
+}
+
+
+class Diagnostic:
+    """The fields of a server ErrorResponse (``diag.message_primary``, ``diag.constraint_name``, ...)."""
+
+    __slots__ = ("_fields",)
+
+    if TYPE_CHECKING:  # the properties below are generated at import time; declared here for type checkers
+        severity: str | None
+        severity_nonlocalized: str | None
+        sqlstate: str | None
+        message_primary: str | None
+        message_detail: str | None
+        message_hint: str | None
+        statement_position: str | None
+        internal_position: str | None
+        internal_query: str | None
+        context: str | None
+        schema_name: str | None
+        table_name: str | None
+        column_name: str | None
+        datatype_name: str | None
+        constraint_name: str | None
+        source_file: str | None
+        source_line: str | None
+        source_function: str | None
+
+    def __init__(self, fields: dict[str, str]) -> None:
+        self._fields = dict(fields)
+
+    def __repr__(self) -> str:
+        return f"Diagnostic(sqlstate={self.sqlstate!r}, message_primary={self.message_primary!r})"
+
+
+def _field_getter(code: str) -> Callable[[Diagnostic], str | None]:
+    def get(self: Diagnostic) -> str | None:
+        return self._fields.get(code)
+    return get
+
+
+for _name, _code in _DIAG_FIELDS.items():
+    setattr(Diagnostic, _name, property(_field_getter(_code)))
+
+
+class Error(Exception):
+    sqlstate: str | None = None
+
+    def __init__(self, *args: Any, diag: Diagnostic | None = None) -> None:
+        super().__init__(*args)
+        self.diag = diag if diag is not None else Diagnostic({})
+
+
+class InterfaceError(Error): ...
+class DatabaseError(Error): ...
+class DataError(DatabaseError): ...
+class OperationalError(DatabaseError): ...
+class IntegrityError(DatabaseError): ...
+class InternalError(DatabaseError): ...
+class ProgrammingError(DatabaseError): ...
+class NotSupportedError(DatabaseError): ...
+
+
+class PipelineAborted(OperationalError):
+    """A statement in a pipeline did not run because an earlier statement in the same batch failed."""
+
+
+# SQLSTATE class (first two characters) -> DB-API base class. Any class not listed here
+# (03, 09, 0B, 0F, 0L, 0P, 72, ...) maps to plain DatabaseError. The table is checked against
+# unknown codes are covered by tests/test_unit.py (test_unknown_sqlstates_fall_back_to_their_dbapi_class).
+_BASE_BY_CLASS: dict[str, type[Error]] = {
+    "08": OperationalError, "0A": NotSupportedError, "20": ProgrammingError, "21": ProgrammingError,
+    "22": DataError, "23": IntegrityError, "24": InternalError, "25": InternalError,
+    "26": ProgrammingError, "27": OperationalError, "28": OperationalError, "2B": InternalError,
+    "2D": InternalError, "2F": OperationalError, "34": ProgrammingError, "38": OperationalError,
+    "39": OperationalError, "3B": OperationalError, "3D": ProgrammingError, "3F": ProgrammingError,
+    "40": OperationalError, "42": ProgrammingError, "44": ProgrammingError, "53": OperationalError,
+    "54": OperationalError, "55": OperationalError, "57": OperationalError, "58": OperationalError,
+    "F0": OperationalError, "HV": OperationalError, "P0": ProgrammingError, "XX": InternalError,
+}
+
+# "<sqlstate> <ClassName>". Codes ending in 000 are class-level parents of the rest of their class.
+_SPEC = """
+08000 ConnectionException
+08001 SqlclientUnableToEstablishSqlconnection
+08003 ConnectionDoesNotExist
+08004 SqlserverRejectedEstablishmentOfSqlconnection
+08006 ConnectionFailure
+08P01 ProtocolViolation
+0A000 FeatureNotSupported
+21000 CardinalityViolation
+22000 DataException
+22001 StringDataRightTruncation
+22002 NullValueNoIndicatorParameter
+22003 NumericValueOutOfRange
+22007 InvalidDatetimeFormat
+22008 DatetimeFieldOverflow
+22012 DivisionByZero
+22021 CharacterNotInRepertoire
+22023 InvalidParameterValue
+22026 StringDataLengthMismatch
+22P02 InvalidTextRepresentation
+22P03 InvalidBinaryRepresentation
+22P05 UntranslatableCharacter
+23000 IntegrityConstraintViolation
+23001 RestrictViolation
+23502 NotNullViolation
+23503 ForeignKeyViolation
+23505 UniqueViolation
+23514 CheckViolation
+23P01 ExclusionViolation
+24000 InvalidCursorState
+25000 InvalidTransactionState
+25001 ActiveSqlTransaction
+25006 ReadOnlySqlTransaction
+25P01 NoActiveSqlTransaction
+25P02 InFailedSqlTransaction
+25P03 IdleInTransactionSessionTimeout
+26000 InvalidSqlStatementName
+28000 InvalidAuthorizationSpecification
+28P01 InvalidPassword
+2BP01 DependentObjectsStillExist
+2D000 InvalidTransactionTermination
+34000 InvalidCursorName
+38000 ExternalRoutineException
+3B000 SavepointException
+3D000 InvalidCatalogName
+3F000 InvalidSchemaName
+40000 TransactionRollback
+40001 SerializationFailure
+40002 TransactionIntegrityConstraintViolation
+40003 StatementCompletionUnknown
+40P01 DeadlockDetected
+42000 SyntaxErrorOrAccessRuleViolation
+42501 InsufficientPrivilege
+42601 SyntaxError
+42602 InvalidName
+42611 InvalidColumnDefinition
+42622 NameTooLong
+42701 DuplicateColumn
+42702 AmbiguousColumn
+42703 UndefinedColumn
+42704 UndefinedObject
+42710 DuplicateObject
+42712 DuplicateAlias
+42723 DuplicateFunction
+42725 AmbiguousFunction
+42803 GroupingError
+42804 DatatypeMismatch
+42809 WrongObjectType
+42830 InvalidForeignKey
+42846 CannotCoerce
+42883 UndefinedFunction
+42939 ReservedName
+42P01 UndefinedTable
+42P02 UndefinedParameter
+42P03 DuplicateCursor
+42P04 DuplicateDatabase
+42P05 DuplicatePreparedStatement
+42P06 DuplicateSchema
+42P07 DuplicateTable
+42P08 AmbiguousParameter
+42P09 AmbiguousAlias
+42P10 InvalidColumnReference
+42P16 InvalidTableDefinition
+42P17 InvalidObjectDefinition
+42P18 IndeterminateDatatype
+44000 WithCheckOptionViolation
+53000 InsufficientResources
+53100 DiskFull
+53200 OutOfMemory
+53300 TooManyConnections
+54000 ProgramLimitExceeded
+55000 ObjectNotInPrerequisiteState
+55006 ObjectInUse
+55P02 CantChangeRuntimeParam
+55P03 LockNotAvailable
+57000 OperatorIntervention
+57014 QueryCanceled
+57P01 AdminShutdown
+57P02 CrashShutdown
+57P03 CannotConnectNow
+57P04 DatabaseDropped
+58000 SystemError
+58030 IoError
+P0000 PlpgsqlError
+P0001 RaiseException
+P0002 NoDataFound
+P0003 TooManyRows
+P0004 AssertFailure
+XX000 InternalError_
+XX001 DataCorrupted
+XX002 IndexCorrupted
+"""
+
+_BY_SQLSTATE: dict[str, type[Error]] = {}
+_ENTRIES = [line.split() for line in _SPEC.strip().splitlines()]
+
+# The hierarchy is flat: every specific class derives directly from its DB-API
+# family (UniqueViolation -> IntegrityError).
+for _code, _name in _ENTRIES:
+    _cls = type(_name, (_BASE_BY_CLASS.get(_code[:2], DatabaseError),),
+                {"sqlstate": _code, "__module__": __name__})
+    globals()[_name] = _cls
+    _BY_SQLSTATE[_code] = _cls
+
+__all__ = [
+    "Warning", "Error", "InterfaceError", "DatabaseError", "DataError", "OperationalError",
+    "IntegrityError", "InternalError", "ProgrammingError", "NotSupportedError", "Diagnostic", "PipelineAborted",
+    *(name for _, name in _ENTRIES),
+]
+
+
+def error_from_fields(fields: dict[str, str]) -> Error:
+    """Build the right exception from the fields of a server ErrorResponse."""
+    sqlstate = fields.get("C") or "XX000"
+    cls = _BY_SQLSTATE.get(sqlstate) or _BASE_BY_CLASS.get(sqlstate[:2], DatabaseError)
+    lines = [fields.get("M", "unknown error")]
+    for key, label in (("D", "DETAIL"), ("H", "HINT"), ("W", "CONTEXT")):
+        if key in fields:
+            lines.append(f"{label}:  {fields[key]}")
+    exc = cls("\n".join(lines), diag=Diagnostic(fields))
+    exc.sqlstate = sqlstate
+    return exc

pgtern/hooks.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Query hooks: observe every statement without monkeypatching.
+
+A hook is a callable ``hook(info) -> context manager``. The context manager is
+entered before the statement is sent and exited when its result (or error) has
+been read, with the usual ``__exit__`` exception arguments. ``info.rowcount`` is
+filled in before a successful exit.
+
+Hooks never receive bound parameter values. ``info.statement`` is the query text;
+for ``sql.Composed`` queries any ``sql.Literal`` is replaced by ``?``.
+
+    @contextmanager
+    def trace(info):
+        with tracer.start_as_current_span("postgresql") as span:
+            span.set_attribute("db.statement", info.statement)
+            yield
+
+    unregister = pgtern.add_query_hook(trace)
+"""
+from __future__ import annotations
+
+from contextlib import AbstractContextManager
+from dataclasses import dataclass
+from typing import Callable
+
+__all__ = ["QueryInfo", "add_query_hook"]
+
+
+@dataclass
+class QueryInfo:
+    statement: str
+    host: str
+    port: int
+    dbname: str
+    user: str
+    rowcount: int = -1
+
+
+QueryHook = Callable[[QueryInfo], AbstractContextManager]
+_HOOKS: list[QueryHook] = []
+
+
+def add_query_hook(hook: QueryHook) -> Callable[[], None]:
+    """Register a hook for every connection; returns a function that unregisters it."""
+    _HOOKS.append(hook)
+
+    def remove() -> None:
+        try:
+            _HOOKS.remove(hook)
+        except ValueError:
+            pass
+
+    return remove

pgtern/pipeline.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pipeline mode: send many statements before waiting for any reply.
+
+    with conn.pipeline() as pipe:
+        a = pipe.execute("INSERT INTO t (v) VALUES (%s)", [1])
+        b = pipe.execute("SELECT count(*) FROM t")
+    # leaving the block synchronizes: now a and b have their results
+    print(b.fetchone())
+
+Statements are sent immediately and their replies read, in order, at a *sync point*. Sync
+points happen when you call ``sync()``, when the block ends, and every ``batch_size``
+statements.
+
+* Transactions: with ``autocommit=False`` the first statement of a batch is preceded by an
+  implicit BEGIN, as usual. With ``autocommit=True`` every sync point is a commit point, so
+  statements in one batch succeed or fail together. Use ``batch_size`` or ``sync()`` to choose.
+* Errors: when one statement fails the server skips the rest of the batch. Those statements
+  get :class:`pgtern.errors.PipelineAborted`, and the failing statement raises its own error
+  when read. Inside a transaction you must roll back afterwards, as with any failure.
+* Only text format is used (``binary=True`` is not available in pipelines).
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from . import _protocol as P
+from . import _query, _types, errors, hooks
+from ._cursor import Cursor
+
+log = logging.getLogger("pgtern")
+
+_IDLE = "I"
+
+
+def _enter_hooks(conn: Any, item: "_Item") -> tuple[hooks.QueryInfo | None, list]:
+    if item.hidden or not hooks._HOOKS:
+        return None, []
+    info = hooks.QueryInfo(item.redacted, conn._host, conn._port, conn._params.dbname, conn._params.user)
+    entered = []
+    for hook in tuple(hooks._HOOKS):
+        try:
+            cm = hook(info)
+            cm.__enter__()
+            entered.append(cm)
+        except Exception:
+            log.exception("a pgtern query hook failed on entry; ignoring it")
+    return info, entered
+
+
+def _exit_hooks(entered: list, error: BaseException | None) -> None:
+    exc_info = (type(error), error, None) if error is not None else (None, None, None)
+    for cm in reversed(entered):
+        try:
+            cm.__exit__(*exc_info)
+        except Exception:
+            log.exception("a pgtern query hook failed on exit; ignoring it")
+
+
+class _Item:
+    __slots__ = ("cursor", "sql", "redacted", "hidden")
+
+    def __init__(self, cursor: Cursor | None, sql: str, redacted: str, hidden: bool = False) -> None:
+        self.cursor = cursor
+        self.sql = sql
+        self.redacted = redacted
+        self.hidden = hidden
+
+
+def _encode(sql: str, params: Any) -> bytes:
+    if params is None:
+        return P.extended_query(sql, [], [], sync=False)
+    pg_sql, values = _query.translate(sql, params)
+    dumped = [_types.dump(v) for v in values]
+    return P.extended_query(pg_sql, [oid for oid, _ in dumped], [data for _, data in dumped], sync=False)
+
+
+class Pipeline:
+    """Synchronous pipeline. Create it with :meth:`pgtern.Connection.pipeline`."""
+
+    def __init__(self, conn: Any, batch_size: int = 64) -> None:
+        if batch_size < 1:
+            raise ValueError("batch_size must be at least 1")
+        self._conn = conn
+        self._batch_size = batch_size
+        self._queue: list[_Item] = []
+        self._begin_queued = False
+        self._open = False
+
+    # ---------------------------------------------------------------- block
+    def __enter__(self) -> "Pipeline":
+        conn = self._conn
+        with conn._lock:
+            conn._check_open()
+            if conn._pipeline is not None:
+                raise errors.ProgrammingError("a pipeline is already active on this connection")
+            conn._pipeline = self
+        self._open = True
+        return self
+
+    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        try:
+            if exc_type is None:
+                self.sync()
+            elif not self._conn.closed:
+                try:
+                    self.sync()  # read the replies so the connection stays usable
+                except Exception:  # noqa: BLE001 - the original exception is what matters
+                    pass
+        finally:
+            self._open = False
+            with self._conn._lock:
+                if self._conn._pipeline is self:
+                    self._conn._pipeline = None
+
+    # ---------------------------------------------------------------- queueing
+    def execute(self, query: Any, params: Any = None) -> Cursor:
+        """Queue a statement; returns a cursor that holds its result after the next sync."""
+        conn = self._conn
+        with conn._lock:
+            conn._check_open()
+            if not self._open:
+                raise errors.ProgrammingError("the pipeline is not active (use it in a 'with' block)")
+            sql, redacted = conn._query_text(query)
+            if len(self._queue) >= self._batch_size:
+                self.sync()
+            if not conn._autocommit and conn._tx_status == _IDLE and not self._begin_queued:
+                conn._send(P.extended_query("BEGIN", [], [], sync=False))
+                self._queue.append(_Item(None, "BEGIN", "BEGIN", hidden=True))
+                self._begin_queued = True
+            cursor = Cursor(conn)
+            cursor._pipeline_pending = True
+            conn._send(_encode(sql, params))
+            self._queue.append(_Item(cursor, sql, redacted))
+            return cursor
+
+    # ---------------------------------------------------------------- sync point
+    def sync(self) -> None:
+        """Send Sync, read every queued reply and fill in the cursors."""
+        conn = self._conn
+        with conn._lock:
+            conn._check_open()
+            items, self._queue = self._queue, []
+            self._begin_queued = False
+            conn._send(P.SYNC)
+            failure: errors.Error | None = None
+            collected: list[tuple[Cursor, list]] = []
+            for item in items:
+                if failure is not None:
+                    if item.cursor is not None:
+                        item.cursor._pipeline_pending = False
+                        item.cursor._pipeline_error = errors.PipelineAborted(
+                            "not executed: an earlier statement in the pipeline failed")
+                    continue
+                info, entered = _enter_hooks(conn, item)
+                results, error = self._read_reply(conn)
+                if info is not None and results is not None:
+                    info.rowcount = results[-1].rowcount if results else -1
+                _exit_hooks(entered, error)
+                if error is not None:
+                    failure = error
+                    if item.cursor is not None:
+                        item.cursor._pipeline_pending = False
+                        item.cursor._pipeline_error = error
+                elif item.cursor is not None and results is not None:
+                    collected.append((item.cursor, results))
+            while True:  # the rest of a failed batch sends nothing but ReadyForQuery
+                tag, payload = conn._next_message()
+                if tag == 0x5A:
+                    conn._tx_status = chr(payload[0])
+                    break
+                if tag == 0x45:
+                    fields = P.parse_error_fields(payload)
+                    if fields.get("V", fields.get("S")) in ("FATAL", "PANIC"):
+                        conn._mark_broken()
+                        raise errors.error_from_fields(fields)
+                else:
+                    conn._handle_async_message(tag, payload)
+            conn._prime_loaders([r for _, res in collected for r in res])
+            for cursor, res in collected:
+                cursor._pipeline_pending = False
+                cursor._set_results(res)
+
+    @staticmethod
+    def _read_reply(conn: Any) -> tuple[list | None, errors.Error | None]:
+        """Read one statement's messages. Returns (results, None), or (None, error)."""
+        collector = P.QueryCollector()
+        while True:
+            tag, payload = conn._next_message()
+            if tag == 0x45:
+                fields = P.parse_error_fields(payload)
+                if fields.get("V", fields.get("S")) in ("FATAL", "PANIC"):
+                    conn._mark_broken()
+                    raise errors.error_from_fields(fields)
+                return None, errors.error_from_fields(fields)
+            if tag in (0x43, 0x49):  # C CommandComplete, I EmptyQueryResponse: this statement is done
+                collector.handle(tag, payload)
+                return collector.results, None
+            if tag in (0x4E, 0x53, 0x41):
+                conn._handle_async_message(tag, payload)
+            else:
+                collector.handle(tag, payload)