slowpoke-python 0.1.3__py3-none-any.whl

slowpoke/_tracer.py

@@ -0,0 +1,235 @@
+import contextvars
+import json
+import math
+import os
+import re
+import sys
+import time
+
+VERSION = "0.1.3"
+
+SERVER = 2
+CLIENT = 3
+CONSUMER = 5
+
+# The trace of the request or job running in this thread or asyncio task. Worker threads and
+# sync_to_async/run_in_threadpool copy the context, so their queries land in the same trace object.
+_current = contextvars.ContextVar("slowpoke_trace", default=None)
+
+_STATEMENT = re.compile(r"[\s(]*([A-Za-z]+)")
+_DB_SYSTEMS = {"postgres": "postgresql", "mssql": "microsoft.sql_server", "sqlserver": "microsoft.sql_server"}
+
+
+def current():
+    return _current.get()
+
+
+def _random_id(size):
+    return os.urandom(size).hex()
+
+
+class Trace:
+    __slots__ = ("kind", "name", "start", "end", "error", "trace_id", "span_id", "attributes", "queries", "dropped", "task")
+
+    def __init__(self, kind, name, start, trace_id, span_id):
+        self.kind = kind
+        self.name = name
+        self.start = start
+        self.end = None
+        self.error = False
+        self.trace_id = trace_id
+        self.span_id = span_id
+        self.attributes = []
+        self.queries = []
+        self.dropped = 0
+        self.task = _running_task()
+
+
+class Tracer:
+    """Collects one trace per HTTP request or job and turns it into OTLP/JSON for the agent: a SERVER
+    (or CONSUMER) span and one CLIENT span per query, with the SQL as the driver received it
+    (placeholders, never parameter values) and the application line that ran it.
+
+    Every public method swallows its own errors: observability must never break the app."""
+
+    def __init__(self, origin, submit, service="app", version=VERSION, max_queries=500, max_sql_length=10000,
+                 clock=time.time, ids=_random_id):
+        self.origin = origin
+        self.submit = submit
+        self.service = service
+        self.version = version
+        self.max_queries = max(0, int(max_queries))
+        self.max_sql_length = max(1, int(max_sql_length))
+        self.clock = clock
+        self.ids = ids
+
+    # ------------------------------------------------------------ requests
+
+    def start_request(self, method):
+        try:
+            method = str(method).upper()
+            trace = Trace(SERVER, method, self.clock(), self.ids(16), self.ids(8))
+            trace.attributes.append(_kv("http.request.method", method))
+            _current.set(trace)
+            return trace
+        except Exception:
+            return None
+
+    def finish_request(self, trace, route, path, status):
+        if trace is None or trace.end is not None:
+            return
+        try:
+            trace.end = self.clock()
+            if route:
+                route = "/" + str(route).lstrip("/")
+                trace.name = trace.name + " " + route
+                trace.attributes.append(_kv("http.route", route))
+            else:
+                trace.attributes.append(_kv("url.path", "/" + str(path or "").split("?", 1)[0].lstrip("/")))
+            status = int(status)
+            trace.attributes.append(_kv("http.response.status_code", status))
+            trace.error = status >= 500
+        except Exception:
+            trace.end = trace.end or self.clock()
+        self._done(trace)
+
+    # ------------------------------------------------------------ jobs
+
+    def start_job(self, name, queue=None):
+        return self._start_background(name, "job", queue)
+
+    def start_command(self, name):
+        """A command run by cron: nobody waits for it, which is why nobody notices when it doubles."""
+        return self._start_background(name, "command", None)
+
+    def _start_background(self, name, kind, queue):
+        running = _current.get()
+        if running is not None and running.end is None:
+            return None  # work run inside a request or another job belongs to it
+        try:
+            trace = Trace(CONSUMER, str(name), self.clock(), self.ids(16), self.ids(8))
+            # The agent files it under Jobs by this attribute, instead of among the endpoints.
+            trace.attributes.append(_kv("slowpoke.kind", kind))
+            if queue:
+                trace.attributes.append(_kv("messaging.destination.name", str(queue)))
+            _current.set(trace)
+            return trace
+        except Exception:
+            return None
+
+    def finish_job(self, trace, failed=False):
+        if trace is None or trace.end is not None:
+            return
+        trace.end = self.clock()
+        trace.error = bool(failed)
+        self._done(trace)
+
+    # ------------------------------------------------------------ queries
+
+    def record_query(self, sql, seconds, system):
+        """Called right after a statement ran, from the thread that ran it."""
+        trace = _current.get()
+        if trace is None or trace.end is not None:
+            return  # outside requests and jobs, or after the response
+        try:
+            if len(trace.queries) >= self.max_queries:
+                trace.dropped += 1
+                return
+            end = self.clock()
+            sql = sql if isinstance(sql, str) else str(sql)
+            trace.queries.append((
+                sql[:self.max_sql_length],
+                system,
+                max(trace.start, end - max(0.0, seconds)),
+                end,
+                self.origin.find(trace.task),
+            ))
+        except Exception:
+            pass  # never let observability break the query that was just run
+
+    # ------------------------------------------------------------ sending
+
+    def _done(self, trace):
+        trace.task = None  # the queued trace must not keep the request's coroutine alive
+        try:
+            if _current.get() is trace:
+                _current.set(None)
+        except Exception:
+            pass
+        try:
+            self.submit(trace)
+        except Exception:
+            pass  # the agent is optional: a missing or broken one costs a trace, nothing else
+
+    def encode(self, t):
+        """The OTLP/JSON payload of a finished trace. Runs on the background thread."""
+        root = {
+            "traceId": t.trace_id,
+            "spanId": t.span_id,
+            "name": t.name,
+            "kind": t.kind,
+            "startTimeUnixNano": _nanos(t.start),
+            "endTimeUnixNano": _nanos(t.end),
+            "attributes": list(t.attributes),
+        }
+        if t.dropped > 0:
+            root["attributes"].append(_kv("slowpoke.dropped_queries", t.dropped))
+        if t.error:
+            root["status"] = {"code": 2}
+        spans = [root]
+        for sql, system, start, end, origin in list(t.queries):
+            attributes = [_kv("db.system.name", _DB_SYSTEMS.get(system, system)), _kv("db.query.text", sql)]
+            if origin is not None:
+                attributes.append(_kv("code.file.path", origin[0]))
+                if origin[1] is not None:
+                    attributes.append(_kv("code.line.number", origin[1]))
+            m = _STATEMENT.match(sql)
+            spans.append({
+                "traceId": t.trace_id,
+                "spanId": self.ids(8),
+                "parentSpanId": t.span_id,
+                "name": m.group(1).upper() if m else "QUERY",
+                "kind": CLIENT,
+                "startTimeUnixNano": _nanos(start),
+                "endTimeUnixNano": _nanos(end),
+                "attributes": attributes,
+            })
+        return {"resourceSpans": [{
+            "resource": {"attributes": [
+                _kv("service.name", self.service),
+                _kv("telemetry.sdk.name", "slowpoke-python"),
+                _kv("telemetry.sdk.language", "python"),
+                _kv("telemetry.sdk.version", self.version),
+            ]},
+            "scopeSpans": [{
+                "scope": {"name": "slowpoke", "version": self.version},
+                "spans": spans,
+            }],
+        }]}
+
+    def encode_json(self, trace):
+        text = json.dumps(self.encode(trace), ensure_ascii=False, separators=(",", ":"))
+        return text.encode("utf-8", "replace")
+
+
+def _running_task():
+    asyncio = sys.modules.get("asyncio")
+    if asyncio is None:
+        return None
+    try:
+        return asyncio.current_task() if asyncio._get_running_loop() is not None else None
+    except Exception:
+        return None
+
+
+def _nanos(seconds):
+    # 64-bit integers travel as strings. Microseconds, as the PHP packages, and rounded the same
+    # way: Python's round() goes to the even digit (round(0.5) is 0), math.floor(x + 0.5) does not,
+    # and the same measure must produce the same trace in every language.
+    return "%d000" % math.floor(seconds * 1e6 + 0.5)
+
+
+def _kv(key, value):
+    if isinstance(value, int) and not isinstance(value, bool):
+        return {"key": key, "value": {"intValue": str(value)}}
+    return {"key": key, "value": {"stringValue": str(value)}}

