tracelet-sdk 0.1.0__py3-none-any.whl

tracelet_sdk/__init__.py

@@ -0,0 +1,130 @@
+"""Tag each database query with the request that ran it.
+
+Queries get a sqlcommenter comment (https://google.github.io/sqlcommenter/)::
+
+    SELECT * FROM users WHERE id = %s /*route='GET%20%2Fusers%2F%3Aid',traceparent='00-…-…-01'*/
+
+Tracelet reads these comments to group queries into requests. Nothing is
+sent anywhere: the comment travels with the query to your database.
+
+Framework integrations live in submodules: ``tracelet_sdk.wsgi``,
+``tracelet_sdk.asgi``, ``tracelet_sdk.flask``, ``tracelet_sdk.django`` and
+``tracelet_sdk.sqlalchemy``.
+"""
+
+import contextvars
+import os
+import re
+from contextlib import contextmanager
+from typing import Callable, Dict, Iterator, Optional, Union
+from urllib.parse import quote
+
+__all__ = [
+    "RequestContext",
+    "request",
+    "current_request",
+    "comment",
+    "new_traceparent",
+    "parse_traceparent",
+]
+
+Route = Union[str, Callable[[], Optional[str]], None]
+
+_TRACEPARENT = re.compile(r"^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$")
+
+
+class RequestContext:
+    """The request a query belongs to."""
+
+    __slots__ = ("_route", "traceparent")
+
+    def __init__(self, route: Route = None, traceparent: Optional[str] = None):
+        self._route = route
+        self.traceparent = traceparent if parse_traceparent(traceparent) else new_traceparent(traceparent)
+
+    @property
+    def route(self) -> Optional[str]:
+        """The route pattern. Callables are evaluated now, so frameworks that
+        only know the route after routing can supply it lazily."""
+        route = self._route
+        if callable(route):
+            try:
+                route = route()
+            except Exception:  # never let tagging break a query
+                route = None
+        return route or None
+
+
+_current: "contextvars.ContextVar[Optional[RequestContext]]" = contextvars.ContextVar("tracelet_request", default=None)
+
+
+def parse_traceparent(header: Optional[str]) -> Optional[Dict[str, str]]:
+    """Parse a W3C traceparent header, or return None if it is not valid."""
+    if not isinstance(header, str):
+        return None
+    m = _TRACEPARENT.match(header.strip())
+    if not m or m.group(1) == "ff" or set(m.group(2)) == {"0"} or set(m.group(3)) == {"0"}:
+        return None
+    return {"version": m.group(1), "trace_id": m.group(2), "parent_id": m.group(3), "flags": m.group(4)}
+
+
+def new_traceparent(incoming: Optional[str] = None) -> str:
+    """A traceparent for a new request span. An incoming header's trace id is
+    kept, so Tracelet's traces line up with the rest of your tracing."""
+    parent = parse_traceparent(incoming)
+    trace_id = parent["trace_id"] if parent else os.urandom(16).hex()
+    flags = parent["flags"] if parent else "01"
+    return "00-{}-{}-{}".format(trace_id, os.urandom(8).hex(), flags)
+
+
+@contextmanager
+def request(route: Route = None, traceparent: Optional[str] = None) -> Iterator[RequestContext]:
+    """Run a block as one request: queries inside it are tagged with it.
+    Works across ``await`` and in threads that copy the context."""
+    ctx = RequestContext(route, traceparent)
+    token = _current.set(ctx)
+    try:
+        yield ctx
+    finally:
+        _current.reset(token)
+
+
+def current_request() -> Optional[RequestContext]:
+    """The current request context, or None outside a request."""
+    return _current.get()
+
+
+def _enc(value: str) -> str:
+    # Same as JavaScript's encodeURIComponent, then sqlcommenter's \' escape.
+    return quote(str(value), safe="-_.!~*'()").replace("'", "\\'")
+
+
+def comment(sql: str, fields: Optional[Dict[str, Optional[str]]] = None, *, escape_percent: bool = False) -> str:
+    """Append a sqlcommenter comment with ``fields`` (or the current
+    request's) to ``sql``.
+
+    Set ``escape_percent`` when the driver will %-format the SQL (psycopg,
+    MySQLdb and Django with parameters): the comment's URL-encoded values
+    contain ``%``, which must then be written as ``%%``.
+
+    SQL that already contains a comment is returned unchanged, as the
+    sqlcommenter spec requires.
+    """
+    if not isinstance(sql, str) or "*/" in sql:
+        return sql
+    if fields is None:
+        ctx = _current.get()
+        if ctx is None:
+            return sql
+        fields = {"route": ctx.route, "traceparent": ctx.traceparent}
+    pairs = ["{}='{}'".format(_enc(k), _enc(v)) for k, v in sorted(fields.items()) if v not in (None, "")]
+    if not pairs:
+        return sql
+
+    text = "/*" + ",".join(pairs) + "*/"
+    if escape_percent:
+        text = text.replace("%", "%%")
+    trimmed = sql.rstrip()
+    if trimmed.endswith(";"):
+        return "{} {};".format(trimmed[:-1].rstrip(), text)
+    return "{} {}".format(trimmed, text)

