justanalytics-python 0.1.0__py3-none-any.whl

justanalytics/span.py

@@ -0,0 +1,281 @@
+"""
+Span class for distributed tracing with high-resolution timing.
+
+Each Span records a timed operation within a trace. Spans support:
+- Context manager protocol (``with`` statement)
+- Decorator usage (``@span``)
+- Manual start/end via ``span.end()``
+
+The ``to_dict()`` method serializes the span to the exact format expected by
+``POST /api/ingest/spans`` (matching the Zod schema in the server route).
+
+Timing uses ``time.perf_counter_ns()`` for nanosecond-resolution monotonic
+measurement, converted to integer milliseconds via rounding.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    TypeVar,
+    overload,
+)
+
+from . import context as ctx
+from .types import SpanEvent, SpanKind, SpanPayload, SpanStatus
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _generate_trace_id() -> str:
+    """Generate a 32-char lowercase hex trace ID (128 bits)."""
+    return os.urandom(16).hex()
+
+
+def _generate_span_id() -> str:
+    """Generate a 16-char lowercase hex span ID (64 bits)."""
+    return os.urandom(8).hex()
+
+
+class Span:
+    """
+    Represents a single timed operation within a distributed trace.
+
+    Spans are typically created via ``justanalytics.start_span()`` and ended
+    automatically when the context manager exits. They can also be ended
+    manually via ``span.end()``.
+
+    Example::
+
+        with justanalytics.start_span("process-order") as span:
+            span.set_attribute("order.id", "12345")
+            process_order()
+        # span is automatically ended here
+
+    Attributes:
+        id: 16-char hex span ID.
+        trace_id: 32-char hex trace ID (shared across all spans in a trace).
+        parent_span_id: Parent span ID, or None for root spans.
+        operation_name: Human-readable name for this operation.
+        service_name: Name of the service producing this span.
+        kind: The span kind (client, server, internal, etc.).
+        start_time: Wall-clock start time (UTC datetime).
+    """
+
+    __slots__ = (
+        "id",
+        "trace_id",
+        "parent_span_id",
+        "operation_name",
+        "service_name",
+        "kind",
+        "start_time",
+        "_start_perf_ns",
+        "_end_time",
+        "_duration_ms",
+        "_status",
+        "_status_message",
+        "_attributes",
+        "_events",
+        "_ended",
+        "_span_token",
+        "_trace_token",
+    )
+
+    def __init__(
+        self,
+        operation_name: str,
+        service_name: str,
+        kind: SpanKind = SpanKind.INTERNAL,
+        trace_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+        attributes: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self.id: str = _generate_span_id()
+        self.trace_id: str = trace_id or _generate_trace_id()
+        self.parent_span_id: Optional[str] = parent_span_id
+        self.operation_name: str = operation_name
+        self.service_name: str = service_name
+        self.kind: SpanKind = kind
+        self.start_time: datetime = datetime.now(timezone.utc)
+        self._start_perf_ns: int = time.perf_counter_ns()
+        self._end_time: Optional[datetime] = None
+        self._duration_ms: Optional[int] = None
+        self._status: SpanStatus = SpanStatus.UNSET
+        self._status_message: Optional[str] = None
+        self._attributes: Dict[str, Any] = dict(attributes) if attributes else {}
+        self._events: List[SpanEvent] = []
+        self._ended: bool = False
+        self._span_token: Any = None
+        self._trace_token: Any = None
+
+    # --- Attribute Methods ---
+
+    def set_attribute(self, key: str, value: Any) -> "Span":
+        """Set a single attribute on the span.
+
+        Args:
+            key: Attribute key (e.g., "http.method", "db.statement").
+            value: Attribute value.
+
+        Returns:
+            self, for chaining.
+        """
+        if not self._ended:
+            self._attributes[key] = value
+        return self
+
+    def set_attributes(self, attrs: Dict[str, Any]) -> "Span":
+        """Set multiple attributes on the span.
+
+        Args:
+            attrs: Key-value pairs to merge into attributes.
+
+        Returns:
+            self, for chaining.
+        """
+        if not self._ended:
+            self._attributes.update(attrs)
+        return self
+
+    # --- Status Methods ---
+
+    def set_status(self, status: SpanStatus, message: Optional[str] = None) -> "Span":
+        """Set the span's status.
+
+        Args:
+            status: One of SpanStatus.OK, SpanStatus.ERROR, SpanStatus.UNSET.
+            message: Optional status message (typically for errors).
+
+        Returns:
+            self, for chaining.
+        """
+        if not self._ended:
+            self._status = status
+            self._status_message = message
+        return self
+
+    # --- Event Methods ---
+
+    def add_event(
+        self, name: str, attributes: Optional[Dict[str, Any]] = None
+    ) -> "Span":
+        """Add a timestamped event to the span.
+
+        Events represent notable moments during the span's lifetime
+        (e.g., "cache.miss", "retry.attempt", "payment.processed").
+
+        Args:
+            name: Event name.
+            attributes: Optional event attributes.
+
+        Returns:
+            self, for chaining.
+        """
+        if not self._ended:
+            event = SpanEvent(
+                name=name,
+                timestamp=datetime.now(timezone.utc).isoformat(),
+                attributes=attributes,
+            )
+            self._events.append(event)
+        return self
+
+    # --- Lifecycle ---
+
+    def end(self) -> None:
+        """Mark the span as ended.
+
+        Calculates duration using ``time.perf_counter_ns()`` for sub-millisecond
+        precision, rounded to integer milliseconds.
+
+        Calling ``end()`` multiple times is a no-op (idempotent).
+        """
+        if self._ended:
+            return
+        self._ended = True
+        self._end_time = datetime.now(timezone.utc)
+        elapsed_ns = time.perf_counter_ns() - self._start_perf_ns
+        self._duration_ms = round(elapsed_ns / 1_000_000)
+
+    @property
+    def is_ended(self) -> bool:
+        """Whether this span has been ended."""
+        return self._ended
+
+    # --- Context Manager Protocol ---
+
+    def __enter__(self) -> "Span":
+        """Enter the span context: set this span as active in contextvars."""
+        self._span_token = ctx.set_active_span(self)
+        self._trace_token = ctx.set_trace_id(self.trace_id)
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit the span context: end the span and restore previous context."""
+        if exc_val is not None and not self._ended:
+            msg = str(exc_val) if exc_val else (exc_type.__name__ if exc_type else None)
+            self.set_status(SpanStatus.ERROR, msg)
+        if not self._ended:
+            self.end()
+        # Restore previous context
+        if self._span_token is not None:
+            ctx._active_span_var.reset(self._span_token)
+            self._span_token = None
+        if self._trace_token is not None:
+            ctx._trace_id_var.reset(self._trace_token)
+            self._trace_token = None
+
+    # --- Serialization ---
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize to the wire protocol dict matching the server's Zod schema."""
+        return {
+            "id": self.id,
+            "traceId": self.trace_id,
+            "parentSpanId": self.parent_span_id,
+            "operationName": self.operation_name,
+            "serviceName": self.service_name,
+            "kind": self.kind.value,
+            "startTime": self.start_time.isoformat(),
+            "endTime": self._end_time.isoformat() if self._end_time else None,
+            "duration": self._duration_ms,
+            "status": self._status.value,
+            "statusMessage": self._status_message,
+            "attributes": dict(self._attributes),
+            "events": [e.to_dict() for e in self._events],
+        }
+
+    def to_payload(self) -> SpanPayload:
+        """Serialize to a SpanPayload dataclass."""
+        return SpanPayload(
+            id=self.id,
+            trace_id=self.trace_id,
+            parent_span_id=self.parent_span_id,
+            operation_name=self.operation_name,
+            service_name=self.service_name,
+            kind=self.kind.value,
+            start_time=self.start_time.isoformat(),
+            end_time=self._end_time.isoformat() if self._end_time else None,
+            duration=self._duration_ms,
+            status=self._status.value,
+            status_message=self._status_message,
+            attributes=dict(self._attributes),
+            events=[e.to_dict() for e in self._events],
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"Span(name={self.operation_name!r}, id={self.id!r}, "
+            f"trace_id={self.trace_id!r}, ended={self._ended})"
+        )