slowpoke/asgi.py

@@ -0,0 +1,77 @@
+"""ASGI integration for FastAPI and Starlette:
+
+    from slowpoke.asgi import SlowpokeMiddleware
+    app.add_middleware(SlowpokeMiddleware)   # add it last: it becomes the outermost middleware
+
+Plain ASGI, not BaseHTTPMiddleware: the endpoint runs in the same task and context as the trace.
+"""
+
+import slowpoke
+
+
+class SlowpokeMiddleware:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        tracer = slowpoke.get_tracer() if scope.get("type") == "http" else None
+        trace = tracer.start_request(scope.get("method", "GET")) if tracer is not None else None
+        if trace is None:
+            return await self.app(scope, receive, send)
+
+        path = scope.get("path", "/")
+        root_path = scope.get("root_path", "")
+        status = []
+
+        async def send_with_status(message):
+            if message.get("type") == "http.response.start" and not status:
+                status.append(message.get("status", 200))
+            await send(message)
+
+        try:
+            await self.app(scope, receive, send_with_status)
+        except BaseException:
+            _finish(tracer, trace, scope, root_path, path, 500)
+            raise
+        _finish(tracer, trace, scope, root_path, path, status[0] if status else 500)
+
+
+def _finish(tracer, trace, scope, root_path, path, status):
+    try:
+        route = route_template(scope, root_path)
+    except Exception:
+        route = None
+    tracer.finish_request(trace, route, path, status)
+
+
+def route_template(scope, initial_root_path=""):
+    """The template of the route that handled the request, read after routing: FastAPI leaves the route
+    in scope["route"]; plain Starlette only the endpoint, looked up in the app's routes."""
+    route = scope.get("route")
+    template = getattr(route, "path", None)
+    if isinstance(template, str):
+        # A mounted sub-application knows its routes without the mount prefix.
+        final_root = scope.get("root_path", "") or ""
+        prefix = final_root[len(initial_root_path):] if final_root.startswith(initial_root_path) else ""
+        if prefix and not template.startswith(prefix + "/"):
+            template = prefix.rstrip("/") + template
+        return template
+    endpoint = scope.get("endpoint")
+    app = scope.get("app")
+    if endpoint is None or app is None:
+        return None
+    return _find_endpoint(getattr(app, "routes", None) or [], endpoint, "", 0)
+
+
+def _find_endpoint(routes, endpoint, prefix, depth):
+    if depth > 8:
+        return None
+    for r in routes:
+        if getattr(r, "endpoint", None) is endpoint and isinstance(getattr(r, "path", None), str):
+            return prefix + r.path
+        children = getattr(r, "routes", None)
+        if children and isinstance(getattr(r, "path", None), str):
+            found = _find_endpoint(children, endpoint, prefix + r.path, depth + 1)
+            if found is not None:
+                return found
+    return None

