fintual-backend-devops-go / adapters / trace.py
trace.py
Raw
"""Opt-in HTTP-associated SQL diagnostics, never enabled during benchmarks."""
import json
import os
import time

from django.db import connection


class SQLTraceMiddleware:
    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        statements = {}

        def capture(execute, sql, params, many, context):
            start = time.perf_counter_ns()
            error = None
            try:
                return execute(sql, params, many, context)
            except Exception as exc:
                error = type(exc).__name__
                raise
            finally:
                key = (sql, error)
                item = statements.setdefault(key, {"sql": sql, "duration_ns": 0,
                                                    "error": error, "calls": 0})
                item["calls"] += 1
                item["duration_ns"] += time.perf_counter_ns() - start

        with connection.execute_wrapper(capture):
            response = self.get_response(request)
        record = {"request_id": request.headers.get("X-Harness-Request", ""),
                  "method": request.method, "path": request.path,
                  "status": response.status_code,
                  "count": sum(item["calls"] for item in statements.values()),
                  "statements": list(statements.values())}
        # A single append syscall prevents worker records from interleaving.
        fd = os.open(os.environ["SQL_TRACE"], os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
        try:
            os.write(fd, (json.dumps(record) + "\n").encode())
        finally:
            os.close(fd)
        return response