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

wrapture_instrumentation_postgresql/psycopg/cursor.py

@@ -0,0 +1,258 @@
+"""The cursor seams: the execute family, streamed queries and COPY,
+on the sync and async cursors and on the server-side cursors,
+recorded as database events.
+
+psycopg's cursors are pure Python, so each method is bound in place
+on its class. `Cursor.execute` and `Cursor.executemany` are the doors
+every ordinary query passes through, whichever cursor class the
+application chose: `ClientCursor` and `RawCursor` inherit them, and
+`Connection.execute`, the shortcut, builds a cursor and calls its
+`execute` in Python, so the one binding records it too with no
+double. The server-side cursors override `execute` to DECLARE a
+cursor rather than run the query, so those overrides are bound in
+their own right and record a DECLARE operation; their fetches, which
+FETCH from the portal, are not recorded, the model every database
+target here follows (a query event closes when its execute returns,
+and time spent iterating rows is the application's).
+
+`Cursor.stream()` is a generator: the query is sent on the first
+iteration and rows arrive one at a time. The binding's decorator is
+itself a generator over the driver's, so wrapture records the event
+around the iteration, its duration the time spent inside the
+generator and its item count the rows streamed.
+
+`Cursor.copy()` is a context manager factory: the COPY statement is
+sent on entering and the data flows through the yielded Copy object
+until exit. A plain binding would record the factory call, an
+instant, so the binding records nothing itself (`when=False`) and its
+decorator hands back a wrapping context manager that opens a block
+event on entry and closes it on exit, the block spanning the
+transfer, with the row count annotated at the end.
+
+Every event carries the database contract keys from the common
+module; the SQL text rides as `statement` only when the setting is
+on. Bound parameters are never recorded.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterator
+from types import TracebackType
+from typing import Any
+
+import wrapture
+
+from ..common import SYSTEM, captured, query_data, server_of
+
+
+class RecordedCopy:
+    """The context manager handed back in place of psycopg's copy
+    factory result: a block event around the real one."""
+
+    def __init__(
+        self, label: str, cursor: Any, inner: Any, data: dict[str, Any]
+    ) -> None:
+        self._cursor = cursor
+        self._inner = inner
+        self._block = wrapture.block(label, category="database", data=data, leaf=True)
+
+    def _finish(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        # The cursor's rowcount is fresh once the inner exit has run
+        # the rest of the factory's body, which reads the COPY result.
+
+        rows = getattr(self._cursor, "rowcount", -1)
+        if isinstance(rows, int) and rows >= 0:
+            wrapture.annotate(rows=rows)
+
+        self._block.__exit__(exc_type, exc_value, traceback)
+
+    def __enter__(self) -> Any:
+        self._block.__enter__()
+        try:
+            return self._inner.__enter__()
+        except BaseException as error:
+            self._block.__exit__(type(error), error, error.__traceback__)
+            raise
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> Any:
+        try:
+            outcome = self._inner.__exit__(exc_type, exc_value, traceback)
+        except BaseException as error:
+            self._finish(type(error), error, error.__traceback__)
+            raise
+
+        self._finish(exc_type, exc_value, traceback)
+
+        return outcome
+
+    async def __aenter__(self) -> Any:
+        self._block.__enter__()
+        try:
+            return await self._inner.__aenter__()
+        except BaseException as error:
+            self._block.__exit__(type(error), error, error.__traceback__)
+            raise
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> Any:
+        try:
+            outcome = await self._inner.__aexit__(exc_type, exc_value, traceback)
+        except BaseException as error:
+            self._finish(type(error), error, error.__traceback__)
+            raise
+
+        self._finish(exc_type, exc_value, traceback)
+
+        return outcome
+
+
+def instrument(module: Any, instrumentation: wrapture.Instrumentation) -> None:
+    """Bind the execute family, stream and copy on the cursor classes;
+    register their removal as this trigger's cleanup."""
+
+    settings = instrumentation.settings
+    record_statement = bool(settings["statement"])
+
+    def query_of(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        return args[0] if args else kwargs.get("query", kwargs.get("statement"))
+
+    def data_for(
+        instance: Any, query: Any, operation: str | None = None
+    ) -> dict[str, Any]:
+        return query_data(
+            query, instance, instance.connection.info, record_statement, operation
+        )
+
+    def executes(operation: str | None = None) -> Any:
+        def record(
+            wrapped: Any,
+            instance: Any,
+            args: tuple[Any, ...],
+            kwargs: dict[str, Any],
+        ) -> Any:
+            wrapture.annotate(**data_for(instance, query_of(args, kwargs), operation))
+
+            return wrapped(*args, **kwargs)
+
+        return record
+
+    def executes_async(operation: str | None = None) -> Any:
+        async def record(
+            wrapped: Any,
+            instance: Any,
+            args: tuple[Any, ...],
+            kwargs: dict[str, Any],
+        ) -> Any:
+            wrapture.annotate(**data_for(instance, query_of(args, kwargs), operation))
+
+            return await wrapped(*args, **kwargs)
+
+        return record
+
+    # The stream decorators are generators over the driver's, so the
+    # annotation lands inside the event wrapture records around the
+    # iteration, not before it exists.
+
+    def streams(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Iterator[Any]:
+        wrapture.annotate(**data_for(instance, query_of(args, kwargs)))
+
+        yield from wrapped(*args, **kwargs)
+
+    async def streams_async(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> AsyncIterator[Any]:
+        wrapture.annotate(**data_for(instance, query_of(args, kwargs)))
+
+        async for row in wrapped(*args, **kwargs):
+            yield row
+
+    def copies(label: str) -> Any:
+        def record(
+            wrapped: Any,
+            instance: Any,
+            args: tuple[Any, ...],
+            kwargs: dict[str, Any],
+        ) -> RecordedCopy:
+            data = data_for(instance, query_of(args, kwargs), "COPY")
+
+            return RecordedCopy(label, instance, wrapped(*args, **kwargs), data)
+
+        return record
+
+    def query_binding(owner: Any, name: str) -> wrapture.Binding:
+        return wrapture.binding(
+            owner,
+            name,
+            category="database",
+            leaf=True,
+            capture_args=captured,
+            capture_result=captured,
+        )
+
+    named: dict[str, wrapture.Binding] = {}
+
+    # The execute family on the ordinary cursors, sync and async.
+
+    for prefix, owner, decorator in (
+        ("cursor", module.Cursor, executes),
+        ("async_cursor", module.AsyncCursor, executes_async),
+    ):
+        for method in ("execute", "executemany"):
+            bound = query_binding(owner, method)
+            bound.on_call.decorates(decorator())
+            named[f"{prefix}_{method}"] = bound
+
+    # The server-side cursors' own execute, a DECLARE.
+
+    for prefix, owner, decorator in (
+        ("server_cursor", module.ServerCursor, executes),
+        ("async_server_cursor", module.AsyncServerCursor, executes_async),
+    ):
+        bound = query_binding(owner, "execute")
+        bound.on_call.decorates(decorator("DECLARE"))
+        named[f"{prefix}_execute"] = bound
+
+    # Streamed queries, recorded around the iteration.
+
+    stream = query_binding(module.Cursor, "stream")
+    stream.on_call.decorates(streams)
+    named["cursor_stream"] = stream
+
+    async_stream = query_binding(module.AsyncCursor, "stream")
+    async_stream.on_call.decorates(streams_async)
+    named["async_cursor_stream"] = async_stream
+
+    # COPY: the factory records nothing itself, the wrapping context
+    # manager its decorator returns records the block.
+
+    for prefix, owner in (
+        ("cursor", module.Cursor),
+        ("async_cursor", module.AsyncCursor),
+    ):
+        bound = wrapture.binding(owner, "copy", when=False)
+        bound.on_call.decorates(copies(f"psycopg:{owner.__name__}.copy"))
+        named[f"{prefix}_copy"] = bound
+
+    group = wrapture.bindings(**named)
+    group.apply()
+
+    instrumentation.on_cleanup(group.remove)
+
+
+__all__ = ["RecordedCopy", "SYSTEM", "instrument", "server_of"]

wrapture_instrumentation_postgresql/psycopg/transaction.py

@@ -0,0 +1,126 @@
+"""The `transaction()` block seams: entering and leaving a
+Transaction or AsyncTransaction, recorded as database events.
+
+`with conn.transaction():` is psycopg's transaction block. Entering
+issues BEGIN when no transaction is open on the connection, and a
+SAVEPOINT for a nested block (or when a savepoint name was asked
+for); leaving issues COMMIT or RELEASE SAVEPOINT, or ROLLBACK (to the
+savepoint, for a nested block) when an exception is passing through,
+`force_rollback` was set, or `psycopg.Rollback` was raised inside.
+The block object knows which it did: whether it opened the outermost
+transaction, and its savepoint name, both read after the enter has
+run and before the exit does. These statements go to libpq directly,
+never through a cursor, so the enter and exit are bound in their own
+right, each recording the operation it performed and the savepoint
+name when one was involved.
+"""
+
+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 the enter and exit of the sync and async transaction
+    blocks; register their removal as this trigger's cleanup."""
+
+    def entered(instance: Any) -> dict[str, Any]:
+        # Known only once the enter has run: whether this block began
+        # the transaction or nested inside one.
+
+        outer = bool(getattr(instance, "_outer_transaction", False))
+        savepoint = getattr(instance, "_savepoint_name", "") or None
+
+        data: dict[str, Any] = {"operation": "BEGIN" if outer else "SAVEPOINT"}
+        if savepoint:
+            data["savepoint"] = savepoint
+
+        return data
+
+    def leaving(
+        instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> dict[str, Any]:
+        exc_value = args[1] if len(args) > 1 else kwargs.get("exc_val")
+        outer = bool(getattr(instance, "_outer_transaction", False))
+        savepoint = getattr(instance, "_savepoint_name", "") or None
+        commits = exc_value is None and not getattr(instance, "force_rollback", False)
+
+        if commits:
+            operation = "COMMIT" if outer else "RELEASE"
+        else:
+            operation = "ROLLBACK"
+
+        data: dict[str, Any] = {"operation": operation}
+        if savepoint:
+            data["savepoint"] = savepoint
+
+        return data
+
+    def enters(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        wrapture.annotate(system=SYSTEM, **server_of(instance.connection.info))
+
+        outcome = wrapped(*args, **kwargs)
+        wrapture.annotate(**entered(instance))
+
+        return outcome
+
+    async def enters_async(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        wrapture.annotate(system=SYSTEM, **server_of(instance.connection.info))
+
+        outcome = await wrapped(*args, **kwargs)
+        wrapture.annotate(**entered(instance))
+
+        return outcome
+
+    def exits(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        wrapture.annotate(
+            system=SYSTEM,
+            **server_of(instance.connection.info),
+            **leaving(instance, args, kwargs),
+        )
+
+        return wrapped(*args, **kwargs)
+
+    async def exits_async(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        wrapture.annotate(
+            system=SYSTEM,
+            **server_of(instance.connection.info),
+            **leaving(instance, args, kwargs),
+        )
+
+        return await wrapped(*args, **kwargs)
+
+    def boundary(owner: Any, name: str, decorator: Any) -> wrapture.Binding:
+        binding = wrapture.binding(
+            owner,
+            name,
+            category="database",
+            leaf=True,
+            capture_args=captured,
+            capture_result=captured,
+        )
+        binding.on_call.decorates(decorator)
+
+        return binding
+
+    group = wrapture.bindings(
+        enter=boundary(module.Transaction, "__enter__", enters),
+        exit=boundary(module.Transaction, "__exit__", exits),
+        async_enter=boundary(module.AsyncTransaction, "__aenter__", enters_async),
+        async_exit=boundary(module.AsyncTransaction, "__aexit__", exits_async),
+    )
+    group.apply()
+
+    instrumentation.on_cleanup(group.remove)

wrapture_instrumentation_postgresql/psycopg2/README.md

@@ -0,0 +1,131 @@
+# psycopg2 instrumentation
+
+Query and transaction tracing for
+[psycopg2](https://www.psycopg.org/docs/), the long-standing
+PostgreSQL adapter. Entry point name `psycopg2`, the package it
+patches (psycopg2-binary installs the same package); supports
+psycopg2 2.9 and later, below 3; fully removable.
+
+## Enabling it
+
+An `[[instrument]]` entry in `wrapture.toml` (with at least one sink
+to hear the events):
+
+```toml
+[[instrument]]
+name = "psycopg2"
+
+[[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("psycopg2"):
+    ...
+```
+
+## What you see
+
+One `database` leaf per operation: the connection being opened, each
+query however it was issued (`execute`, `executemany`, `callproc`,
+and the extras' batch helpers above them), each COPY (`copy_from`,
+`copy_to`, `copy_expert`), and each transaction boundary (`commit`,
+`rollback`, and the connection's commit-or-rollback context manager,
+whose exit records which of the two it performed):
+
+```
+psycopg2:connect()  -> '<connection>'
+psycopg2.extensions:cursor.execute(query='<31 chars>', vars='<1 values>')
+psycopg2.extensions:connection.commit()
+```
+
+- `psycopg2.extensions.connection` and `.cursor` are C types no
+  patch can touch, and psycopg2's own C entry points (`register_type`
+  and everything built on it, `quote_ident`, `Json.prepare`,
+  `lobject`) type-check the objects handed to them, so a proxy is
+  ruled out too. The instrumentation instead uses the mechanism
+  psycopg2 provides for its own extras: it binds `psycopg2.connect`
+  to substitute the requested connection class with a recording
+  subclass of it, and that subclass hands out cursors that are
+  recording subclasses of whatever cursor class was asked for. Your
+  own factories keep working and are simply recorded: a
+  `cursor_factory` named at connect or per cursor (`DictCursor`,
+  `RealDictCursor`, `NamedTupleCursor`), a `connection_factory` with
+  a cursor default of its own (`RealDictConnection`,
+  `LoggingConnection`), all of them real subclasses, so isinstance
+  checks and the C type checks pass and reprs read as before.
+
+- Every event carries the database contract keys `system`
+  (`postgresql`) and `operation` (the SQL's leading keyword, or
+  `CONNECT`, `COMMIT`, `ROLLBACK`, `CALL`, `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. A `callproc`
+  names the procedure as `procedure`; a `copy_from` or `copy_to`
+  names the table as `collection`.
+
+- The connection's context manager (`with conn:`) commits, or rolls
+  back when an exception is on its way through, and does not close
+  the connection; its exit records which it did. A cursor's context
+  manager only closes the cursor and records nothing.
+
+- 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 sequence the driver has yet to
+  consume, which is never iterated); the SQL text reduces to its
+  length unless the `statement` setting is on; a COPY's file reduces
+  to its type; and the connect event captures none of its arguments,
+  which carry the password. psycopg2 interpolates parameters
+  client-side, below the seam, so the recorded text is always the
+  template with its `%s` placeholders, never the query as sent.
+
+- The extras' batch helpers (`execute_values`, `execute_batch`) call
+  `execute` once per page or batch, and record one event each, which
+  is what happens on the wire. A failing statement records the
+  driver's exception (`psycopg2.errors.UndefinedTable`, say) on the
+  event as it escapes. Fetching is not recorded: a query event closes
+  when its execute returns, and a named (server-side) cursor's
+  FETCHes go unrecorded for the same reason, its execute being the
+  DECLARE.
+
+- An asynchronous connection (`async_=True`) returns from `execute`
+  before the query runs and the application polls; the event then
+  measures the send only.
+
+## 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 = "psycopg2"
+statement = true
+```
+
+## With the sqlalchemy instrumentation
+
+An instrumented psycopg2 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: `cursor.execute` under
+`do_execute`, `psycopg2:connect` under the dialect's `connect`, and
+the driver's `executemany` (or the `execute` calls of the dialect's
+`execute_values` batching, for compiled inserts) under the psycopg2
+dialect's own `do_executemany` override, which the sqlalchemy target
+binds in its own right.
+Raw psycopg2 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 settings the dialect reads when it opens a
+connection, and the pool's reset-on-return rollback.
+
+## How it patches
+
+For the implementation detail see the module docstring of
+[factories.py](factories.py).

wrapture_instrumentation_postgresql/psycopg2/__init__.py

@@ -0,0 +1,49 @@
+"""Instrumentation for psycopg2: every query, the connections opened
+and the transaction boundaries recorded as database events, through
+recording subclasses of the driver's connection and cursor types
+injected by way of the factory hooks psycopg2 itself provides.
+
+This module imports only wrapture. Everything that touches psycopg2
+lives in the sibling factories module, importing only wrapture at
+top level and reaching psycopg2 through the package the hook is
+handed or a lazy import inside a method, so loading this class when a
+config loads never imports psycopg2 ahead of the hook meant to fire
+on its import.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import wrapture
+from wrapture import Setting
+
+from . import factories
+
+
+class Psycopg2Instrumentation(wrapture.Instrumentation):
+    """Query and transaction tracing for psycopg2."""
+
+    description = "Query and transaction tracing for psycopg2."
+
+    target = "psycopg2"
+    supports = ">=2.9,<3"
+    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("psycopg2")
+    def psycopg2(self, name: str, module: Any) -> None:
+        """Bind the connect factory and the recording classes' methods
+        once psycopg2 exists."""
+
+        factories.instrument(module, self)