slowpoke/celery.py

@@ -0,0 +1,67 @@
+"""Celery integration: one trace per task a worker runs.
+
+    # wherever the Celery app is created
+    import slowpoke.celery
+
+    slowpoke.celery.install()
+
+A task dispatched from a request is not traced there: it is traced where it runs, in the worker.
+"""
+
+import slowpoke
+
+# Bounded on purpose: a worker that somehow never sees the end of its tasks must not grow a map
+# forever. Past this many open tasks the next ones are simply not traced.
+MAX_OPEN = 512
+
+_open = {}
+_state = {"installed": False}
+
+
+def install():
+    """Connects the worker signals. Idempotent, and never raises: a monitoring package must not be
+    able to stop a worker from starting."""
+    try:
+        if _state["installed"]:
+            return True
+        from celery import signals
+
+        signals.task_prerun.connect(_prerun, dispatch_uid="slowpoke")
+        signals.task_postrun.connect(_postrun, dispatch_uid="slowpoke")
+        _state["installed"] = True
+        return True
+    except Exception:
+        return False
+
+
+def _prerun(task_id=None, task=None, **kwargs):
+    try:
+        tracer = slowpoke.get_tracer()
+        if tracer is None or task_id is None or len(_open) >= MAX_OPEN:
+            return
+        trace = tracer.start_job(getattr(task, "name", None) or "task", _queue(task))
+        if trace is not None:
+            _open[task_id] = trace
+    except Exception:
+        pass  # never let observability break the worker
+
+
+def _postrun(task_id=None, state=None, **kwargs):
+    try:
+        trace = _open.pop(task_id, None)
+        if trace is None:
+            return
+        tracer = slowpoke.get_tracer()
+        if tracer is not None:
+            # RETRY is this attempt giving up: to whoever is waiting for the work it is a failure.
+            tracer.finish_job(trace, failed=str(state) in ("FAILURE", "RETRY"))
+    except Exception:
+        pass
+
+
+def _queue(task):
+    try:
+        info = getattr(getattr(task, "request", None), "delivery_info", None) or {}
+        return info.get("routing_key") or None
+    except Exception:
+        return None

slowpoke/django/__init__.py

