pgtern 0.3.0__py3-none-any.whl

pgtern/pool.py

@@ -0,0 +1,309 @@
+# SPDX-License-Identifier: Apache-2.0
+"""A synchronous, thread-safe connection pool. Built in: no separate package to install."""
+from __future__ import annotations
+
+import collections
+import logging
+import threading
+import time
+from contextlib import contextmanager
+from typing import Any, Callable, Iterator
+
+from . import errors
+from ._conn import Connection
+
+log = logging.getLogger("pgtern.pool")
+
+__all__ = ["ConnectionPool", "PoolClosed", "PoolTimeout"]
+
+
+class PoolClosed(errors.OperationalError):
+    """The pool was closed."""
+
+
+class PoolTimeout(errors.OperationalError):
+    """No connection became available in time."""
+
+
+class _Idle:
+    __slots__ = ("conn", "returned")
+
+    def __init__(self, conn: Connection, returned: float) -> None:
+        self.conn = conn
+        self.returned = returned
+
+
+class ConnectionPool:
+    """Pool of :class:`Connection`.
+
+    ``min_size`` connections are kept open by a background thread (which also retires
+    connections older than ``max_lifetime`` and surplus ones idle longer than ``max_idle``);
+    up to ``max_size`` are opened on demand. A connection is checked for liveness
+    (no round trip) every time it is handed out, a failed connection is replaced, and
+    a connection returned mid-transaction is rolled back.
+    """
+
+    def __init__(self, conninfo: str = "", *, kwargs: dict[str, Any] | None = None, min_size: int = 4,
+                 max_size: int | None = None, open: bool = True, timeout: float = 30.0,
+                 max_lifetime: float = 3600.0, max_idle: float = 600.0,
+                 check: Callable[[Connection], None] | None = None, name: str | None = None,
+                 connection_class: type[Connection] = Connection) -> None:
+        max_size = min_size if max_size is None else max_size
+        if min_size < 0 or max_size < 1 or max_size < min_size:
+            raise ValueError(f"bad pool sizes: min_size={min_size}, max_size={max_size}")
+        self.conninfo = conninfo
+        self.kwargs = dict(kwargs or {})
+        self.min_size, self.max_size = min_size, max_size
+        self.timeout, self.max_lifetime, self.max_idle = timeout, max_lifetime, max_idle
+        self.name = name or f"pool-{id(self):x}"
+        self._check = check
+        self._connection_class = connection_class
+        self._default_autocommit = bool(self.kwargs.get("autocommit", False))
+        self._cond = threading.Condition()
+        self._idle: collections.deque[_Idle] = collections.deque()
+        self._n = 0            # open + being opened
+        self._connecting = 0
+        self._waiting = 0
+        self._closed = True
+        self._thread: threading.Thread | None = None
+        self._wake = threading.Event()
+        if open:
+            self.open()
+
+    # ---------------------------------------------------------------- lifecycle
+    def open(self, wait: bool = False, timeout: float = 30.0) -> None:
+        with self._cond:
+            if not self._closed:
+                return
+            self._closed = False
+        self._wake.clear()
+        self._thread = threading.Thread(target=self._maintain, name=f"pgtern-{self.name}", daemon=True)
+        self._thread.start()
+        if wait:
+            self.wait(timeout)
+
+    def wait(self, timeout: float = 30.0) -> None:
+        """Block until ``min_size`` connections are open."""
+        deadline = time.monotonic() + timeout
+        with self._cond:
+            while self._n - self._connecting < self.min_size:
+                if self._closed:
+                    raise PoolClosed("the pool is closed")
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise PoolTimeout(f"pool initialization incomplete after {timeout:.2f} sec")
+                self._cond.wait(min(remaining, 0.2))
+
+    def close(self, timeout: float = 5.0) -> None:
+        with self._cond:
+            if self._closed:
+                return
+            self._closed = True
+            idle = list(self._idle)
+            self._idle.clear()
+            self._n -= len(idle)
+            self._cond.notify_all()
+        self._wake.set()
+        for item in idle:
+            item.conn.close()
+        thread = self._thread
+        if thread is not None and thread is not threading.current_thread():
+            thread.join(timeout)
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    def __enter__(self) -> "ConnectionPool":
+        self.open()
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
+    def get_stats(self) -> dict[str, int]:
+        with self._cond:
+            return {
+                "pool_min": self.min_size, "pool_max": self.max_size,
+                "pool_size": self._n - self._connecting, "pool_available": len(self._idle),
+                "requests_waiting": self._waiting,
+            }
+
+    # ---------------------------------------------------------------- internals
+    def _connect(self) -> Connection:
+        conn = self._connection_class.connect(self.conninfo, **self.kwargs)
+        conn._pool_created = time.monotonic()  # type: ignore[attr-defined]
+        return conn
+
+    def _reap(self) -> list[Connection]:
+        """Pick idle connections to retire. Caller holds the lock and closes the result."""
+        now = time.monotonic()
+        keep: collections.deque[_Idle] = collections.deque()
+        retire: list[Connection] = []
+        for item in self._idle:
+            conn = item.conn
+            too_old = now - conn._pool_created > self.max_lifetime  # type: ignore[attr-defined]
+            surplus_idle = now - item.returned > self.max_idle and self._n - len(retire) > self.min_size
+            if conn.closed or too_old or surplus_idle:
+                retire.append(conn)
+            else:
+                keep.append(item)
+        self._idle = keep
+        self._n -= len(retire)
+        return retire
+
+    def _maintain(self) -> None:
+        backoff = 0.5
+        while True:
+            create = False
+            with self._cond:
+                if self._closed:
+                    return
+                retire = self._reap()
+                if self._n < self.min_size:
+                    self._n += 1
+                    self._connecting += 1
+                    create = True
+            for conn in retire:
+                conn.close()
+            if create:
+                try:
+                    conn = self._connect()
+                except Exception as exc:
+                    with self._cond:
+                        self._n -= 1
+                        self._connecting -= 1
+                    log.warning("pool %s: couldn't open a connection: %s", self.name, exc)
+                    self._wake.wait(backoff)
+                    self._wake.clear()
+                    backoff = min(backoff * 2, 30.0)
+                    continue
+                backoff = 0.5
+                with self._cond:
+                    self._connecting -= 1
+                    stored = not self._closed
+                    if stored:
+                        self._idle.append(_Idle(conn, time.monotonic()))
+                    else:
+                        self._n -= 1
+                    self._cond.notify_all()
+                if not stored:
+                    conn.close()
+                continue
+            self._wake.wait(5.0)
+            self._wake.clear()
+
+    def _usable(self, conn: Connection) -> bool:
+        if conn.closed or time.monotonic() - conn._pool_created > self.max_lifetime:  # type: ignore[attr-defined]
+            return False
+        if not conn._check_alive():
+            return False
+        if self._check is not None:
+            try:
+                self._check(conn)
+            except Exception:
+                return False
+        return True
+
+    def _discard(self, conn: Connection) -> None:
+        with self._cond:
+            self._n -= 1
+            self._cond.notify_all()
+        self._wake.set()
+        conn.close()
+
+    # ---------------------------------------------------------------- public API
+    def getconn(self, timeout: float | None = None) -> Connection:
+        timeout = self.timeout if timeout is None else timeout
+        deadline = time.monotonic() + timeout
+        backoff = 0.1
+        while True:
+            candidate: _Idle | None = None
+            create = False
+            with self._cond:
+                if self._closed:
+                    raise PoolClosed("the pool is closed")
+                if self._idle:
+                    candidate = self._idle.pop()
+                elif self._n < self.max_size:
+                    self._n += 1
+                    self._connecting += 1
+                    create = True
+                else:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        raise PoolTimeout(f"couldn't get a connection after {timeout:.2f} sec")
+                    self._waiting += 1
+                    try:
+                        self._cond.wait(remaining)
+                    finally:
+                        self._waiting -= 1
+                    continue
+            if candidate is not None:
+                if self._usable(candidate.conn):
+                    return candidate.conn
+                self._discard(candidate.conn)
+                continue
+            assert create
+            try:
+                conn = self._connect()
+            except Exception as exc:
+                with self._cond:
+                    self._n -= 1
+                    self._connecting -= 1
+                    self._cond.notify_all()
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise PoolTimeout(f"couldn't get a connection after {timeout:.2f} sec: {exc}") from exc
+                time.sleep(min(backoff, remaining))
+                backoff = min(backoff * 2, 2.0)
+                continue
+            with self._cond:
+                self._connecting -= 1
+                closed = self._closed
+                if closed:
+                    self._n -= 1
+            if closed:
+                conn.close()
+                raise PoolClosed("the pool is closed")
+            return conn
+
+    def putconn(self, conn: Connection) -> None:
+        ok = not conn.closed
+        if ok:
+            try:
+                conn._reset_for_pool()
+                if conn.autocommit != self._default_autocommit:
+                    conn.autocommit = self._default_autocommit
+            except Exception:
+                ok = False
+        with self._cond:
+            fresh = time.monotonic() - conn._pool_created <= self.max_lifetime  # type: ignore[attr-defined]
+            if ok and fresh and not self._closed:
+                self._idle.append(_Idle(conn, time.monotonic()))
+                self._cond.notify_all()
+                return
+            self._n -= 1
+            self._cond.notify_all()
+        self._wake.set()
+        conn.close()
+
+    @contextmanager
+    def connection(self, timeout: float | None = None) -> Iterator[Connection]:
+        """Borrow a connection: commit on success, roll back on error, return it to the pool."""
+        conn = self.getconn(timeout)
+        try:
+            try:
+                yield conn
+            except BaseException:
+                try:
+                    if not conn.closed:
+                        conn._reset_for_pool()
+                except Exception:
+                    log.warning("rollback failed while returning a connection", exc_info=True)
+                raise
+            else:
+                if not conn.closed:
+                    conn.commit()
+        finally:
+            self.putconn(conn)

