pgtern 0.3.0__py3-none-any.whl

pgtern/_copy.py

@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+"""COPY support: text-format helpers and the synchronous ``Copy`` object."""
+from __future__ import annotations
+
+from contextlib import ExitStack
+from typing import Any, Iterator, Sequence
+
+from . import _protocol as P
+from . import _types, errors
+
+_ESCAPES = {ord("\\"): "\\\\", ord("\t"): "\\t", ord("\n"): "\\n", ord("\r"): "\\r"}
+_UNESCAPE = {b"b": b"\b", b"f": b"\f", b"n": b"\n", b"r": b"\r", b"t": b"\t", b"v": b"\v", b"\\": b"\\"}
+
+
+def format_row(row: Sequence[Any]) -> bytes:
+    """One row in COPY text format (tab separated, ``\\N`` for NULL), newline terminated."""
+    fields = []
+    for value in row:
+        _oid, data = _types.dump(value)
+        fields.append("\\N" if data is None else data.decode("utf-8").translate(_ESCAPES))
+    return ("\t".join(fields) + "\n").encode("utf-8")
+
+
+def _unescape(field: bytes) -> str | None:
+    if field == b"\\N":
+        return None
+    if b"\\" not in field:
+        return field.decode("utf-8")
+    out = bytearray()
+    i, n = 0, len(field)
+    while i < n:
+        c = field[i:i + 1]
+        if c != b"\\" or i + 1 >= n:
+            out += c
+            i += 1
+            continue
+        nxt = field[i + 1:i + 2]
+        if nxt in _UNESCAPE:
+            out += _UNESCAPE[nxt]
+            i += 2
+        elif nxt == b"x":
+            j = i + 2
+            while j < n and j < i + 4 and chr(field[j]) in "0123456789abcdefABCDEF":
+                j += 1
+            out.append(int(field[i + 2:j], 16)) if j > i + 2 else out.extend(b"x")
+            i = max(j, i + 2)
+        elif nxt in b"01234567":
+            j = i + 1
+            while j < n and j < i + 4 and field[j:j + 1] in (b"0", b"1", b"2", b"3", b"4", b"5", b"6", b"7"):
+                j += 1
+            out.append(int(field[i + 1:j], 8) & 0xFF)
+            i = j
+        else:
+            out += nxt
+            i += 2
+    return out.decode("utf-8")
+
+
+def parse_rows(buffer: bytearray, final: bool = False) -> list[tuple]:
+    """Consume complete text-format lines from ``buffer``; return them as tuples of str/None."""
+    rows: list[tuple] = []
+    while True:
+        end = buffer.find(b"\n")
+        if end < 0:
+            if final and buffer:
+                end = len(buffer)
+            else:
+                return rows
+        line = bytes(buffer[:end])
+        del buffer[:end + 1]
+        rows.append(tuple(_unescape(f) for f in line.split(b"\t")))
+
+
+class Copy:
+    """``with cursor.copy("COPY t FROM STDIN") as copy: copy.write_row(...)`` or
+    ``with cursor.copy("COPY t TO STDOUT") as copy: for chunk in copy: ...``."""
+
+    def __init__(self, cursor: Any, query: Any, params: Any = None) -> None:
+        self._cursor = cursor
+        self._conn = cursor.connection
+        self._query = query
+        self._params = params
+        self._kind = ""
+        self._stack = ExitStack()
+        self._done = False
+        self._buffer = bytearray()
+        self._results: list = []
+
+    # ---------------------------------------------------------------- context
+    def __enter__(self) -> "Copy":
+        conn = self._conn
+        conn._lock.acquire()
+        try:
+            self._kind = conn._start_copy(self._query, self._params, self._stack)
+        except BaseException:
+            self._stack.close()
+            conn._lock.release()
+            raise
+        return self
+
+    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        conn = self._conn
+        try:
+            if conn.closed:
+                return
+            if self._kind == "in":
+                if exc_type is not None:
+                    conn._send(P.copy_fail(f"error in Python: {exc_type.__name__}"))
+                    try:
+                        conn._read_cycle()
+                    except errors.Error:
+                        pass  # the expected "COPY from stdin failed"
+                else:
+                    conn._send(P.COPY_DONE)
+                    self._results = conn._read_cycle()
+            else:
+                while not self._done:
+                    self._next_chunk()
+                self._results = self._finish()
+            if exc_type is None:
+                self._cursor._set_results(self._results)
+        finally:
+            try:
+                self._stack.__exit__(exc_type, exc, tb)
+            finally:
+                conn._lock.release()
+
+    # ---------------------------------------------------------------- COPY FROM STDIN
+    def write(self, data: bytes | bytearray | memoryview | str) -> None:
+        if self._kind != "in":
+            raise errors.ProgrammingError("write() needs a COPY ... FROM STDIN statement")
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+        data = bytes(data)
+        for i in range(0, max(len(data), 1), 1 << 20):
+            self._conn._send(P.copy_data(data[i:i + (1 << 20)]))
+
+    def write_row(self, row: Sequence[Any]) -> None:
+        """Write one row in text format (``COPY ... FROM STDIN`` default, or ``WITH (FORMAT text)``)."""
+        self.write(format_row(row))
+
+    # ---------------------------------------------------------------- COPY TO STDOUT
+    def _next_chunk(self) -> bytes | None:
+        conn = self._conn
+        while True:
+            tag, payload = conn._next_message()
+            if tag == 0x64:  # d CopyData
+                return payload
+            if tag == 0x63:  # c CopyDone
+                self._done = True
+                return None
+            if tag == 0x45:
+                self._done = True
+                self._pending_error = P.parse_error_fields(payload)
+                return None
+            conn._handle_async_message(tag, payload)
+
+    def _finish(self) -> list:
+        err = getattr(self, "_pending_error", None)
+        results = self._conn._read_cycle()
+        if err is not None:
+            raise errors.error_from_fields(err)
+        return results
+
+    def read(self) -> bytes:
+        """The next chunk of COPY data, or ``b""`` at the end."""
+        if self._kind != "out":
+            raise errors.ProgrammingError("read() needs a COPY ... TO STDOUT statement")
+        if self._done:
+            return b""
+        chunk = self._next_chunk()
+        return chunk if chunk is not None else b""
+
+    def __iter__(self) -> Iterator[bytes]:
+        while True:
+            chunk = self.read()
+            if not chunk:
+                return
+            yield chunk
+
+    def rows(self) -> Iterator[tuple]:
+        """Parsed rows (tuples of ``str``/``None``) of a text-format ``COPY ... TO STDOUT``."""
+        for chunk in self:
+            self._buffer += chunk
+            yield from parse_rows(self._buffer)
+        yield from parse_rows(self._buffer, final=True)

