tracely-sdk 0.1.0__py3-none-any.whl

tracely/instrumentation/dbapi.py

@@ -0,0 +1,264 @@
+"""Database query instrumentation for SQL and MongoDB.
+
+Creates child Span objects linked to the active root span via context
+propagation, enabling hierarchical trace trees (FR56, FR57).
+
+Provides instrumentors for:
+- SQLAlchemy (via before/after_cursor_execute events)
+- Django ORM (via query callback)
+- MongoDB (via pymongo command monitoring)
+
+All instrumentors are fail-silent and never crash the host application.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any, Callable
+
+from tracely.context import get_current_span
+from tracely.span import Span
+
+logger = logging.getLogger("tracely")
+
+
+def _extract_operation(statement: str) -> str:
+    """Extract the SQL operation (SELECT, INSERT, UPDATE, DELETE) from a statement."""
+    normalized = statement.strip().upper()
+    for op in ("SELECT", "INSERT", "UPDATE", "DELETE"):
+        if normalized.startswith(op):
+            return op
+    return "UNKNOWN"
+
+
+class SQLAlchemyInstrumentor:
+    """Tracks SQLAlchemy query execution as child spans.
+
+    Creates a child Span linked to the active root/parent span when a
+    query executes. If no active span exists, records a standalone dict.
+    """
+
+    def __init__(
+        self,
+        on_span: Callable[[dict[str, Any]], None] | None = None,
+    ) -> None:
+        self._on_span = on_span
+        self._timings: dict[int, tuple[float, Span | None]] = {}
+
+    def before_cursor_execute(
+        self,
+        conn: Any,
+        cursor: Any,
+        statement: str,
+        parameters: Any,
+        context: Any,
+        executemany: bool,
+    ) -> None:
+        """Record start time and create child span before query execution."""
+        try:
+            parent = get_current_span()
+            span: Span | None = None
+            if parent is not None:
+                operation = _extract_operation(statement)
+                span = Span(
+                    name=f"SQL {operation}",
+                    parent=parent,
+                    kind="CLIENT",
+                )
+                span.set_attribute("db.system", "sql")
+                span.set_attribute("db.statement", statement)
+                span.set_attribute("db.operation", operation)
+            self._timings[id(context)] = (time.perf_counter(), span)
+        except Exception:
+            logger.debug("Error in before_cursor_execute", exc_info=True)
+
+    def after_cursor_execute(
+        self,
+        conn: Any,
+        cursor: Any,
+        statement: str,
+        parameters: Any,
+        context: Any,
+        executemany: bool,
+    ) -> None:
+        """Record query span after execution completes."""
+        try:
+            entry = self._timings.pop(id(context), None)
+            if entry is None:
+                return
+
+            start, span = entry
+            elapsed_ms = (time.perf_counter() - start) * 1000
+
+            if span is not None:
+                span.set_attribute("duration_ms", str(round(elapsed_ms, 3)))
+                span.set_status("OK")
+                span.end()
+
+                if self._on_span is not None:
+                    self._on_span(span.to_dict())
+            elif self._on_span is not None:
+                self._on_span({
+                    "span_type": "db",
+                    "db.system": "sql",
+                    "db.statement": statement,
+                    "db.operation": _extract_operation(statement),
+                    "duration_ms": round(elapsed_ms, 3),
+                    "error": False,
+                })
+        except Exception:
+            logger.debug("Error in after_cursor_execute", exc_info=True)
+
+
+class DjangoORMInstrumentor:
+    """Tracks Django ORM query execution as child spans."""
+
+    def __init__(
+        self,
+        on_span: Callable[[dict[str, Any]], None] | None = None,
+    ) -> None:
+        self._on_span = on_span
+
+    def on_query(
+        self,
+        sql: str,
+        duration_ms: float,
+        vendor: str = "sql",
+    ) -> None:
+        """Record a completed Django ORM query as a child span."""
+        try:
+            parent = get_current_span()
+            operation = _extract_operation(sql)
+
+            if parent is not None:
+                span = Span(
+                    name=f"SQL {operation}",
+                    parent=parent,
+                    kind="CLIENT",
+                )
+                span.set_attribute("db.system", vendor)
+                span.set_attribute("db.statement", sql)
+                span.set_attribute("db.operation", operation)
+                span.set_attribute("duration_ms", str(round(duration_ms, 3)))
+                span.set_status("OK")
+                span.end()
+
+                if self._on_span is not None:
+                    self._on_span(span.to_dict())
+            elif self._on_span is not None:
+                self._on_span({
+                    "span_type": "db",
+                    "db.system": vendor,
+                    "db.statement": sql,
+                    "db.operation": operation,
+                    "duration_ms": round(duration_ms, 3),
+                    "error": False,
+                })
+        except Exception:
+            logger.debug("Error in DjangoORMInstrumentor.on_query", exc_info=True)
+
+
+class MongoInstrumentor:
+    """Tracks MongoDB command execution as child spans."""
+
+    def __init__(
+        self,
+        on_span: Callable[[dict[str, Any]], None] | None = None,
+    ) -> None:
+        self._on_span = on_span
+        self._inflight: dict[int, dict[str, Any]] = {}
+
+    def on_command_start(
+        self,
+        command_name: str,
+        database_name: str,
+        request_id: int,
+    ) -> None:
+        """Record the start of a MongoDB command and create a child span."""
+        try:
+            parent = get_current_span()
+            span: Span | None = None
+            if parent is not None:
+                span = Span(
+                    name=f"MongoDB {command_name}",
+                    parent=parent,
+                    kind="CLIENT",
+                )
+                span.set_attribute("db.system", "mongodb")
+                span.set_attribute("db.operation", command_name)
+                span.set_attribute("db.name", database_name)
+
+            self._inflight[request_id] = {
+                "command_name": command_name,
+                "database_name": database_name,
+                "start": time.perf_counter(),
+                "span": span,
+            }
+        except Exception:
+            logger.debug("Error in MongoInstrumentor.on_command_start", exc_info=True)
+
+    def on_command_success(
+        self,
+        request_id: int,
+        duration_ms: float,
+    ) -> None:
+        """Record a successful MongoDB command completion."""
+        try:
+            info = self._inflight.pop(request_id, None)
+            if info is None:
+                return
+
+            span: Span | None = info.get("span")
+            if span is not None:
+                span.set_attribute("duration_ms", str(round(duration_ms, 3)))
+                span.set_status("OK")
+                span.end()
+
+                if self._on_span is not None:
+                    self._on_span(span.to_dict())
+            elif self._on_span is not None:
+                self._on_span({
+                    "span_type": "db",
+                    "db.system": "mongodb",
+                    "db.operation": info["command_name"],
+                    "db.name": info["database_name"],
+                    "duration_ms": round(duration_ms, 3),
+                    "error": False,
+                })
+        except Exception:
+            logger.debug("Error in MongoInstrumentor.on_command_success", exc_info=True)
+
+    def on_command_failure(
+        self,
+        request_id: int,
+        duration_ms: float,
+        failure: str,
+    ) -> None:
+        """Record a failed MongoDB command."""
+        try:
+            info = self._inflight.pop(request_id, None)
+            if info is None:
+                return
+
+            span: Span | None = info.get("span")
+            if span is not None:
+                span.set_attribute("duration_ms", str(round(duration_ms, 3)))
+                span.set_attribute("error.message", failure)
+                span.set_status("ERROR", failure)
+                span.end()
+
+                if self._on_span is not None:
+                    self._on_span(span.to_dict())
+            elif self._on_span is not None:
+                self._on_span({
+                    "span_type": "db",
+                    "db.system": "mongodb",
+                    "db.operation": info["command_name"],
+                    "db.name": info["database_name"],
+                    "duration_ms": round(duration_ms, 3),
+                    "error": True,
+                    "error.message": failure,
+                })
+        except Exception:
+            logger.debug("Error in MongoInstrumentor.on_command_failure", exc_info=True)

tracely/instrumentation/django_inst.py

@@ -0,0 +1,155 @@
+"""Django auto-instrumentation (Django middleware).
+
+Provides a Django-style middleware that creates structured Span objects
+for each HTTP request, with full trace hierarchy support via context
+propagation. Captures full request/response data (FR6/FR7).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable
+
+from tracely.capture import build_url, capture_request_data, capture_response_data
+from tracely.context import _span_context
+from tracely.instrumentation.base import BaseInstrumentor
+from tracely.span import Span
+from tracely.span_processor import on_span_end, on_span_start
+
+logger = logging.getLogger("tracely")
+
+
+def _extract_django_headers(meta: dict[str, Any]) -> dict[str, str]:
+    """Extract HTTP headers from Django request.META dict.
+
+    Django stores headers as HTTP_<NAME> with underscores replacing hyphens.
+    CONTENT_TYPE and CONTENT_LENGTH are special cases without HTTP_ prefix.
+    """
+    headers: dict[str, str] = {}
+    for key, value in meta.items():
+        if key.startswith("HTTP_"):
+            header_name = key[5:].lower().replace("_", "-")
+            headers[header_name] = str(value)
+        elif key == "CONTENT_TYPE":
+            headers["content-type"] = str(value)
+        elif key == "CONTENT_LENGTH":
+            headers["content-length"] = str(value)
+    return headers
+
+
+class TracelyDjangoMiddleware:
+    """Django middleware that creates root spans for HTTP requests.
+
+    Follows Django's middleware protocol: __init__(get_response) + __call__(request).
+    Creates a Span object with trace_id and span_id, sets it as the active span,
+    and captures HTTP attributes including full request/response data.
+    """
+
+    def __init__(
+        self,
+        get_response: Callable[..., Any],
+        on_span: Callable[[dict[str, Any]], None] | None = None,
+        service_name: str | None = None,
+        on_end: Callable[[Span], None] | None = None,
+    ) -> None:
+        self.get_response = get_response
+        self._on_span = on_span
+        self._service_name = service_name
+        self._on_end = on_end
+
+    def __call__(self, request: Any) -> Any:
+        method = getattr(request, "method", "UNKNOWN")
+        path = getattr(request, "path", "/")
+        meta = getattr(request, "META", {})
+        query = meta.get("QUERY_STRING", "")
+
+        # Build full URL from Django request
+        scheme = getattr(request, "scheme", "http")
+        host = meta.get("HTTP_HOST", "localhost")
+        full_url = build_url(scheme, host, path, query)
+
+        # Extract request headers and body
+        req_headers = _extract_django_headers(meta)
+        req_content_type = req_headers.get("content-type", "")
+        req_body = getattr(request, "body", b"")
+
+        span = Span(
+            name=f"{method} {path}",
+            kind="SERVER",
+            service_name=self._service_name,
+            on_end=self._on_end or on_span_end,
+        )
+        span.set_attribute("http.route", path)
+        span.set_attribute("http.query", query)
+
+        # AR3: Export pending_span immediately for real-time dashboard
+        on_span_start(span)
+
+        with _span_context(span):
+            try:
+                response = self.get_response(request)
+
+                # Extract response data
+                resp_headers_dict: dict[str, str] = {}
+                try:
+                    for name, value in response.items():
+                        resp_headers_dict[str(name).lower()] = str(value)
+                except Exception:
+                    pass
+                resp_content_type = resp_headers_dict.get("content-type", "")
+                resp_body = getattr(response, "content", b"")
+
+                # Capture request data (FR6)
+                capture_request_data(
+                    span,
+                    method=method,
+                    url=full_url,
+                    headers=req_headers,
+                    body=req_body,
+                    content_type=req_content_type,
+                    query_params=query,
+                )
+
+                # Capture response data (FR7)
+                capture_response_data(
+                    span,
+                    status_code=getattr(response, "status_code", 0),
+                    headers=resp_headers_dict,
+                    body=resp_body,
+                    content_type=resp_content_type,
+                )
+
+                return response
+            except Exception as exc:
+                span.set_status("ERROR", str(exc))
+                span.set_attribute("error", "true")
+                span.set_attribute("error.type", type(exc).__name__)
+                span.set_attribute("error.message", str(exc))
+                raise
+            finally:
+                span.end()
+                if self._on_span is not None:
+                    try:
+                        self._on_span(span.to_dict())
+                    except Exception:
+                        logger.debug("Error in on_span callback", exc_info=True)
+
+
+class DjangoInstrumentor(BaseInstrumentor):
+    """Instruments Django applications with middleware."""
+
+    def __init__(self, framework_info: Any) -> None:
+        super().__init__(framework_info)
+        self._active = False
+
+    def activate(self) -> None:
+        self._active = True
+        logger.info("TRACELY: Django instrumentation activated")
+
+    def deactivate(self) -> None:
+        self._active = False
+        logger.debug("TRACELY: Django instrumentation deactivated")
+
+    @property
+    def is_active(self) -> bool:
+        return self._active