justanalytics/trace_context.py

@@ -0,0 +1,124 @@
+"""
+W3C Trace Context (traceparent) header parsing and serialization.
+
+Format: "00-{traceId}-{spanId}-{flags}"
+- version: "00" (only supported version)
+- traceId: 32-char lowercase hex (128-bit, must not be all zeros)
+- spanId: 16-char lowercase hex (64-bit, must not be all zeros)
+- flags: 2-char hex ("01" = sampled, "00" = not sampled)
+
+References:
+- W3C Trace Context: https://www.w3.org/TR/trace-context/
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Optional
+
+# Validation patterns
+_TRACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
+_SPAN_ID_RE = re.compile(r"^[0-9a-f]{16}$")
+_FLAGS_RE = re.compile(r"^[0-9a-f]{2}$")
+
+# Invalid all-zero IDs per W3C spec
+_ZERO_TRACE_ID = "0" * 32
+_ZERO_SPAN_ID = "0" * 16
+
+
+@dataclass(frozen=True)
+class TraceparentData:
+    """Parsed W3C traceparent header data."""
+
+    version: str
+    trace_id: str
+    parent_span_id: str
+    trace_flags: str
+
+    @property
+    def sampled(self) -> bool:
+        """Whether the trace is sampled (flags bit 0 set)."""
+        return int(self.trace_flags, 16) & 0x01 == 1
+
+
+def parse_traceparent(header: str) -> Optional[TraceparentData]:
+    """
+    Parse a W3C traceparent header string.
+
+    Returns None for invalid or malformed headers (never raises).
+
+    Args:
+        header: The traceparent header string to parse.
+
+    Returns:
+        Parsed TraceparentData, or None if invalid.
+
+    Example::
+
+        data = parse_traceparent("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
+        assert data.trace_id == "4bf92f3577b34da6a3ce929d0e0e4736"
+        assert data.parent_span_id == "00f067aa0ba902b7"
+        assert data.sampled is True
+    """
+    if not header or not isinstance(header, str):
+        return None
+
+    parts = header.split("-")
+    if len(parts) != 4:
+        return None
+
+    version, trace_id, parent_span_id, trace_flags = parts
+
+    # Only version "00" is supported
+    if version != "00":
+        return None
+
+    # Validate traceId: 32 hex chars, not all zeros
+    if not _TRACE_ID_RE.match(trace_id) or trace_id == _ZERO_TRACE_ID:
+        return None
+
+    # Validate parentSpanId: 16 hex chars, not all zeros
+    if not _SPAN_ID_RE.match(parent_span_id) or parent_span_id == _ZERO_SPAN_ID:
+        return None
+
+    # Validate traceFlags: 2 hex chars
+    if not _FLAGS_RE.match(trace_flags):
+        return None
+
+    return TraceparentData(
+        version=version,
+        trace_id=trace_id,
+        parent_span_id=parent_span_id,
+        trace_flags=trace_flags,
+    )
+
+
+def serialize_traceparent(
+    trace_id: str,
+    span_id: str,
+    sampled: bool = True,
+) -> str:
+    """
+    Serialize a W3C traceparent header string.
+
+    Produces: "00-{traceId}-{spanId}-{flags}"
+
+    Args:
+        trace_id: 32-character hex trace ID.
+        span_id: 16-character hex span ID.
+        sampled: Whether the trace is sampled (default True).
+
+    Returns:
+        W3C traceparent header string.
+
+    Example::
+
+        header = serialize_traceparent(
+            "4bf92f3577b34da6a3ce929d0e0e4736",
+            "00f067aa0ba902b7",
+        )
+        assert header == "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
+    """
+    flags = "01" if sampled else "00"
+    return f"00-{trace_id}-{span_id}-{flags}"