wrapture-instrumentation-postgresql 1.0.0.dev1__py3-none-any.whl

wrapture_instrumentation_postgresql/common.py

@@ -0,0 +1,169 @@
+"""What every PostgreSQL driver's events have in common: the database
+category's contract keys, the operation name derived from the SQL,
+and the capture policy that keeps queries and their data out of the
+record.
+
+This module imports only wrapture, so a target subpackage can import
+it at load time without dragging any driver in.
+
+Every event carries `system` ("postgresql") and `operation` (the SQL's
+leading keyword, or CONNECT, COMMIT, ROLLBACK and their kin), plus
+`database`, `host` and `port` read from the driver's connection info,
+so each event says which server it went to. The SQL text is recorded
+as `statement` only when the target's `statement` setting is on, as
+the application handed it to the driver and never with its bound
+parameters, which no setting captures.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+SYSTEM = "postgresql"
+
+# The argument names under which the drivers take SQL text, and those
+# under which they take its parameters; the capture policy reduces
+# the first to a length and the second to a count.
+
+_QUERY_NAMES = frozenset({"query", "sql", "statement", "command"})
+_PARAMETER_NAMES = frozenset(
+    {"params", "params_seq", "vars", "vars_list", "parameters", "args", "records"}
+)
+
+# File-like arguments (a COPY's source or sink) reduce to their type:
+# their repr can name a path, and their contents are the data.
+
+_FILE_NAMES = frozenset({"file", "source", "output"})
+
+
+def operation_of(sql: str) -> str:
+    """The SQL's leading keyword, uppercased: the low-cardinality
+    operation name the database contract carries."""
+
+    head = sql.split(None, 1)
+
+    # A statement may end in its keyword ("COMMIT;"): the terminator is
+    # not part of the operation.
+
+    return head[0].upper().rstrip(";") if head else "?"
+
+
+def statement_of(query: Any, context: Any = None) -> str | None:
+    """The SQL text of a query as the driver was handed it: a string
+    as is, bytes decoded, a composed query (psycopg's `sql.SQL` and
+    kin) rendered against the cursor or connection it will run on,
+    anything else None."""
+
+    if isinstance(query, str):
+        return query
+
+    if isinstance(query, (bytes, bytearray, memoryview)):
+        try:
+            return bytes(query).decode()
+        except UnicodeDecodeError:
+            return None
+
+    render = getattr(query, "as_string", None)
+    if callable(render):
+        try:
+            rendered = render(context)
+        except Exception:
+            return None
+        return rendered if isinstance(rendered, str) else None
+
+    return None
+
+
+def server_of(info: Any) -> dict[str, Any]:
+    """The `database`, `host` and `port` keys from a driver's
+    connection info object (psycopg's and psycopg2's both expose
+    `dbname`, `host` and `port`), whichever of them it can supply."""
+
+    data: dict[str, Any] = {}
+
+    for key, attribute in (("database", "dbname"), ("host", "host"), ("port", "port")):
+        try:
+            value = getattr(info, attribute)
+        except Exception:
+            continue
+
+        if value not in (None, ""):
+            data[key] = value
+
+    return data
+
+
+def statement_data(
+    query: Any,
+    context: Any,
+    server: dict[str, Any],
+    record_statement: bool,
+    operation: str | None = None,
+) -> dict[str, Any]:
+    """The data for one query event: the contract keys, the server
+    keys given, and the statement text when the setting asks for it.
+    The operation is the one given, else the query's leading keyword."""
+
+    text = statement_of(query, context)
+
+    data: dict[str, Any] = {"system": SYSTEM}
+
+    if operation is not None:
+        data["operation"] = operation
+    elif text is not None:
+        data["operation"] = operation_of(text)
+
+    data.update(server)
+
+    if record_statement and text is not None:
+        data["statement"] = text
+
+    return data
+
+
+def query_data(
+    query: Any,
+    context: Any,
+    info: Any,
+    record_statement: bool,
+    operation: str | None = None,
+) -> dict[str, Any]:
+    """statement_data() with the server keys read from a driver's
+    connection info object (psycopg's and psycopg2's)."""
+
+    return statement_data(query, context, server_of(info), record_statement, operation)
+
+
+def captured(name: str | None, value: Any) -> Any:
+    """SQL text reduces to its length, parameters to a count or their
+    type (a parameter sequence may be a generator the driver has yet
+    to consume, and is never iterated), a COPY's file to its type, a
+    context manager exit's exception value and traceback to their
+    types, and every unnamed value to its type: the query and its
+    data never reach the record through argument capture."""
+
+    if name in _QUERY_NAMES:
+        text = statement_of(value)
+        if text is not None:
+            return f"<{len(text)} chars>"
+        return f"<{type(value).__name__}>"
+
+    if name in _PARAMETER_NAMES:
+        if isinstance(value, (list, tuple, dict)):
+            return f"<{len(value)} values>"
+        return f"<{type(value).__name__}>"
+
+    if name in _FILE_NAMES:
+        return f"<{type(value).__name__}>"
+
+    # An exception's message is application data like any other; the
+    # exit event's exception, when one escapes, is recorded properly
+    # on the event itself.
+
+    if name in ("exc_value", "exc_val", "traceback", "exc_tb") and value is not None:
+        return f"<{type(value).__name__}>"
+
+    if name is None:
+        return f"<{type(value).__name__}>"
+
+    return value