tracely/instrumentation/fastapi_inst.py

@@ -0,0 +1,203 @@
+"""FastAPI auto-instrumentation (ASGI middleware).
+
+Wraps FastAPI (or any ASGI) applications with middleware that creates
+structured Span objects for each HTTP request, with full trace hierarchy
+support via context propagation. Captures full request/response data (FR6/FR7).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable
+
+from tracely.capture import build_url, capture_request_data, capture_response_data
+from tracely.context import _span_context
+from tracely.instrumentation.base import BaseInstrumentor
+from tracely.span import Span
+from tracely.span_processor import on_span_end, on_span_start
+
+logger = logging.getLogger("tracely")
+
+# Type aliases for ASGI protocol
+Scope = dict[str, Any]
+Receive = Callable[..., Any]
+Send = Callable[..., Any]
+ASGIApp = Callable[..., Any]
+
+
+def _extract_header_value(
+    headers: list[tuple[bytes, bytes]], name: bytes,
+) -> str:
+    """Extract a single header value from ASGI header list (case-insensitive)."""
+    name_lower = name.lower()
+    for hdr_name, hdr_value in headers:
+        if hdr_name.lower() == name_lower:
+            return hdr_value.decode("utf-8", errors="replace")
+    return ""
+
+
+class TracelyASGIMiddleware:
+    """ASGI middleware that creates root spans for HTTP requests.
+
+    Creates a Span object with trace_id and span_id, sets it as the
+    active span via context propagation, and captures HTTP attributes
+    including full request/response data (headers, body, URL).
+    Non-HTTP scopes (lifespan, websocket) pass through untouched.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        on_span: Callable[[dict[str, Any]], None] | None = None,
+        service_name: str | None = None,
+        on_end: Callable[[Span], None] | None = None,
+    ) -> None:
+        self.app = app
+        self._on_span = on_span
+        self._service_name = service_name
+        self._on_end = on_end
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        method = scope.get("method", "UNKNOWN")
+        path = scope.get("path", "/")
+        query = scope.get("query_string", b"").decode("utf-8", errors="replace")
+        raw_headers: list[tuple[bytes, bytes]] = scope.get("headers", [])
+
+        # Build full URL from ASGI scope
+        scheme = scope.get("scheme", "http")
+        server = scope.get("server")
+        if server:
+            host = server[0]
+            port = server[1]
+            # Include port only if non-standard
+            if (scheme == "https" and port != 443) or (scheme == "http" and port != 80):
+                host = f"{host}:{port}"
+        else:
+            host = "localhost"
+        full_url = build_url(scheme, host, path, query)
+
+        # Extract content type from request headers
+        req_content_type = _extract_header_value(raw_headers, b"content-type")
+
+        span = Span(
+            name=f"{method} {path}",
+            kind="SERVER",
+            service_name=self._service_name,
+            on_end=self._on_end or on_span_end,
+        )
+        span.set_attribute("http.route", path)
+        span.set_attribute("http.query", query)
+
+        # AR3: Export pending_span immediately for real-time dashboard
+        on_span_start(span)
+
+        # Buffer request body so the app can still read it
+        request_body = b""
+        body_consumed = False
+
+        async def receive_wrapper() -> dict[str, Any]:
+            nonlocal request_body, body_consumed
+            message = await receive()
+            if message.get("type") == "http.request" and not body_consumed:
+                chunk = message.get("body", b"")
+                request_body += chunk
+                if not message.get("more_body", False):
+                    body_consumed = True
+            return message
+
+        # Capture response data via send wrapper
+        status_code = 0
+        response_headers: list[tuple[bytes, bytes]] = []
+        response_body_chunks: list[bytes] = []
+        resp_content_type = ""
+
+        async def send_wrapper(message: dict[str, Any]) -> None:
+            nonlocal status_code, response_headers, resp_content_type
+            if message.get("type") == "http.response.start":
+                status_code = message.get("status", 0)
+                response_headers = message.get("headers", [])
+                resp_content_type = _extract_header_value(response_headers, b"content-type")
+            elif message.get("type") == "http.response.body":
+                chunk = message.get("body", b"")
+                if chunk:
+                    response_body_chunks.append(chunk)
+            await send(message)
+
+        with _span_context(span):
+            try:
+                await self.app(scope, receive_wrapper, send_wrapper)
+            except Exception as exc:
+                span.set_status("ERROR", str(exc))
+                span.set_attribute("error", "true")
+                span.set_attribute("error.type", type(exc).__name__)
+                span.set_attribute("error.message", str(exc))
+                raise
+            finally:
+                # Capture request data (FR6)
+                capture_request_data(
+                    span,
+                    method=method,
+                    url=full_url,
+                    headers=raw_headers,
+                    body=request_body,
+                    content_type=req_content_type,
+                    query_params=query,
+                )
+
+                # Capture response data (FR7)
+                response_body = b"".join(response_body_chunks)
+                capture_response_data(
+                    span,
+                    status_code=status_code,
+                    headers=response_headers,
+                    body=response_body,
+                    content_type=resp_content_type,
+                )
+
+                span.end()
+                if self._on_span is not None:
+                    try:
+                        self._on_span(span.to_dict())
+                    except Exception:
+                        logger.debug("Error in on_span callback", exc_info=True)
+
+
+class FastAPIInstrumentor(BaseInstrumentor):
+    """Instruments FastAPI applications with ASGI middleware.
+
+    On activate(), registers the middleware factory so it can be applied
+    to FastAPI apps. The actual wrapping happens when the user's app
+    is detected or when middleware is explicitly applied.
+    """
+
+    def __init__(self, framework_info: Any) -> None:
+        super().__init__(framework_info)
+        self._active = False
+
+    def activate(self) -> None:
+        self._active = True
+        logger.info("TRACELY: FastAPI instrumentation activated")
+
+    def deactivate(self) -> None:
+        self._active = False
+        logger.debug("TRACELY: FastAPI instrumentation deactivated")
+
+    @property
+    def is_active(self) -> bool:
+        return self._active
+
+    @staticmethod
+    def wrap_app(
+        app: ASGIApp,
+        on_span: Callable[[dict[str, Any]], None] | None = None,
+        service_name: str | None = None,
+        on_end: Callable[[Span], None] | None = None,
+    ) -> TracelyASGIMiddleware:
+        """Wrap an ASGI app with TRACELY middleware."""
+        return TracelyASGIMiddleware(
+            app=app, on_span=on_span, service_name=service_name, on_end=on_end,
+        )