@@ -0,0 +1,226 @@
+"""Django integration: add "slowpoke.django" to INSTALLED_APPS and
+"slowpoke.django.SlowpokeMiddleware" at the top of MIDDLEWARE."""
+
+import functools
+import os
+import time
+
+from asgiref.sync import iscoroutinefunction, markcoroutinefunction
+from django.db import connections
+from django.db.backends.signals import connection_created
+
+import slowpoke
+from slowpoke._tracer import current
+
+_state = {"signal": False, "commands": False}
+
+# Commands that never end: traced as one run they would hold a trace open for as long as the
+# process lives, and the requests or tasks inside them would be lost. Names from Django itself and
+# from the workers people usually run with manage.py.
+LONG_RUNNING = (
+    "runserver", "runserver_plus", "testserver", "runworker", "shell", "dbshell", "test",
+    "qcluster", "rqworker", "rqscheduler", "process_tasks", "run_huey", "celery", "mail_queue",
+)
+
+
+def install():
+    """Hooks every database connection, the ones already open and the ones opened later in other threads
+    or async contexts. Idempotent: a statement is never recorded twice."""
+    try:
+        from django.conf import settings
+
+        base = getattr(settings, "BASE_DIR", None)
+        slowpoke.set_defaults(code_root=str(base) if base else None)
+        if not _state["signal"]:
+            connection_created.connect(_on_connection_created, dispatch_uid="slowpoke")
+            _state["signal"] = True
+        _wrap_open_connections()
+        _trace_commands()
+    except Exception:
+        pass
+
+
+def _trace_commands():
+    """One trace per management command, which on a server means cron. SLOWPOKE_COMMANDS=false
+    turns it off; the commands that never end are left alone whatever it says."""
+    if _state["commands"] or os.environ.get("SLOWPOKE_COMMANDS", "").strip().lower() in ("0", "false", "no", "off"):
+        return
+    try:
+        from django.core.management.base import BaseCommand
+    except Exception:
+        return
+    if getattr(BaseCommand.execute, "_slowpoke", False):
+        _state["commands"] = True
+        return
+    original = BaseCommand.execute
+
+    @functools.wraps(original)
+    def execute(self, *args, **options):
+        name = command_name(self)
+        if name in LONG_RUNNING:
+            return original(self, *args, **options)
+        with slowpoke.command(name):
+            return original(self, *args, **options)
+
+    execute._slowpoke = True
+    BaseCommand.execute = execute
+    _state["commands"] = True
+
+
+def command_name(cmd):
+    """The name a person would type: "close_orders", from myapp.management.commands.close_orders."""
+    module = getattr(type(cmd), "__module__", "") or ""
+    return module.rsplit(".", 1)[-1] or type(cmd).__name__
+
+
+def _wrap_open_connections():
+    try:
+        for conn in connections.all(initialized_only=True):
+            _wrap(conn)
+    except TypeError:  # before Django 4.1
+        for conn in connections.all():
+            _wrap(conn)
+
+
+def _on_connection_created(sender, connection, **kwargs):
+    _wrap(connection)
+
+
+def _wrap(conn):
+    wrappers = conn.execute_wrappers
+    if _execute_wrapper not in wrappers:
+        # First, not last: connection.execute_wrapper() pops the last one when its block ends.
+        wrappers.insert(0, _execute_wrapper)
+
+
+def _execute_wrapper(execute, sql, params, many, context):
+    if current() is None:
+        return execute(sql, params, many, context)
+    started = time.perf_counter()
+    try:
+        return execute(sql, params, many, context)
+    finally:
+        seconds = time.perf_counter() - started
+        tracer = slowpoke.get_tracer()
+        if tracer is not None:
+            try:
+                vendor = context["connection"].vendor
+            except Exception:
+                vendor = "unknown"
+            # sql only, never params: Django passes %s placeholders to the driver
+            tracer.record_query(sql, seconds, vendor)
+
+
+class SlowpokeMiddleware:
+    """Opens one trace per request. Put it first in MIDDLEWARE so the timing covers the others."""
+
+    sync_capable = True
+    async_capable = True
+
+    def __init__(self, get_response):
+        self.get_response = get_response
+        self.is_async = iscoroutinefunction(get_response)
+        if self.is_async:
+            markcoroutinefunction(self)
+        install()
+
+    def __call__(self, request):
+        if self.is_async:
+            return self.__acall__(request)
+        tracer, trace = _start(request)
+        try:
+            response = self.get_response(request)
+        except BaseException:
+            _finish(tracer, trace, request, 500)
+            raise
+        _finish(tracer, trace, request, getattr(response, "status_code", 200))
+        return response
+
+    async def __acall__(self, request):
+        tracer, trace = _start(request)
+        try:
+            response = await self.get_response(request)
+        except BaseException:
+            _finish(tracer, trace, request, 500)
+            raise
+        _finish(tracer, trace, request, getattr(response, "status_code", 200))
+        return response
+
+
+def _start(request):
+    tracer = slowpoke.get_tracer()
+    if tracer is None:
+        return None, None
+    try:
+        _wrap_open_connections()
+    except Exception:
+        pass
+    return tracer, tracer.start_request(request.method)
+
+
+def _finish(tracer, trace, request, status):
+    if trace is None:
+        return
+    try:
+        match = getattr(request, "resolver_match", None)
+        route = route_template(match.route) if match is not None and getattr(match, "route", None) else None
+        path = request.path
+    except Exception:
+        route, path = None, "/"
+    tracer.finish_request(trace, route, path, status)
+
+
+def route_template(route):
+    """path() routes already read like templates ("orders/<int:pk>/", the agent normalizes them);
+    re_path() gives a regular expression, turned here into "legacy/{slug}/"."""
+    if not route or ("(" not in route and not route.startswith("^") and not route.endswith("$")):
+        return route or ""
+    if route.startswith("^"):
+        route = route[1:]
+    for end in ("\\Z", "$"):
+        if route.endswith(end) and not route.endswith("\\" + end):
+            route = route[: -len(end)]
+            break
+    out = []
+    i, n = 0, len(route)
+    while i < n:
+        c = route[i]
+        if c == "\\" and i + 1 < n:
+            out.append(route[i + 1])
+            i += 2
+        elif c == "(":
+            name = "param"
+            if route.startswith("(?P<", i):
+                close = route.find(">", i)
+                if close > 0:
+                    name = route[i + 4:close]
+            i = _group_end(route, i)
+            while i < n and route[i] in "?*+":
+                i += 1
+            out.append("{%s}" % name)
+        else:
+            out.append(c)
+            i += 1
+    return "".join(out)
+
+
+def _group_end(pattern, i):
+    """Index after the group that opens at pattern[i], skipping escapes and character classes."""
+    depth, in_class, n = 0, False, len(pattern)
+    while i < n:
+        c = pattern[i]
+        if c == "\\":
+            i += 2
+            continue
+        if in_class:
+            in_class = c != "]"
+        elif c == "[":
+            in_class = True
+        elif c == "(":
+            depth += 1
+        elif c == ")":
+            depth -= 1
+            if depth == 0:
+                return i + 1
+        i += 1
+    return n