pgtern/_cursor.py

@@ -0,0 +1,261 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DB-API cursor. Results are fully buffered on the client."""
+from __future__ import annotations
+
+from typing import Any, Callable, Iterator, Sequence
+
+from . import errors
+from ._copy import Copy
+from ._protocol import Column, Result
+from ._types import load_row, make_loaders
+from .sql import SQL, Composable, quote_ident
+
+
+class Cursor:
+    def __init__(self, connection: Any, *, row_factory: Callable | None = None) -> None:
+        self.connection = connection
+        self.row_factory = row_factory or connection.row_factory
+        self.arraysize = 1
+        self._results: list[Result] = []
+        self._index = 0
+        self._pos = 0
+        self._make_row: Callable | None = None
+        self._loaders: list | None = None
+        self._total_rowcount: int | None = None
+        self._closed = False
+        self._pipeline_pending = False  # set by a pipeline until its batch is synchronized
+        self._pipeline_error: Exception | None = None
+
+    # ---------------------------------------------------------------- state
+    def _current(self) -> Result | None:
+        if self._pipeline_pending:
+            raise errors.ProgrammingError("the pipeline is not synchronized yet: leave the pipeline block first")
+        if self._pipeline_error is not None:
+            raise self._pipeline_error
+        return self._results[self._index] if self._index < len(self._results) else None
+
+    @property
+    def description(self) -> list[Column] | None:
+        cur = self._current()
+        return list(cur.columns) if cur is not None and cur.columns is not None else None
+
+    @property
+    def rowcount(self) -> int:
+        if self._total_rowcount is not None:
+            return self._total_rowcount
+        cur = self._current()
+        return cur.rowcount if cur is not None else -1
+
+    @property
+    def statusmessage(self) -> str | None:
+        cur = self._current()
+        return cur.tag if cur is not None else None
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    def close(self) -> None:
+        self._closed = True
+        self._results = []
+
+    def __enter__(self) -> "Cursor":
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
+    def _reset_position(self, index: int) -> None:
+        self._index = index
+        self._pos = 0
+        self._make_row = None
+        self._loaders = None
+
+    # ---------------------------------------------------------------- execute
+    def execute(self, query: Any, params: Any = None, *, prepare: bool | None = None,
+                binary: bool | None = None) -> "Cursor":
+        """Run a query. ``binary=True`` sends and receives binary format where the type supports
+        it (one extra round trip). ``prepare`` is accepted for DB-API compatibility and ignored:
+        pgtern always uses the unnamed statement."""
+        if self._closed:
+            raise errors.InterfaceError("the cursor is closed")
+        self._set_results(self.connection._execute(query, params, bool(binary)))
+        return self
+
+    def _set_results(self, results: list[Result]) -> None:
+        self._results = results
+        self._total_rowcount = None
+        self._reset_position(0)
+
+    def copy(self, query: Any, params: Any = None) -> Copy:
+        """Run ``COPY ... FROM STDIN`` / ``TO STDOUT``; use the result as a context manager."""
+        if self._closed:
+            raise errors.InterfaceError("the cursor is closed")
+        return Copy(self, query, params)
+
+    def executemany(self, query: Any, params_seq: Sequence[Any]) -> None:
+        """Run ``query`` once per parameter set as one pipelined batch, in one transaction
+        (see ``Connection._run_many``). ``rowcount`` is the total of the statements' rowcounts."""
+        if self._closed:
+            raise errors.InterfaceError("the cursor is closed")
+        total = self.connection._run_many(query, params_seq)
+        self._set_results([])
+        self._total_rowcount = total
+
+    def nextset(self) -> bool | None:
+        if self._index + 1 < len(self._results):
+            self._reset_position(self._index + 1)
+            return True
+        return None
+
+    # ---------------------------------------------------------------- fetch
+    def _prepare_rows(self) -> tuple[Result, list[Any]]:
+        cur = self._current()
+        if cur is None or cur.columns is None:
+            raise errors.ProgrammingError("the last operation didn't produce a result")
+        if self._loaders is None:
+            self._loaders = make_loaders([c.type_code for c in cur.columns], self.connection._dates_ok,
+                                         self.connection._loader_cache, cur.formats)
+        if self._make_row is None:
+            self._make_row = self.row_factory(self)
+        return cur, self._loaders
+
+    def fetchone(self) -> Any:
+        cur, loaders = self._prepare_rows()
+        if self._pos >= len(cur.rows):
+            return None
+        raw = cur.rows[self._pos]
+        self._pos += 1
+        return self._make_row(load_row(raw, loaders))  # type: ignore[misc]
+
+    def fetchmany(self, size: int | None = None) -> list:
+        cur, loaders = self._prepare_rows()
+        count = self.arraysize if size is None else size
+        chunk = cur.rows[self._pos:self._pos + count]
+        self._pos += len(chunk)
+        return [self._make_row(load_row(raw, loaders)) for raw in chunk]  # type: ignore[misc]
+
+    def fetchall(self) -> list:
+        cur, loaders = self._prepare_rows()
+        chunk = cur.rows[self._pos:]
+        self._pos = len(cur.rows)
+        return [self._make_row(load_row(raw, loaders)) for raw in chunk]  # type: ignore[misc]
+
+    def __iter__(self) -> Iterator:
+        return self
+
+    def __next__(self) -> Any:
+        row = self.fetchone()
+        if row is None:
+            raise StopIteration
+        return row
+
+
+class ServerCursor(Cursor):
+    """A named cursor living in the server (``DECLARE``): rows are fetched in batches,
+    so large results never have to fit in memory. Needs a transaction unless ``withhold``."""
+
+    itersize = 100
+
+    def __init__(self, connection: Any, name: str, *, row_factory: Callable | None = None,
+                 scrollable: bool | None = None, withhold: bool = False) -> None:
+        super().__init__(connection, row_factory=row_factory)
+        self.name = name
+        self._scrollable = scrollable
+        self._withhold = withhold
+        self._binary = False
+        self._declared = False
+        self._columns: list[Column] | None = None
+        self._fetched = 0
+        self._batch: Result | None = None
+
+    # -------------------------------------------------------------- state
+    @property
+    def description(self) -> list[Column] | None:
+        return list(self._columns) if self._columns is not None else None
+
+    @property
+    def rowcount(self) -> int:
+        return -1 if not self._declared else self._fetched
+
+    @property
+    def rownumber(self) -> int | None:
+        return self._fetched if self._declared else None
+
+    @property
+    def statusmessage(self) -> str | None:
+        return self._batch.tag if self._batch is not None else None
+
+    def close(self) -> None:
+        if self._declared and not self.connection.closed and self.connection.info.transaction_status.name != "INERROR":
+            try:
+                self.connection._execute(f"CLOSE {quote_ident(self.name)}", None)
+            except errors.Error:
+                pass
+        self._declared = False
+        super().close()
+
+    # -------------------------------------------------------------- execute
+    def execute(self, query: Any, params: Any = None, *, binary: bool | None = None, **_: Any) -> "ServerCursor":
+        if self._closed:
+            raise errors.InterfaceError("the cursor is closed")
+        self._binary = bool(binary)
+        if self._declared:
+            raise errors.ProgrammingError("the server-side cursor was already executed; create a new one")
+        prefix = f"DECLARE {quote_ident(self.name)} "
+        if self._scrollable is not None:
+            prefix += "SCROLL " if self._scrollable else "NO SCROLL "
+        prefix += "CURSOR " + ("WITH HOLD " if self._withhold else "") + "FOR "
+        body = query if isinstance(query, Composable) else SQL(query if isinstance(query, str) else bytes(query).decode())
+        self.connection._execute(SQL(prefix) + body, params)
+        self._declared = True
+        self._fetched = 0
+        first = self.connection._execute(f"FETCH FORWARD 0 FROM {quote_ident(self.name)}", None, self._binary)[0]
+        self._columns = first.columns
+        self._set_batch(first)
+        return self
+
+    def executemany(self, query: Any, params_seq: Sequence[Any]) -> None:
+        raise errors.NotSupportedError("executemany() can't be used with a server-side cursor")
+
+    def _set_batch(self, result: Result) -> None:
+        self._batch = result
+        self._set_results([result])
+
+    def _fetch(self, command: str) -> Result:
+        if not self._declared:
+            raise errors.ProgrammingError("the cursor has not been executed")
+        result = self.connection._execute(f"{command} FROM {quote_ident(self.name)}", None, self._binary)[0]
+        if result.columns is None:
+            result.columns = self._columns
+        self._set_batch(result)
+        return result
+
+    # -------------------------------------------------------------- fetch
+    def fetchone(self) -> Any:
+        rows = self.fetchmany(1)
+        return rows[0] if rows else None
+
+    def fetchmany(self, size: int | None = None) -> list:
+        count = self.arraysize if size is None else size
+        result = self._fetch(f"FETCH FORWARD {int(count)}")
+        self._fetched += len(result.rows)
+        return Cursor.fetchall(self)
+
+    def fetchall(self) -> list:
+        result = self._fetch("FETCH FORWARD ALL")
+        self._fetched += len(result.rows)
+        return Cursor.fetchall(self)
+
+    def __iter__(self) -> Iterator:
+        while True:
+            rows = self.fetchmany(self.itersize)
+            yield from rows
+            if len(rows) < self.itersize:
+                return
+
+    def scroll(self, value: int, mode: str = "relative") -> None:
+        if mode not in ("relative", "absolute"):
+            raise ValueError(f"bad mode: {mode}; it should be 'relative' or 'absolute'")
+        self._fetch(f"MOVE {mode.upper()} {int(value)}")
+        self._fetched = (self._fetched + value) if mode == "relative" else value