tracelet_sdk/asgi.py

@@ -0,0 +1,36 @@
+"""ASGI middleware: one request context per HTTP request (FastAPI, Starlette, Django ASGI)."""
+
+from typing import Any, Callable, Dict, Optional
+
+from . import new_traceparent, request
+
+RouteGetter = Callable[[Dict[str, Any]], Optional[str]]
+
+
+def default_route(scope: Dict[str, Any]) -> Optional[str]:
+    """The route pattern FastAPI records in the scope after routing."""
+    route = scope.get("route")
+    path = getattr(route, "path", None) or getattr(route, "path_format", None)
+    return "{} {}".format(scope.get("method", ""), path).strip() if path else None
+
+
+class TraceletMiddleware:
+    """Wrap any ASGI app. ``route`` is called with the ASGI scope each time a
+    query runs; the default reads the route FastAPI matched."""
+
+    def __init__(self, app: Callable, route: Optional[RouteGetter] = default_route):
+        self.app = app
+        self.route = route
+
+    async def __call__(self, scope: Dict[str, Any], receive: Callable, send: Callable) -> None:
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+        incoming = None
+        for name, value in scope.get("headers") or []:
+            if name == b"traceparent":
+                incoming = value.decode("latin-1")
+                break
+        getter = (lambda: self.route(scope)) if self.route else None
+        with request(route=getter, traceparent=new_traceparent(incoming)):
+            await self.app(scope, receive, send)

tracelet_sdk/django.py

@@ -0,0 +1,37 @@
+"""Django integration. Add to MIDDLEWARE::
+
+    MIDDLEWARE = ["tracelet_sdk.django.TraceletMiddleware", ...]
+
+Every query Django runs during a request is tagged with the URL pattern
+("GET /users/<int:pk>/"). No other setup is needed.
+"""
+
+from contextlib import ExitStack
+from typing import Any, Callable
+
+from . import comment, new_traceparent, request
+
+
+def _execute_wrapper(execute: Callable, sql: str, params: Any, many: bool, context: Any) -> Any:
+    # With parameters, Django's backends %-format the SQL, so the comment's
+    # % signs must be doubled. Without parameters they are sent as-is.
+    return execute(comment(sql, escape_percent=params is not None), params, many, context)
+
+
+class TraceletMiddleware:
+    def __init__(self, get_response: Callable):
+        self.get_response = get_response
+
+    def __call__(self, http_request: Any) -> Any:
+        from django.db import connections
+
+        def route():
+            match = getattr(http_request, "resolver_match", None)
+            pattern = getattr(match, "route", None)
+            return "{} /{}".format(http_request.method, pattern) if pattern is not None else None
+
+        incoming = http_request.META.get("HTTP_TRACEPARENT")
+        with request(route=route, traceparent=new_traceparent(incoming)), ExitStack() as stack:
+            for conn in connections.all():
+                stack.enter_context(conn.execute_wrapper(_execute_wrapper))
+            return self.get_response(http_request)