slowpoke/django/apps.py

@@ -0,0 +1,12 @@
+from django.apps import AppConfig
+
+
+class SlowpokeConfig(AppConfig):
+    name = "slowpoke.django"
+    label = "slowpoke"
+    verbose_name = "Slowpoke"
+
+    def ready(self):
+        from . import install
+
+        install()

slowpoke/flask.py

@@ -0,0 +1,55 @@
+"""Flask integration (2.x and 3.x):
+
+    import slowpoke.flask, slowpoke.sqlalchemy
+    slowpoke.flask.init_app(app)
+    slowpoke.sqlalchemy.instrument(db.engine)   # or instrument() for every engine
+"""
+
+from flask import request
+
+import slowpoke
+
+_KEY = "slowpoke.trace"
+_STATUS = "slowpoke.status"
+
+
+def init_app(app):
+    """Idempotent. The trace opens before every other before_request hook, so a hook that answers
+    early (authentication, rate limits) is traced too."""
+    before = app.before_request_funcs.setdefault(None, [])
+    if _before not in before:
+        before.insert(0, _before)
+    if _after not in app.after_request_funcs.setdefault(None, []):
+        app.after_request(_after)
+    if _teardown not in app.teardown_request_funcs.setdefault(None, []):
+        app.teardown_request(_teardown)
+    return app
+
+
+def _before():
+    try:
+        tracer = slowpoke.get_tracer()
+        if tracer is not None:
+            request.environ[_KEY] = (tracer, tracer.start_request(request.method))
+    except Exception:
+        pass
+
+
+def _after(response):
+    try:
+        request.environ[_STATUS] = response.status_code
+    except Exception:
+        pass
+    return response
+
+
+def _teardown(exc):
+    try:
+        tracer, trace = request.environ.pop(_KEY, (None, None))
+        if trace is None:
+            return
+        status = request.environ.get(_STATUS) or (500 if exc is not None else 200)
+        rule = request.url_rule
+        tracer.finish_request(trace, rule.rule if rule is not None else None, request.path, status)
+    except Exception:
+        pass