pgtern/_dsn.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Connection strings: ``postgresql://`` URLs and ``key=value`` strings, plus ``PG*`` environment defaults."""
+from __future__ import annotations
+
+import getpass
+import os
+from dataclasses import dataclass
+from urllib.parse import parse_qsl, unquote
+
+from . import errors
+
+SSL_MODES = ("disable", "allow", "prefer", "require", "verify-ca", "verify-full")
+
+_ENV = {
+    "host": "PGHOST", "port": "PGPORT", "user": "PGUSER", "password": "PGPASSWORD",
+    "dbname": "PGDATABASE", "sslmode": "PGSSLMODE", "sslrootcert": "PGSSLROOTCERT",
+    "sslcert": "PGSSLCERT", "sslkey": "PGSSLKEY", "connect_timeout": "PGCONNECT_TIMEOUT",
+    "application_name": "PGAPPNAME", "options": "PGOPTIONS",
+}
+_KNOWN = set(_ENV) | {"database"}
+
+
+@dataclass
+class ConnParams:
+    hosts: list[tuple[str, int]]
+    user: str
+    password: str | None
+    dbname: str
+    sslmode: str = "prefer"
+    sslrootcert: str | None = None
+    sslcert: str | None = None
+    sslkey: str | None = None
+    connect_timeout: float | None = None
+    application_name: str | None = None
+    options: str | None = None
+
+
+def _parse_url(url: str) -> dict[str, str]:
+    scheme, _, rest = url.partition("://")
+    if scheme not in ("postgres", "postgresql"):
+        raise errors.ProgrammingError(f"invalid connection URL scheme {scheme!r}")
+    rest = rest.partition("#")[0]
+    rest, _, query = rest.partition("?")
+    authority, _, path = rest.partition("/")
+    userinfo, _, hostpart = authority.rpartition("@")
+    out: dict[str, str] = {}
+    if userinfo:
+        user, colon, password = userinfo.partition(":")
+        if user:
+            out["user"] = unquote(user)
+        if colon:
+            out["password"] = unquote(password)
+    hosts: list[str] = []
+    ports: list[str] = []
+    for item in hostpart.split(","):
+        if item.startswith("["):
+            end = item.index("]")
+            host, port = item[1:end], (item[end + 2:] if item[end + 1:end + 2] == ":" else "")
+        else:
+            host, _, port = item.partition(":")
+        hosts.append(unquote(host))
+        ports.append(port)
+    if any(hosts):
+        out["host"] = ",".join(hosts)
+    if any(ports):
+        out["port"] = ",".join(ports)
+    if path:
+        out["dbname"] = unquote(path)
+    for key, value in parse_qsl(query, keep_blank_values=True):
+        out[key] = value
+    return out
+
+
+def _parse_keywords(s: str) -> dict[str, str]:
+    out: dict[str, str] = {}
+    i, n = 0, len(s)
+    while i < n:
+        while i < n and s[i].isspace():
+            i += 1
+        if i >= n:
+            break
+        j = i
+        while j < n and s[j] not in "= \t\r\n":
+            j += 1
+        key = s[i:j]
+        while j < n and s[j].isspace():
+            j += 1
+        if j >= n or s[j] != "=":
+            raise errors.ProgrammingError(f"missing '=' after {key!r} in connection info string")
+        j += 1
+        while j < n and s[j].isspace():
+            j += 1
+        if j < n and s[j] == "'":
+            j += 1
+            buf: list[str] = []
+            while True:
+                if j >= n:
+                    raise errors.ProgrammingError("unterminated quoted string in connection info string")
+                c = s[j]
+                if c == "\\" and j + 1 < n:
+                    buf.append(s[j + 1])
+                    j += 2
+                elif c == "'":
+                    j += 1
+                    break
+                else:
+                    buf.append(c)
+                    j += 1
+            value = "".join(buf)
+        else:
+            k = j
+            while k < n and not s[k].isspace():
+                k += 1
+            value, j = s[j:k], k
+        out[key] = value
+        i = j
+    return out
+
+
+def parse_conninfo(conninfo: str) -> dict[str, str]:
+    conninfo = conninfo.strip()
+    if conninfo.startswith(("postgresql://", "postgres://")):
+        return _parse_url(conninfo)
+    return _parse_keywords(conninfo)
+
+
+def _default_host(port: int) -> str:
+    for directory in ("/var/run/postgresql", "/tmp"):
+        if os.path.exists(f"{directory}/.s.PGSQL.{port}"):
+            return directory
+    return "localhost"
+
+
+def build_params(conninfo: str = "", **kwargs: object) -> ConnParams:
+    values = parse_conninfo(conninfo) if conninfo else {}
+    for key, value in kwargs.items():
+        if value is not None:
+            values[key] = str(value)
+    if "database" in values:
+        values.setdefault("dbname", values.pop("database"))
+    unknown = set(values) - _KNOWN
+    if unknown:
+        raise errors.ProgrammingError(f"invalid connection option {sorted(unknown)[0]!r}")
+    for key, env in _ENV.items():
+        if not values.get(key) and os.environ.get(env):
+            values[key] = os.environ[env]
+
+    user = values.get("user") or getpass.getuser()
+    host_text = values.get("host", "")
+    port_text = values.get("port", "")
+    ports = [p for p in port_text.split(",")] if port_text else []
+    try:
+        port_numbers = [int(p) if p else 5432 for p in ports] or [5432]
+    except ValueError:
+        raise errors.ProgrammingError(f"invalid port {port_text!r}") from None
+    host_list = host_text.split(",") if host_text else [_default_host(port_numbers[0])]
+    hosts: list[tuple[str, int]] = []
+    for index, host in enumerate(host_list):
+        port = port_numbers[index] if len(port_numbers) > 1 else port_numbers[0]
+        hosts.append((host or _default_host(port), port))
+
+    sslmode = values.get("sslmode") or "prefer"
+    if sslmode not in SSL_MODES:
+        raise errors.ProgrammingError(f"invalid sslmode value {sslmode!r}")
+    timeout_text = values.get("connect_timeout")
+    try:
+        timeout = float(timeout_text) if timeout_text else None
+    except ValueError:
+        raise errors.ProgrammingError(f"invalid connect_timeout {timeout_text!r}") from None
+    return ConnParams(
+        hosts=hosts,
+        user=user,
+        password=values.get("password") or None,
+        dbname=values.get("dbname") or user,
+        sslmode=sslmode,
+        sslrootcert=values.get("sslrootcert") or None,
+        sslcert=values.get("sslcert") or None,
+        sslkey=values.get("sslkey") or None,
+        connect_timeout=timeout if timeout and timeout > 0 else None,
+        application_name=values.get("application_name") or None,
+        options=values.get("options") or None,
+    )