tracelet_sdk/flask.py

@@ -0,0 +1,21 @@
+"""Flask integration: ``tracelet_sdk.flask.instrument(app)``."""
+
+from typing import Any, Dict, Optional
+
+from .wsgi import TraceletMiddleware
+
+
+def _route(environ: Dict[str, Any]) -> Optional[str]:
+    from flask import has_request_context, request
+
+    if has_request_context() and request.url_rule is not None:
+        return "{} {}".format(request.method, request.url_rule.rule)
+    return None
+
+
+def instrument(app: Any) -> Any:
+    """Tag queries with the Flask route that ran them ("GET /users/<int:id>").
+    Combine with ``tracelet_sdk.sqlalchemy.instrument`` (Flask-SQLAlchemy:
+    ``instrument(db.engine)``) or call ``tracelet_sdk.comment`` yourself."""
+    app.wsgi_app = TraceletMiddleware(app.wsgi_app, route=_route)
+    return app

tracelet_sdk/sqlalchemy.py

@@ -0,0 +1,20 @@
+"""SQLAlchemy integration: ``tracelet_sdk.sqlalchemy.instrument(engine)``."""
+
+from typing import Any
+
+from . import comment
+
+
+def instrument(engine: Any) -> Any:
+    """Tag every statement the engine runs inside a request context.
+    Works with any driver; % signs are escaped for drivers that %-format."""
+    from sqlalchemy import event
+
+    target = getattr(engine, "sync_engine", engine)  # AsyncEngine wraps a sync engine
+
+    @event.listens_for(target, "before_cursor_execute", retval=True)
+    def _tag(conn, cursor, statement, parameters, context, executemany):
+        formats = conn.dialect.paramstyle in ("format", "pyformat") and parameters is not None
+        return comment(statement, escape_percent=formats), parameters
+
+    return engine

tracelet_sdk/wsgi.py

@@ -0,0 +1,26 @@
+"""WSGI middleware: one request context per HTTP request."""
+
+from typing import Any, Callable, Dict, Optional
+
+from . import new_traceparent, request
+
+RouteGetter = Callable[[Dict[str, Any]], Optional[str]]
+
+
+class TraceletMiddleware:
+    """Wrap any WSGI app. ``route`` is called with the WSGI environ each
+    time a query runs and should return the matched route pattern (for
+    example ``"GET /users/<id>"``). Without it, queries are still grouped by
+    request, just without a route. Raw URLs are never used, so ids and other
+    values in paths don't end up in your stats."""
+
+    def __init__(self, app: Callable, route: Optional[RouteGetter] = None):
+        self.app = app
+        self.route = route
+
+    def __call__(self, environ: Dict[str, Any], start_response: Callable):
+        getter = (lambda: self.route(environ)) if self.route else None
+        with request(route=getter, traceparent=new_traceparent(environ.get("HTTP_TRACEPARENT"))):
+            # Queries made while a streamed response is being iterated, after
+            # this returns, are not tagged.
+            return self.app(environ, start_response)

tracelet_sdk-0.1.0.dist-info/METADATA