pgtern/rows.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Row factories: callables taking a cursor and returning a ``make_row(values)`` function."""
+from __future__ import annotations
+
+from collections import namedtuple
+from typing import Any, Callable
+
+__all__ = ["tuple_row", "dict_row", "namedtuple_row"]
+
+
+def _identity(values: tuple) -> tuple:
+    return values
+
+
+def tuple_row(cursor: Any) -> Callable[[tuple], tuple]:
+    """Rows as plain tuples (the default)."""
+    return _identity
+
+
+def dict_row(cursor: Any) -> Callable[[tuple], dict[str, Any]]:
+    """Rows as ``dict`` keyed by column name (a repeated name keeps the last value)."""
+    names = [c.name for c in (cursor.description or ())]
+    return lambda values: dict(zip(names, values))
+
+
+def namedtuple_row(cursor: Any) -> Callable[[tuple], Any]:
+    """Rows as ``namedtuple`` instances (unusable names are renamed ``_0``, ``_1``, ...)."""
+    names = [c.name for c in (cursor.description or ())]
+    factory: Any = namedtuple  # typed as Any: the field list is only known at run time
+    row_type = factory("Row", names, rename=True)
+    return row_type._make

pgtern/sql.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Safe composition of SQL text: ``SQL``, ``Identifier``, ``Literal``, ``Placeholder``.
+
+Use this for statements that cannot take bind parameters (``CREATE ROLE ... PASSWORD``,
+``GRANT ... TO <role>``). Everything else should use ``%s`` parameters.
+"""
+from __future__ import annotations
+
+import binascii
+import string
+import uuid
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal
+from typing import Any, Iterable, Iterator
+
+from . import errors
+from .types import Json, Jsonb
+
+__all__ = ["Composable", "Composed", "SQL", "Identifier", "Literal", "Placeholder",
+           "quote_ident", "quote_literal"]
+
+
+def quote_ident(name: str) -> str:
+    if "\x00" in name:
+        raise errors.DataError("PostgreSQL identifiers cannot contain NUL (0x00) bytes")
+    return '"' + name.replace('"', '""') + '"'
+
+
+def quote_literal(text: str) -> str:
+    """Quote a string as a PostgreSQL literal, independent of ``standard_conforming_strings``.
+
+    Backslashes force the ``E'...'`` form with the backslash doubled, so the result
+    means the same thing whichever way the server setting is configured.
+    """
+    if "\x00" in text:
+        raise errors.DataError("PostgreSQL text fields cannot contain NUL (0x00) bytes")
+    body = text.replace("'", "''")
+    if "\\" in body:
+        return " E'" + body.replace("\\", "\\\\") + "'"
+    return "'" + body + "'"
+
+
+class Composable:
+    __slots__ = ("_obj",)
+
+    def __init__(self, obj: Any) -> None:
+        self._obj = obj
+
+    def as_string(self, context: Any = None) -> str:
+        raise NotImplementedError
+
+    def as_redacted(self) -> str:
+        """The statement with every Literal hidden: safe to log or put on a trace span."""
+        return self.as_string()
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self._obj!r})"
+
+    def __add__(self, other: Any) -> "Composed":
+        if isinstance(other, Composed):
+            return Composed([self, *other._obj])
+        if isinstance(other, Composable):
+            return Composed([self, other])
+        return NotImplemented
+
+    def __mul__(self, n: int) -> "Composed":
+        return Composed([self] * n)
+
+    def __eq__(self, other: object) -> bool:
+        return type(self) is type(other) and self._obj == other._obj  # type: ignore[attr-defined]
+
+    def __hash__(self) -> int:
+        return hash((type(self), repr(self._obj)))
+
+
+class Composed(Composable):
+    def __init__(self, seq: Iterable[Composable]) -> None:
+        items = list(seq)
+        for item in items:
+            if not isinstance(item, Composable):
+                raise TypeError(f"Composed elements must be Composable, got {item!r}")
+        super().__init__(items)
+
+    def __iter__(self) -> Iterator[Composable]:
+        return iter(self._obj)
+
+    def as_string(self, context: Any = None) -> str:
+        return "".join(item.as_string(context) for item in self._obj)
+
+    def as_redacted(self) -> str:
+        return "".join(item.as_redacted() for item in self._obj)
+
+    def join(self, joiner: "str | SQL") -> "Composed":
+        sep = SQL(joiner) if isinstance(joiner, str) else joiner
+        out: list[Composable] = []
+        for index, item in enumerate(self._obj):
+            if index:
+                out.append(sep)
+            out.append(item)
+        return Composed(out)
+
+
+class SQL(Composable):
+    def __init__(self, text: str) -> None:
+        if not isinstance(text, str):
+            raise TypeError(f"SQL values must be strings, got {text!r}")
+        super().__init__(text)
+
+    def as_string(self, context: Any = None) -> str:
+        return self._obj
+
+    def format(self, *args: Composable, **kwargs: Composable) -> Composed:
+        parts: list[Composable] = []
+        auto = 0
+        for literal, field, spec, conversion in string.Formatter().parse(self._obj):
+            if spec or conversion:
+                raise ValueError("no format specification supported by SQL")
+            if literal:
+                parts.append(SQL(literal))
+            if field is None:
+                continue
+            if field == "":
+                value = args[auto]
+                auto += 1
+            elif field.isdigit():
+                value = args[int(field)]
+            else:
+                value = kwargs[field]
+            if not isinstance(value, Composable):
+                raise TypeError(f"SQL.format() arguments must be Composable, got {value!r}")
+            parts.append(value)
+        return Composed(parts)
+
+    def join(self, seq: Iterable[Composable]) -> Composed:
+        return Composed(list(seq)).join(self)
+
+
+class Identifier(Composable):
+    def __init__(self, *strings: str) -> None:
+        if not strings:
+            raise TypeError("Identifier requires at least one string")
+        for s in strings:
+            if not isinstance(s, str):
+                raise TypeError(f"SQL identifier parts must be strings, got {s!r}")
+        super().__init__(strings)
+
+    def as_string(self, context: Any = None) -> str:
+        return ".".join(quote_ident(s) for s in self._obj)
+
+
+class Literal(Composable):
+    def as_string(self, context: Any = None) -> str:
+        return _literal(self._obj)
+
+    def as_redacted(self) -> str:
+        return "?"
+
+
+class Placeholder(Composable):
+    def __init__(self, name: str = "") -> None:
+        super().__init__(name)
+
+    def as_string(self, context: Any = None) -> str:
+        return f"%({self._obj})s" if self._obj else "%s"
+
+
+def _literal(value: Any) -> str:
+    if value is None:
+        return "NULL"
+    if value is True or value is False:
+        return "true" if value else "false"
+    if isinstance(value, int):
+        return f" {value}" if value < 0 else str(value)
+    if isinstance(value, float):
+        if value != value or value in (float("inf"), float("-inf")):
+            return quote_literal(str(value).replace("inf", "Infinity").replace("nan", "NaN")) + "::float8"
+        return f" {value!r}" if value < 0 else repr(value)
+    if isinstance(value, Decimal):
+        if not value.is_finite():
+            return quote_literal(str(value)) + "::numeric"
+        return f" {value}" if value < 0 else str(value)
+    if isinstance(value, str):
+        return quote_literal(value)
+    if isinstance(value, (bytes, bytearray, memoryview)):
+        return "'\\x" + binascii.hexlify(bytes(value)).decode() + "'::bytea"
+    if isinstance(value, datetime):
+        return quote_literal(value.isoformat(" ")) + ("::timestamptz" if value.tzinfo else "::timestamp")
+    if isinstance(value, date):
+        return quote_literal(value.isoformat()) + "::date"
+    if isinstance(value, time):
+        return quote_literal(value.isoformat()) + ("::timetz" if value.tzinfo else "::time")
+    if isinstance(value, timedelta):
+        return quote_literal(f"{value.days} days {value.seconds} seconds {value.microseconds} microseconds") + "::interval"
+    if isinstance(value, uuid.UUID):
+        return quote_literal(str(value)) + "::uuid"
+    if isinstance(value, Json):
+        return quote_literal(value.encode()) + ("::jsonb" if isinstance(value, Jsonb) else "::json")
+    if isinstance(value, list):
+        return "ARRAY[" + ",".join(_literal(v) for v in value) + "]" if value else "'{}'"
+    raise errors.ProgrammingError(f"cannot make a SQL literal from {type(value).__name__!r}")

pgtern/types/__init__.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wrapper types that tell pgtern how to send a Python value to PostgreSQL."""
+from __future__ import annotations
+
+import json as _json
+from typing import Any, Callable
+
+__all__ = ["Json", "Jsonb"]
+
+
+class Json:
+    """Send a Python object as a PostgreSQL ``json`` value."""
+
+    __slots__ = ("obj", "dumps")
+
+    def __init__(self, obj: Any, dumps: Callable[[Any], str] | None = None) -> None:
+        self.obj = obj
+        self.dumps = dumps or _json.dumps
+
+    def encode(self) -> str:
+        return self.dumps(self.obj)
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self.obj!r})"
+
+
+class Jsonb(Json):
+    """Send a Python object as a PostgreSQL ``jsonb`` value."""
+
+    __slots__ = ()

pgtern/types/json.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``pgtern.types.json``: JSON wrappers for query parameters."""
+from . import Json, Jsonb
+
+__all__ = ["Json", "Jsonb"]