wrapture_instrumentation_postgresql/psycopg/README.md

@@ -0,0 +1,144 @@
+# psycopg instrumentation
+
+Query and transaction tracing for
+[psycopg](https://www.psycopg.org/psycopg3/) (version 3), the current
+PostgreSQL adapter. Entry point name `psycopg`, the package it
+patches; supports psycopg 3.1 and later, below 4; fully removable.
+Sync and async classes alike, whichever of the pure Python, C or
+binary implementations is installed.
+
+## Enabling it
+
+An `[[instrument]]` entry in `wrapture.toml` (with at least one sink
+to hear the events):
+
+```toml
+[[instrument]]
+name = "psycopg"
+
+[[sink]]
+type = "printer"
+```
+
+run under wrapture's runner (`python -m wrapture -m myapp`), or in a
+test through the context manager:
+
+```python
+with wrapture.instrumentation("psycopg"):
+    ...
+```
+
+## What you see
+
+One `database` leaf per operation: the connection being opened, each
+query however it was issued (a cursor's `execute` or `executemany`,
+or the connection's `execute` shortcut), each streamed query, each
+COPY, and each transaction boundary (`commit`, `rollback`, the
+connection's commit-or-rollback context manager, and a
+`transaction()` block's begin and end):
+
+```
+psycopg:connect()  -> '<Connection>'
+psycopg:Cursor.execute(query='<29 chars>', params='<1 values>', prepare=None, binary=None)  -> '<Cursor>'
+psycopg:Transaction.__enter__()  -> '<Transaction>'
+psycopg:Transaction.__exit__(exc_type=None, exc_val=None, exc_tb=None)  -> '<bool>'
+psycopg:Connection.commit()
+```
+
+- psycopg's classes are pure Python, so the instrumentation binds
+  their methods in place: `Cursor.execute` and `executemany` (which
+  every cursor class inherits, `ClientCursor` and `RawCursor`
+  included, and which the connection's `execute` shortcut calls, so
+  it records once), `Cursor.stream`, `Cursor.copy`, the server-side
+  cursors' own `execute` (a DECLARE), `Connection.connect` (and the
+  `psycopg.connect` spelling of it), `commit`, `rollback`, the
+  connection's context manager exit, and the `transaction()` block's
+  enter and exit. The async classes are bound the same way and
+  record around the await. Connections from a `psycopg_pool` pool
+  record like any other, since the bindings sit on the classes.
+
+- Every event carries the database contract keys `system`
+  (`postgresql`) and `operation` (the SQL's leading keyword, or
+  `CONNECT`, `COMMIT`, `ROLLBACK`, `BEGIN`, `SAVEPOINT`, `RELEASE`,
+  `DECLARE`, `COPY`), which wrapture's OpenTelemetry export maps to
+  `db.system.name` and `db.operation.name`, plus the `database`,
+  `host` and `port` the connection reached, from the connection's
+  own info, so every span says which server it went to.
+
+- A `transaction()` block records what it did: `BEGIN` on entering
+  when nothing was open on the connection, `SAVEPOINT` (with the
+  savepoint name) for a nested block, and on leaving `COMMIT` or
+  `RELEASE`, or `ROLLBACK` when an exception passed through,
+  `force_rollback` was set or `psycopg.Rollback` was raised inside.
+  The connection's own context manager records its `COMMIT` or
+  `ROLLBACK` the same way. The `BEGIN` psycopg sends implicitly
+  before the first statement of a transaction goes with that
+  statement and folds into its event.
+
+- A streamed query (`cursor.stream()`) records one event around the
+  whole iteration: its duration is the time spent inside the
+  generator and `items` the rows streamed; the event closes when the
+  iteration ends or the generator is abandoned.
+
+- A COPY (`with cursor.copy(...) as copy:`) records one block event
+  spanning the transfer, from entering the block to leaving it,
+  labelled `psycopg:Cursor.copy` (or `AsyncCursor.copy`), with the
+  operation `COPY` and the rows copied as `rows`.
+
+- The capture policy is deliberate about sensitive data: bound
+  parameters are never recorded, under any setting (they reduce to a
+  count, or to a type name for a parameter sequence the driver has
+  yet to consume, which is never iterated); the SQL text reduces to
+  its length unless the `statement` setting is on; and the connect
+  event captures none of its arguments, which carry the password.
+  There is no obfuscation at this layer, because rewriting SQL to
+  strip literals is a losing game outside a real lexer: record the
+  text where queries are parameterized, leave it off where they are
+  not.
+
+- A failing statement records the driver's exception
+  (`psycopg.errors.UndefinedTable`, say) on the event as it escapes.
+  Fetching is not recorded: a query event closes when its execute
+  returns, so time spent iterating rows afterwards is not attributed
+  to the database, and a server-side cursor's FETCHes from its portal
+  go unrecorded for the same reason, its DECLARE being the event.
+
+- In pipeline mode (`with conn.pipeline():`) an execute returns as
+  soon as the query is queued and the results arrive at the sync
+  point, so the statement events are short and the wait sits in the
+  pipeline's exit; every statement still records.
+
+## Settings
+
+| Setting | Default | Controls |
+| ------- | ------- | -------- |
+| `statement` | `false` | Whether each query event records the SQL text as handed to the driver, as `statement` (a composed `sql.SQL(...)` query rendered as it will be sent). Off by default because the driver cannot tell a literal an application interpolated from a placeholder; turn it on when your queries are parameterized, the text then carrying placeholders rather than data. A `sql.Literal` composed into a query is recorded as written, so prefer placeholders there too. Bound parameters are never recorded either way. |
+
+```toml
+[[instrument]]
+name = "psycopg"
+statement = true
+```
+
+## With the sqlalchemy instrumentation
+
+An instrumented psycopg beneath the core package's `sqlalchemy`
+target composes through that target's `leaf` setting. With the
+default `leaf = true` each statement is one event and the driver's
+own events stay out of the tree; with `leaf = false` the driver's
+events nest beneath each statement, `psycopg:Cursor.execute` under
+`do_execute`, `psycopg:connect` under the dialect's `connect`, the
+async engine's `AsyncCursor.execute` likewise. Raw psycopg use beside
+the engine records at the top level either way. A little of the
+dialect's own housekeeping also shows up regardless, because it runs
+straight against the driver outside the recorded seams: the type
+lookups the psycopg dialect makes when it opens a connection, and
+the pool's reset-on-return rollback.
+
+## How it patches
+
+For the implementation detail see the module docstrings of
+[cursor.py](cursor.py) (the execute family, stream and COPY),
+[connection.py](connection.py) (connect and the connection's own
+boundaries) and [transaction.py](transaction.py) (the `transaction()`
+block).

wrapture_instrumentation_postgresql/psycopg/__init__.py

@@ -0,0 +1,57 @@
+"""Instrumentation for psycopg (version 3): every query, the
+connections opened and the transaction boundaries recorded as
+database events, by bindings on the pure Python classes the driver
+is made of.
+
+This module imports only wrapture. Everything that touches psycopg
+lives in the sibling modules, one per kind of seam (cursor.py for the
+execute family, stream and COPY; connection.py for connect, commit,
+rollback and the connection's context manager; transaction.py for
+`transaction()` blocks), each importing only wrapture at top level
+and reaching psycopg through the package the hook is handed, so
+loading this class when a config loads never imports psycopg ahead of
+the hook meant to fire on its import.
+
+One trigger suffices: importing psycopg initialises every submodule
+the seams live in, so by the time the hook fires all the classes
+exist under the package.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import wrapture
+from wrapture import Setting
+
+from . import connection, cursor, transaction
+
+
+class PsycopgInstrumentation(wrapture.Instrumentation):
+    """Query and transaction tracing for psycopg (version 3)."""
+
+    description = "Query and transaction tracing for psycopg (version 3)."
+
+    target = "psycopg"
+    supports = ">=3.1,<4"
+    removable = True
+
+    settings = {
+        "statement": Setting(
+            False,
+            "record the SQL text as handed to the driver on each query"
+            " event; off by default because the driver cannot tell a"
+            " literal an application interpolated from a placeholder,"
+            " and the text is only safe to record when queries are"
+            " parameterized",
+        ),
+    }
+
+    @wrapture.instrumentation_hook("psycopg")
+    def psycopg(self, name: str, module: Any) -> None:
+        """Bind the cursor, connection and transaction seams once
+        psycopg exists."""
+
+        cursor.instrument(module, self)
+        connection.instrument(module, self)
+        transaction.instrument(module, self)

wrapture_instrumentation_postgresql/psycopg/connection.py

@@ -0,0 +1,168 @@
+"""The connection seams: connections being opened, and the
+transaction boundaries the connection itself performs, recorded as
+database events.
+
+`Connection.connect` and `AsyncConnection.connect` are the
+classmethods every connection comes from: the module-level
+`psycopg.connect` is the same classmethod under another name, bound
+at import, so it is bound as a module attribute too, as sqlite3's two
+spellings of `connect` are; a pool (psycopg_pool) calls the class's
+method, so pooled connections record through it. The connect event's
+arguments are never captured (the conninfo carries a password); the
+server it reached is annotated from the connection's info afterwards.
+
+`commit()` and `rollback()` go to libpq directly rather than through
+a cursor, so they are bound in their own right. The connection's
+context manager exit commits, or rolls back when an exception is on
+its way through, then closes the connection unless it belongs to a
+pool; the exit is bound and records which of the two it performed.
+The BEGIN psycopg issues implicitly before the first statement of a
+transaction is sent inside that statement's execute, so it folds into
+that event.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import wrapture
+
+from ..common import SYSTEM, captured, server_of
+
+
+def instrument(module: Any, instrumentation: wrapture.Instrumentation) -> None:
+    """Bind connect, commit, rollback and the context manager exit on
+    the sync and async connection classes; register their removal as
+    this trigger's cleanup."""
+
+    def opens(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        wrapture.annotate(system=SYSTEM, operation="CONNECT")
+
+        connection = wrapped(*args, **kwargs)
+        wrapture.annotate(**server_of(connection.info))
+
+        return connection
+
+    async def opens_async(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        wrapture.annotate(system=SYSTEM, operation="CONNECT")
+
+        connection = await wrapped(*args, **kwargs)
+        wrapture.annotate(**server_of(connection.info))
+
+        return connection
+
+    def performs(operation: str) -> Any:
+        def record(
+            wrapped: Any,
+            instance: Any,
+            args: tuple[Any, ...],
+            kwargs: dict[str, Any],
+        ) -> Any:
+            wrapture.annotate(
+                system=SYSTEM, operation=operation, **server_of(instance.info)
+            )
+
+            return wrapped(*args, **kwargs)
+
+        return record
+
+    def performs_async(operation: str) -> Any:
+        async def record(
+            wrapped: Any,
+            instance: Any,
+            args: tuple[Any, ...],
+            kwargs: dict[str, Any],
+        ) -> Any:
+            wrapture.annotate(
+                system=SYSTEM, operation=operation, **server_of(instance.info)
+            )
+
+            return await wrapped(*args, **kwargs)
+
+        return record
+
+    def exit_operation(args: tuple[Any, ...], kwargs: dict[str, Any]) -> str:
+        # The exit commits unless an exception is on its way through.
+
+        exc_type = args[0] if args else kwargs.get("exc_type")
+
+        return "COMMIT" if exc_type is None else "ROLLBACK"
+
+    def leaves(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        wrapture.annotate(
+            system=SYSTEM,
+            operation=exit_operation(args, kwargs),
+            **server_of(instance.info),
+        )
+
+        return wrapped(*args, **kwargs)
+
+    async def leaves_async(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        wrapture.annotate(
+            system=SYSTEM,
+            operation=exit_operation(args, kwargs),
+            **server_of(instance.info),
+        )
+
+        return await wrapped(*args, **kwargs)
+
+    def database_binding(
+        owner: Any, name: str, capture_args: Any = captured
+    ) -> wrapture.Binding:
+        return wrapture.binding(
+            owner,
+            name,
+            category="database",
+            leaf=True,
+            capture_args=capture_args,
+            capture_result=captured,
+        )
+
+    named: dict[str, wrapture.Binding] = {}
+
+    # The connect classmethods, and the module-level spelling of the
+    # sync one; none captures its arguments.
+
+    connect = database_binding(module.Connection, "connect", "none")
+    connect.on_call.decorates(opens)
+    named["connect"] = connect
+
+    async_connect = database_binding(module.AsyncConnection, "connect", "none")
+    async_connect.on_call.decorates(opens_async)
+    named["async_connect"] = async_connect
+
+    module_connect = database_binding(module, "connect", "none")
+    module_connect.on_call.decorates(opens)
+    named["module_connect"] = module_connect
+
+    # The transaction boundaries the connection performs itself.
+
+    for method, operation in (("commit", "COMMIT"), ("rollback", "ROLLBACK")):
+        bound = database_binding(module.Connection, method)
+        bound.on_call.decorates(performs(operation))
+        named[method] = bound
+
+        bound = database_binding(module.AsyncConnection, method)
+        bound.on_call.decorates(performs_async(operation))
+        named[f"async_{method}"] = bound
+
+    closes = database_binding(module.Connection, "__exit__")
+    closes.on_call.decorates(leaves)
+    named["exit"] = closes
+
+    async_closes = database_binding(module.AsyncConnection, "__aexit__")
+    async_closes.on_call.decorates(leaves_async)
+    named["async_exit"] = async_closes
+
+    group = wrapture.bindings(**named)
+    group.apply()
+
+    instrumentation.on_cleanup(group.remove)