@@ -0,0 +1,68 @@
+Metadata-Version: 2.4
+Name: tracelet-sdk
+Version: 0.1.0
+Summary: Tag your database queries with the request that ran them, so Tracelet can show per-request waterfalls and catch N+1 loops.
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/moseleydev/tracelet
+Keywords: tracelet,sqlcommenter,postgres,database,performance,n+1
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+
+# tracelet-sdk (Python)
+
+Tags each database query with the request that ran it, so [Tracelet](https://github.com/moseleydev/tracelet) can show per-request waterfalls and catch N+1 loops. Zero dependencies. Python 3.8+.
+
+```bash
+pip install tracelet-sdk
+```
+
+## Django
+
+```python
+MIDDLEWARE = [
+    "tracelet_sdk.django.TraceletMiddleware",
+    # ...
+]
+```
+
+That's all. Every query Django runs during a request carries the URL pattern (`GET /users/<int:pk>/`) and a trace id.
+
+## Flask
+
+```python
+from tracelet_sdk.flask import instrument
+from tracelet_sdk.sqlalchemy import instrument as instrument_engine
+
+instrument(app)
+instrument_engine(db.engine)  # Flask-SQLAlchemy, or any SQLAlchemy engine
+```
+
+## FastAPI / Starlette / any ASGI app
+
+```python
+from tracelet_sdk.asgi import TraceletMiddleware
+from tracelet_sdk.sqlalchemy import instrument
+
+app.add_middleware(TraceletMiddleware)  # reads the matched FastAPI route
+instrument(engine)
+```
+
+Plain WSGI apps can use `tracelet_sdk.wsgi.TraceletMiddleware(app, route=...)`.
+
+## Anything else
+
+```python
+import tracelet_sdk
+
+with tracelet_sdk.request(route="POST /import"):
+    cursor.execute(tracelet_sdk.comment("SELECT * FROM jobs WHERE id = %s", escape_percent=True), (job_id,))
+```
+
+Pass `escape_percent=True` when the driver %-formats the query (psycopg, MySQLdb and Django with parameters). The comment's URL-encoded values contain `%`, which must then be doubled. The Django and SQLAlchemy integrations do this for you.
+
+## Notes
+
+- **Traces:** an incoming `traceparent` header is continued, so traces line up with OpenTelemetry.
+- **No raw URLs:** paths are never recorded, so ids and other values in them don't end up in your stats.
+- **Existing comments:** SQL that already contains a comment is left unchanged, as the sqlcommenter spec requires.
+- **Streaming responses:** in WSGI, queries made while a streamed response is iterated are not tagged.

tracelet_sdk-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+tracelet_sdk/__init__.py,sha256=S3WLomwaB4867o4J-iZdrcA1ZORpi46X-ML9WWJqSxA,4687
+tracelet_sdk/asgi.py,sha256=hn5tBrA0omUHsZKGysd6MV7Iy9K1D4fTwVxW_kF0kfA,1463
+tracelet_sdk/django.py,sha256=1PSh_8GOFJR2zMuQEehmpoILxn1yM8fJMyb4iDatjzs,1467
+tracelet_sdk/flask.py,sha256=u6vCCDqOPg_HfWpPyrkfqyN1eBBaMFJQLX3XAaDcs9s,746
+tracelet_sdk/sqlalchemy.py,sha256=9nQxBbigru7aI4A3l7wfc0DxJocbKenqk9hBvWE-arA,757
+tracelet_sdk/wsgi.py,sha256=Cb-tNthK8Ieo9dJ9S8hfBqNwLtE3Z_zttSYfvg70PYU,1132
+tracelet_sdk-0.1.0.dist-info/METADATA,sha256=lCRT5xHxI2kzbvyHrNB1zIRH42sxkyCsHSWWaWoLDJE,2313
+tracelet_sdk-0.1.0.dist-info/WHEEL,sha256=a-udL2m4V_6Yn8b-uKzWb12RjRxwwMo6ohYFM24tpaY,91
+tracelet_sdk-0.1.0.dist-info/top_level.txt,sha256=EX3uds2sf1D-gOtjjJudnyfhgiQdxa-RyFzzDSJbrbg,13
+tracelet_sdk-0.1.0.dist-info/RECORD,,

tracelet_sdk-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

tracelet_sdk-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+tracelet_sdk