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

wrapture_instrumentation_postgresql/__init__.py

@@ -0,0 +1,26 @@
+"""
+Instrumentation for the PostgreSQL client libraries, applied through
+wrapture.
+
+Each target lives in its own subpackage here (psycopg and psycopg2
+now, asyncpg later), holding one wrapture.Instrumentation subclass
+registered in the wrapture.instrumentation entry point group under the
+bare target name. This module carries only the version: importing it
+loads no instrumentation and no target.
+"""
+
+
+def _format_version(parts: tuple[str, ...]) -> str:
+    base = ".".join(parts[:3])
+
+    if len(parts) == 3:
+        return base
+
+    suffix = parts[3]
+    return (
+        f"{base}.{suffix}" if suffix.startswith(("dev", "post")) else f"{base}{suffix}"
+    )
+
+
+__version_info__ = ("1", "0", "0", "dev1")
+__version__ = _format_version(__version_info__)

wrapture_instrumentation_postgresql/asyncpg/README.md

@@ -0,0 +1,136 @@
+# asyncpg instrumentation
+
+Query and transaction tracing for
+[asyncpg](https://magicstack.github.io/asyncpg/), the asyncio
+PostgreSQL client. Entry point name `asyncpg`, the package it
+patches; supports asyncpg 0.29 and later, below 1.0; fully
+removable.
+
+## Enabling it
+
+An `[[instrument]]` entry in `wrapture.toml` (with at least one sink
+to hear the events):
+
+```toml
+[[instrument]]
+name = "asyncpg"
+
+[[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("asyncpg"):
+    ...
+```
+
+## What you see
+
+One `database` leaf per operation, recorded around its await: the
+connection being opened, each query through a connection's `execute`,
+`executemany`, `fetch`, `fetchrow`, `fetchval` or `fetchmany`, each
+of a prepared statement's own fetches, each round trip a server-side
+cursor makes, each COPY, and each transaction boundary:
+
+```
+asyncpg:connect()  -> '<Connection>'
+asyncpg.connection:Connection.execute(query='<6 chars>', args='<0 values>', timeout=None)  -> '<str>'
+asyncpg.connection:Connection.fetch(query='<43 chars>', args='<1 values>', timeout=None, record_class=None)  -> '<list>'
+asyncpg.prepared_stmt:PreparedStatement.fetchval(args='<1 values>', column=0, timeout=None)  -> '<int>'
+```
+
+- asyncpg's public surface is pure Python coroutine methods, so the
+  instrumentation binds them in place on the classes and records
+  each around its await; the Cython protocol beneath does the socket
+  work and is where the seams stop. There is no callback-style
+  completion to make spans overlap: one call, one span, closed when
+  the await returns or raises. Two tasks querying at once give
+  sibling spans that overlap in wall time, which is what they are.
+
+- Every event carries the database contract keys `system`
+  (`postgresql`) and `operation` (the SQL's leading keyword, or
+  `CONNECT`, `DECLARE`, `FETCH`, `COPY`, `EXPLAIN`), which
+  wrapture's OpenTelemetry export maps to `db.system.name` and
+  `db.operation.name`, plus the `database`, `host` and `port` the
+  connection reached. A COPY to or from a table names it as
+  `collection`.
+
+- Transactions need nothing of their own: `async with
+  conn.transaction():` issues its BEGIN, COMMIT, ROLLBACK, SAVEPOINT,
+  RELEASE SAVEPOINT and ROLLBACK TO through `Connection.execute`, so
+  each boundary is an ordinary execute event with its own operation.
+
+- A prepared statement (`await conn.prepare(query)`) records each of
+  its `fetch`, `fetchrow`, `fetchval`, `fetchmany`, `executemany` and
+  `explain` calls, with the statement's SQL; the `prepare` itself is
+  not a query and is not recorded.
+
+- A server-side cursor (`conn.cursor(query)` inside a transaction,
+  iterated with `async for` or driven by `fetch`, `fetchrow` and
+  `forward`) never passes through the connection's public methods,
+  so its round trips are recorded where they happen: one `DECLARE`
+  event when the cursor is bound (with the first batch of rows, for
+  an iteration), one `FETCH` event per batch fetched, and a `MOVE`
+  for a `forward`. The paths name asyncpg's private `BaseCursor`
+  methods, since those are the round trips.
+
+- The capture policy is deliberate about sensitive data: query
+  arguments are never recorded, under any setting (they reduce to a
+  count); the SQL text reduces to its length unless the `statement`
+  setting is on; a COPY's records, source and output reduce to a
+  count or a type; and the connect event captures none of its
+  arguments, which carry the password. asyncpg sends arguments
+  server-side, so the recorded text always carries `$1` placeholders
+  rather than data, unless the application interpolated its own.
+
+- A failing statement records the driver's exception
+  (`asyncpg.UndefinedTableError`, say) on the event as it escapes; a
+  timeout or cancellation arrives the same way. A pool made after the
+  instrumentation is applied opens its connections through the
+  bound `connect`, on the task that awaited it, so those record too.
+
+## Settings
+
+| Setting | Default | Controls |
+| ------- | ------- | -------- |
+| `statement` | `false` | Whether each query event records the SQL text as handed to the driver, as `statement`. 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 `$n` placeholders rather than data. Query arguments are never recorded either way. |
+
+```toml
+[[instrument]]
+name = "asyncpg"
+statement = true
+```
+
+## With the sqlalchemy instrumentation
+
+An instrumented asyncpg beneath the core package's `sqlalchemy`
+target (the `postgresql+asyncpg` async engine) 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` they nest beneath each statement. The
+asyncpg dialect prepares every statement and fetches through the
+prepared statement, so what appears beneath `do_execute` is the
+prepared statement's own fetch, and the commit's `execute` beneath
+`_commit_impl`. Raw asyncpg use beside the engine records at the top
+level either way, and the dialect's setup queries show up beside the
+tree regardless, since they run outside the recorded seams.
+
+## Known gap
+
+A connection taken from a pool (`async with pool.acquire() as conn:`)
+does not yet record its queries: asyncpg's pool proxy calls the
+connection's methods through the class, `Connection.fetchval(
+connection, query)`, a calling convention wrapture's signature check
+does not yet handle. A fix is on wrapture's side and this note goes
+when it lands. The pool's connects record, and so does everything on
+a connection you opened yourself.
+
+## How it patches
+
+For the implementation detail see the module docstrings of
+[connection.py](connection.py) (connect and the Connection methods),
+[prepared.py](prepared.py) (prepared statements) and
+[cursor.py](cursor.py) (server-side cursors).

wrapture_instrumentation_postgresql/asyncpg/__init__.py

@@ -0,0 +1,75 @@
+"""Instrumentation for asyncpg: 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, each
+recording around its await.
+
+This module imports only wrapture. Everything that touches asyncpg
+lives in the sibling modules, one per asyncpg module patched
+(connection.py for asyncpg and asyncpg.connection, prepared.py for
+asyncpg.prepared_stmt, cursor.py for asyncpg.cursor), each importing
+only wrapture at top level and reaching asyncpg through the module
+the hook is handed, so loading this class when a config loads never
+imports asyncpg ahead of the hooks meant to fire on its import.
+
+Importing asyncpg initialises every one of those modules, so all four
+hooks fire from the one import; each binds through the module it is
+handed and cleans up its own bindings.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import wrapture
+from wrapture import Setting
+
+from . import connection, cursor, prepared
+
+
+class AsyncpgInstrumentation(wrapture.Instrumentation):
+    """Query and transaction tracing for asyncpg."""
+
+    description = "Query and transaction tracing for asyncpg."
+
+    target = "asyncpg"
+    supports = ">=0.29,<1"
+    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("asyncpg")
+    def asyncpg(self, name: str, module: Any) -> None:
+        """Bind the package's own spelling of connect once asyncpg
+        exists."""
+
+        connection.instrument_package(module, self)
+
+    @wrapture.instrumentation_hook("asyncpg.connection")
+    def asyncpg_connection(self, name: str, module: Any) -> None:
+        """Bind connect and the Connection methods once
+        asyncpg.connection exists."""
+
+        connection.instrument(module, self)
+
+    @wrapture.instrumentation_hook("asyncpg.prepared_stmt")
+    def asyncpg_prepared_stmt(self, name: str, module: Any) -> None:
+        """Bind the prepared statement's own fetch family once
+        asyncpg.prepared_stmt exists."""
+
+        prepared.instrument(module, self)
+
+    @wrapture.instrumentation_hook("asyncpg.cursor")
+    def asyncpg_cursor(self, name: str, module: Any) -> None:
+        """Bind the server-side cursor round trips once asyncpg.cursor
+        exists."""
+
+        cursor.instrument(module, self)

wrapture_instrumentation_postgresql/asyncpg/connection.py

@@ -0,0 +1,220 @@
+"""The connection seams: connections being opened, and Connection's
+execute, fetch and COPY families, recorded as database events around
+their awaits.
+
+asyncpg's public surface is pure Python coroutine methods on
+`Connection`, so each is bound in place on the class and records
+around its await: `execute` (with or without arguments, a script of
+several statements included), `executemany`, `fetch`, `fetchrow`,
+`fetchval`, `fetchmany` (0.30 and later), and the four COPY methods.
+The Cython protocol beneath them does the socket work and is where
+the seams stop. A `PoolConnectionProxy` forwards each of these to the
+real Connection's class attribute, so pooled connections record
+through the same bindings.
+
+`asyncpg.connection.connect` is the coroutine function every
+connection comes from; `asyncpg.connect` is the same function under
+the package's name, bound at import, so it is bound as a module
+attribute too, as sqlite3's two spellings of `connect` are. A pool
+looks `connection.connect` up when it is created, so a pool made after
+the instrumentation is applied opens its connections through the
+binding, on the task that awaits it. The connect event's arguments are
+never captured (the dsn and the keyword form carry the password); the
+server reached is annotated from the connection afterwards.
+
+Transactions need no bindings of their own: `Transaction.start()`,
+its commit and its rollback issue BEGIN, COMMIT, ROLLBACK, SAVEPOINT,
+RELEASE SAVEPOINT and ROLLBACK TO through `Connection.execute`, so
+each boundary records as a statement event with its own operation.
+
+Every event carries the database contract keys and the server the
+connection reached, which asyncpg keeps privately (`_params.database`
+and `_addr`, a host and port pair or a socket path, stable since its
+0.1x line) and which are read through a guarded helper. The SQL text
+rides as `statement` only when the setting is on; the query arguments
+(asyncpg's `*args`, sent server-side) are never recorded and reduce
+to a count.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import wrapture
+
+from ..common import SYSTEM, captured, statement_data
+
+
+def server_of(connection: Any) -> dict[str, Any]:
+    """The `database`, `host` and `port` keys for an asyncpg
+    connection, from what it keeps privately; whichever it can
+    supply."""
+
+    data: dict[str, Any] = {}
+
+    params = getattr(connection, "_params", None)
+    database = getattr(params, "database", None)
+    if isinstance(database, str) and database:
+        data["database"] = database
+
+    addr = getattr(connection, "_addr", None)
+    if isinstance(addr, tuple) and len(addr) == 2:
+        host, port = addr
+        if isinstance(host, str):
+            data["host"] = host
+        if isinstance(port, int):
+            data["port"] = port
+    elif isinstance(addr, str) and addr:
+        data["host"] = addr
+
+    return data
+
+
+def query_of(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+    """The query a Connection method was handed: the first positional,
+    else the keyword the method spells it as."""
+
+    if args:
+        return args[0]
+
+    return kwargs.get("query", kwargs.get("command"))
+
+
+def database_binding(
+    owner: Any, name: str, capture_args: Any = captured
+) -> wrapture.Binding:
+    """A ready binding on one of the driver's methods, shared with the
+    prepared statement and cursor modules."""
+
+    return wrapture.binding(
+        owner,
+        name,
+        category="database",
+        leaf=True,
+        capture_args=capture_args,
+        capture_result=captured,
+    )
+
+
+def opens_binding(owner: Any) -> wrapture.Binding:
+    """The connect binding for either spelling of connect."""
+
+    async def opens(
+        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))
+
+        return connection
+
+    binding = database_binding(owner, "connect", "none")
+    binding.on_call.decorates(opens)
+
+    return binding
+
+
+def instrument_package(module: Any, instrumentation: wrapture.Instrumentation) -> None:
+    """Bind the package's own spelling of connect; register its
+    removal as this trigger's cleanup."""
+
+    group = wrapture.bindings(connect=opens_binding(module))
+    group.apply()
+
+    instrumentation.on_cleanup(group.remove)
+
+
+def instrument(module: Any, instrumentation: wrapture.Instrumentation) -> None:
+    """Bind connect and the Connection methods; register their removal
+    as this trigger's cleanup."""
+
+    settings = instrumentation.settings
+    record_statement = bool(settings["statement"])
+
+    def data_for(
+        connection: Any, query: Any, operation: str | None = None
+    ) -> dict[str, Any]:
+        return statement_data(
+            query, connection, server_of(connection), record_statement, operation
+        )
+
+    def queries(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        async def record() -> Any:
+            wrapture.annotate(**data_for(instance, query_of(args, kwargs)))
+
+            return await wrapped(*args, **kwargs)
+
+        return record()
+
+    def copies_table(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        async def record() -> Any:
+            table = args[0] if args else kwargs.get("table_name")
+            data = data_for(instance, None, "COPY")
+            if isinstance(table, str):
+                data["collection"] = table
+            wrapture.annotate(**data)
+
+            return await wrapped(*args, **kwargs)
+
+        return record()
+
+    def copies_query(
+        wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        async def record() -> Any:
+            data = data_for(instance, query_of(args, kwargs), "COPY")
+            wrapture.annotate(**data)
+
+            # With arguments, asyncpg renders them into the query
+            # through a SELECT on this same connection before the COPY:
+            # a bound call beneath this leaf, silenced, whose own
+            # annotation would otherwise leave the event saying
+            # SELECT. The COPY's data is written again on the way out.
+
+            try:
+                return await wrapped(*args, **kwargs)
+            finally:
+                wrapture.annotate(**data)
+
+        return record()
+
+    named: dict[str, wrapture.Binding] = {"connect": opens_binding(module)}
+
+    connection_class = module.Connection
+
+    # The execute and fetch family; fetchmany arrived in 0.30.
+
+    for method in (
+        "execute",
+        "executemany",
+        "fetch",
+        "fetchrow",
+        "fetchval",
+        "fetchmany",
+    ):
+        if method not in vars(connection_class):
+            continue
+        bound = database_binding(connection_class, method)
+        bound.on_call.decorates(queries)
+        named[method] = bound
+
+    # COPY: three methods name a table, one takes a query.
+
+    for method in ("copy_from_table", "copy_to_table", "copy_records_to_table"):
+        bound = database_binding(connection_class, method)
+        bound.on_call.decorates(copies_table)
+        named[method] = bound
+
+    bound = database_binding(connection_class, "copy_from_query")
+    bound.on_call.decorates(copies_query)
+    named["copy_from_query"] = bound
+
+    group = wrapture.bindings(**named)
+    group.apply()
+
+    instrumentation.on_cleanup(group.remove)

wrapture_instrumentation_postgresql/asyncpg/cursor.py

@@ -0,0 +1,88 @@
+"""The server-side cursor seams: each round trip a cursor makes to
+its portal, recorded as a database event around its await.
+
+`connection.cursor(query, *args)` inside a transaction gives a
+server-side cursor, iterated with `async for` or driven by `fetch`,
+`fetchrow` and `forward`. Nothing there passes through Connection's
+public methods: the DECLARE happens inside the cursor's own bind, and
+every FETCH inside its exec. With no binding at all a query run this
+way would leave no event, so the three private methods on
+`BaseCursor` that make the round trips are bound: `_bind_exec` (the
+DECLARE with the first batch, what `async for` does), `_bind` (the
+DECLARE alone, for a cursor awaited then driven by hand) and `_exec`
+(each FETCH). One event per round trip, whichever public method drove
+it; binding `__anext__` instead would make an event per buffered row,
+fifty times over. sqlalchemy's `_commit_impl` is the precedent for a
+private seam where the public ones do not bound the work. `forward`,
+which skips rows with a MOVE of its own rather than through `_exec`,
+is bound as the public method it is.
+
+The cursor knows its query (`_query`) and its connection, which
+supply the statement (when the setting is on) and the server keys.
+The operation is DECLARE for the two binds, FETCH for the exec and
+MOVE for a forward.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import wrapture
+
+from ..common import statement_data
+from .connection import database_binding, server_of
+
+
+def instrument(module: Any, instrumentation: wrapture.Instrumentation) -> None:
+    """Bind the cursor's round trips; register their removal as this
+    trigger's cleanup."""
+
+    settings = instrumentation.settings
+    record_statement = bool(settings["statement"])
+
+    def round_trips(operation: str) -> Any:
+        def decorator(
+            wrapped: Any,
+            instance: Any,
+            args: tuple[Any, ...],
+            kwargs: dict[str, Any],
+        ) -> Any:
+            async def record() -> Any:
+                connection = getattr(instance, "_connection", None)
+                query = getattr(instance, "_query", None)
+                wrapture.annotate(
+                    **statement_data(
+                        query,
+                        connection,
+                        server_of(connection),
+                        record_statement,
+                        operation,
+                    )
+                )
+
+                return await wrapped(*args, **kwargs)
+
+            return record()
+
+        return decorator
+
+    cursor_class = module.BaseCursor
+    named: dict[str, wrapture.Binding] = {}
+
+    for method, operation in (
+        ("_bind_exec", "DECLARE"),
+        ("_bind", "DECLARE"),
+        ("_exec", "FETCH"),
+    ):
+        bound = database_binding(cursor_class, method)
+        bound.on_call.decorates(round_trips(operation))
+        named[method.lstrip("_")] = bound
+
+    forward = database_binding(module.Cursor, "forward")
+    forward.on_call.decorates(round_trips("MOVE"))
+    named["forward"] = forward
+
+    group = wrapture.bindings(**named)
+    group.apply()
+
+    instrumentation.on_cleanup(group.remove)

wrapture_instrumentation_postgresql/asyncpg/prepared.py

@@ -0,0 +1,74 @@
+"""The prepared statement seams: a PreparedStatement's own fetch and
+execute family, recorded as database events around their awaits.
+
+`await connection.prepare(query)` hands back a PreparedStatement whose
+`fetch`, `fetchrow`, `fetchval`, `fetchmany`, `executemany` and
+`explain` bypass Connection's public methods, driving the protocol
+through the statement's own private path, so each is bound in its
+own right. The statement's SQL is public (`get_query()`), and the
+connection it belongs to supplies the server keys. The `prepare` call
+itself is not a query and is not recorded; the round trip that
+prepares happens inside it.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import wrapture
+
+from ..common import statement_data
+from .connection import database_binding, server_of
+
+
+def instrument(module: Any, instrumentation: wrapture.Instrumentation) -> None:
+    """Bind the prepared statement's methods; register their removal
+    as this trigger's cleanup."""
+
+    settings = instrumentation.settings
+    record_statement = bool(settings["statement"])
+
+    def executes(operation: str | None = None) -> Any:
+        def decorator(
+            wrapped: Any,
+            instance: Any,
+            args: tuple[Any, ...],
+            kwargs: dict[str, Any],
+        ) -> Any:
+            async def record() -> Any:
+                connection = getattr(instance, "_connection", None)
+                query = instance.get_query()
+                wrapture.annotate(
+                    **statement_data(
+                        query,
+                        connection,
+                        server_of(connection),
+                        record_statement,
+                        operation,
+                    )
+                )
+
+                return await wrapped(*args, **kwargs)
+
+            return record()
+
+        return decorator
+
+    statement_class = module.PreparedStatement
+    named: dict[str, wrapture.Binding] = {}
+
+    for method in ("fetch", "fetchrow", "fetchval", "fetchmany", "executemany"):
+        if method not in vars(statement_class):
+            continue
+        bound = database_binding(statement_class, method)
+        bound.on_call.decorates(executes())
+        named[method] = bound
+
+    explain = database_binding(statement_class, "explain")
+    explain.on_call.decorates(executes("EXPLAIN"))
+    named["explain"] = explain
+
+    group = wrapture.bindings(**named)
+    group.apply()
+
+    instrumentation.on_cleanup(group.remove)