justanalytics-python 0.1.0__py3-none-any.whl

justanalytics/client.py

@@ -0,0 +1,665 @@
+"""
+Core Client class managing SDK lifecycle, span creation, error capture,
+context propagation, and batched transport.
+
+The client is the central orchestrator:
+- Validates and stores configuration from ``init()``
+- Creates spans via ``start_span()`` (context manager / decorator)
+- Captures errors via ``capture_exception()`` and ``capture_message()``
+- Manages the BatchTransport for periodic flushing
+- Provides ``set_user()``, ``set_tag()`` context APIs
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import platform
+import sys
+import traceback
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from . import context as ctx
+from .span import Span, _generate_span_id, _generate_trace_id
+from .transport import BatchTransport
+from .types import (
+    ErrorLevel,
+    ErrorPayload,
+    LogLevel,
+    LogPayload,
+    MechanismType,
+    MetricPayload,
+    SpanKind,
+    SpanStatus,
+    UserContext,
+)
+
+sdk_logger = logging.getLogger("justanalytics")
+
+# Default JustAnalytics server URL
+DEFAULT_SERVER_URL = "https://justanalytics.up.railway.app"
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class JustAnalyticsClient:
+    """
+    Core SDK client managing lifecycle, instrumentation, and transport.
+
+    Typically used as a singleton via the module-level functions in
+    ``justanalytics.__init__``.
+    """
+
+    def __init__(self) -> None:
+        self._initialized: bool = False
+        self._enabled: bool = True
+        self._service_name: str = ""
+        self._environment: Optional[str] = None
+        self._release: Optional[str] = None
+        self._debug: bool = False
+        self._transport: Optional[BatchTransport] = None
+        self._site_id: str = ""
+
+    def init(
+        self,
+        site_id: str,
+        api_key: str,
+        service_name: str,
+        environment: Optional[str] = None,
+        release: Optional[str] = None,
+        server_url: Optional[str] = None,
+        debug: bool = False,
+        flush_interval_s: float = 2.0,
+        max_batch_size: int = 100,
+        enabled: bool = True,
+    ) -> None:
+        """
+        Initialize the JustAnalytics SDK.
+
+        Must be called before any other SDK method. Calling ``init()`` more
+        than once logs a warning and is ignored.
+
+        Args:
+            site_id: Site ID from the JustAnalytics dashboard.
+            api_key: API key (format: ``ja_sk_...``).
+            service_name: Service name (e.g., "api-server", "worker").
+            environment: Deployment environment (e.g., "production", "staging").
+            release: Release/version string (e.g., "1.2.3", git SHA).
+            server_url: Base URL of JustAnalytics server (default: production URL).
+            debug: Enable debug logging to the ``justanalytics`` logger.
+            flush_interval_s: Flush interval in seconds (default: 2.0).
+            max_batch_size: Max items per batch before immediate flush (default: 100).
+            enabled: Enable/disable the SDK (default: True). When False, all methods are no-ops.
+
+        Raises:
+            ValueError: If required fields (site_id, api_key, service_name) are missing.
+        """
+        if self._initialized:
+            sdk_logger.warning("init() already called. Ignoring.")
+            return
+
+        if not site_id:
+            raise ValueError(
+                "init() requires 'site_id'. Get it from your JustAnalytics dashboard."
+            )
+        if not api_key:
+            raise ValueError(
+                "init() requires 'api_key'. Create one in your JustAnalytics site settings."
+            )
+        if not service_name:
+            raise ValueError(
+                "init() requires 'service_name' (e.g., 'api-server', 'worker')."
+            )
+
+        self._enabled = enabled
+        self._service_name = service_name
+        self._environment = environment
+        self._release = release
+        self._debug = debug
+        self._site_id = site_id
+
+        resolved_url = (
+            server_url
+            or os.environ.get("JUSTANALYTICS_URL")
+            or DEFAULT_SERVER_URL
+        )
+
+        if self._debug:
+            key_prefix = api_key[:10] if len(api_key) >= 10 else api_key
+            sdk_logger.debug(
+                "Initializing SDK: service=%s, server=%s, apiKey=%s..., enabled=%s",
+                service_name,
+                resolved_url,
+                key_prefix,
+                enabled,
+            )
+
+        if self._enabled:
+            self._transport = BatchTransport(
+                server_url=resolved_url,
+                api_key=api_key,
+                site_id=site_id,
+                flush_interval_s=flush_interval_s,
+                max_batch_size=max_batch_size,
+                debug=debug,
+            )
+            self._transport.start()
+
+        self._initialized = True
+
+    @property
+    def is_initialized(self) -> bool:
+        """Whether ``init()`` has been called successfully."""
+        return self._initialized
+
+    # --- Span Creation ---
+
+    @contextmanager
+    def start_span(
+        self,
+        name: str,
+        op: Optional[str] = None,
+        kind: SpanKind = SpanKind.INTERNAL,
+        attributes: Optional[Dict[str, Any]] = None,
+    ) -> Generator[Span, None, None]:
+        """
+        Create a span as a context manager.
+
+        The span is automatically ended when the ``with`` block exits.
+        If an exception occurs, the span status is set to ``error``.
+
+        Nested ``start_span()`` calls automatically form parent-child relationships
+        via ``contextvars``.
+
+        Args:
+            name: Operation name for the span.
+            op: Optional operation type (alias for backwards compat; sets attribute).
+            kind: Span kind (default: INTERNAL).
+            attributes: Initial attributes to set on the span.
+
+        Yields:
+            The created Span instance.
+
+        Example::
+
+            with justanalytics.start_span("process-order") as span:
+                span.set_attribute("order.id", "12345")
+                result = process_order()
+        """
+        if not self._initialized or not self._enabled:
+            # Yield a dummy span that does nothing
+            dummy = Span(
+                operation_name=name,
+                service_name=self._service_name or "unknown",
+                kind=kind,
+            )
+            yield dummy
+            return
+
+        # Determine parent and trace ID from current context
+        parent_span = ctx.get_active_span()
+        if parent_span:
+            trace_id = parent_span.trace_id
+            parent_span_id = parent_span.id
+        else:
+            trace_id = _generate_trace_id()
+            parent_span_id = None
+
+        # Merge context tags into initial attributes
+        context_tags = ctx.get_tags()
+        initial_attrs: Dict[str, Any] = {**context_tags}
+        if self._environment:
+            initial_attrs["environment"] = self._environment
+        if self._release:
+            initial_attrs["release"] = self._release
+        if op:
+            initial_attrs["op"] = op
+        if attributes:
+            initial_attrs.update(attributes)
+
+        span = Span(
+            operation_name=name,
+            service_name=self._service_name,
+            kind=kind,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+            attributes=initial_attrs,
+        )
+
+        # Use the span as context manager (sets contextvars)
+        with span:
+            yield span
+
+        # Enqueue the ended span
+        if self._transport:
+            self._transport.enqueue_span(span.to_dict())
+
+    def span_decorator(
+        self,
+        name: Optional[str] = None,
+        op: Optional[str] = None,
+        kind: SpanKind = SpanKind.INTERNAL,
+        attributes: Optional[Dict[str, Any]] = None,
+    ) -> Callable[[F], F]:
+        """
+        Decorator that wraps a function in a span.
+
+        Args:
+            name: Span name (defaults to the function's qualified name).
+            op: Optional operation type.
+            kind: Span kind.
+            attributes: Initial span attributes.
+
+        Returns:
+            A decorator that wraps the function in a span.
+
+        Example::
+
+            @justanalytics.span_decorator(op="db.query")
+            def get_user(user_id: str):
+                return db.query(...)
+        """
+
+        def decorator(fn: F) -> F:
+            span_name = name or fn.__qualname__
+
+            @wraps(fn)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                with self.start_span(span_name, op=op, kind=kind, attributes=attributes):
+                    return fn(*args, **kwargs)
+
+            return wrapper  # type: ignore[return-value]
+
+        return decorator
+
+    # --- Error Capture ---
+
+    def capture_exception(
+        self,
+        error: BaseException,
+        tags: Optional[Dict[str, str]] = None,
+        extra: Optional[Dict[str, Any]] = None,
+        user: Optional[UserContext] = None,
+        level: ErrorLevel = ErrorLevel.ERROR,
+        fingerprint: Optional[List[str]] = None,
+    ) -> str:
+        """
+        Capture an exception and send it to JustAnalytics.
+
+        Automatically attaches the current trace_id, span_id, user context,
+        and tags from the contextvars context.
+
+        Args:
+            error: The exception to capture.
+            tags: Additional tags (merged with context tags).
+            extra: Arbitrary extra data.
+            user: User context (overrides context user).
+            level: Severity level (default: ERROR).
+            fingerprint: Custom fingerprint for grouping.
+
+        Returns:
+            A unique event_id for correlation, or empty string if SDK is disabled.
+
+        Example::
+
+            try:
+                risky_operation()
+            except Exception as e:
+                justanalytics.capture_exception(e, tags={"module": "payments"})
+        """
+        if not self._initialized or not self._enabled:
+            return ""
+        try:
+            payload = self._build_error_payload(
+                error=error,
+                tags=tags,
+                extra=extra,
+                user=user,
+                level=level,
+                fingerprint=fingerprint,
+                mechanism_type=MechanismType.MANUAL,
+                handled=True,
+            )
+            if self._transport:
+                self._transport.enqueue_error(payload.to_dict())
+            return payload.event_id
+        except Exception:
+            if self._debug:
+                sdk_logger.debug("capture_exception() internal error", exc_info=True)
+            return ""
+
+    def capture_message(
+        self,
+        message: str,
+        level: ErrorLevel = ErrorLevel.INFO,
+        tags: Optional[Dict[str, str]] = None,
+        extra: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        """
+        Capture a message and send it to JustAnalytics.
+
+        Args:
+            message: The message string.
+            level: Severity level (default: INFO).
+            tags: Additional tags (merged with context tags).
+            extra: Arbitrary extra data.
+
+        Returns:
+            A unique event_id for correlation, or empty string if SDK is disabled.
+
+        Example::
+
+            justanalytics.capture_message(
+                "User exceeded rate limit",
+                level=ErrorLevel.WARNING,
+                tags={"userId": user.id},
+            )
+        """
+        if not self._initialized or not self._enabled:
+            return ""
+        try:
+            # Create a synthetic error-like payload for the message
+            payload = self._build_error_payload(
+                error=message,
+                tags=tags,
+                extra=extra,
+                user=None,
+                level=level,
+                fingerprint=None,
+                mechanism_type=MechanismType.MANUAL,
+                handled=True,
+            )
+            # Override type to 'Message' for capture_message
+            payload.error["type"] = "Message"
+            if self._transport:
+                self._transport.enqueue_error(payload.to_dict())
+            return payload.event_id
+        except Exception:
+            if self._debug:
+                sdk_logger.debug("capture_message() internal error", exc_info=True)
+            return ""
+
+    # --- Context ---
+
+    def set_user(
+        self,
+        id: Optional[str] = None,
+        email: Optional[str] = None,
+        username: Optional[str] = None,
+    ) -> None:
+        """
+        Set user context for the current execution scope.
+
+        User information is attached to all spans, errors, and logs
+        created within the current contextvars scope.
+
+        Args:
+            id: User ID.
+            email: User email.
+            username: Username.
+
+        Example::
+
+            justanalytics.set_user(id="user-123", email="alice@example.com")
+        """
+        if not self._initialized or not self._enabled:
+            return
+        ctx.set_user(UserContext(id=id, email=email, username=username))
+
+    def set_tag(self, key: str, value: str) -> None:
+        """
+        Set a tag for the current execution scope.
+
+        Tags are attached as attributes on all spans created within the
+        current contextvars scope.
+
+        Args:
+            key: Tag key.
+            value: Tag value.
+
+        Example::
+
+            justanalytics.set_tag("feature", "checkout")
+            justanalytics.set_tag("region", "us-east-1")
+        """
+        if not self._initialized or not self._enabled:
+            return
+        ctx.set_tag(key, value)
+
+    # --- Logging ---
+
+    def log(
+        self,
+        level: LogLevel,
+        message: str,
+        attributes: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """
+        Send a structured log entry to JustAnalytics.
+
+        Automatically attaches trace_id and span_id from the current context.
+
+        Args:
+            level: Log severity level.
+            message: Log message.
+            attributes: Optional key-value metadata.
+        """
+        if not self._initialized or not self._enabled or not self._transport:
+            return
+        try:
+            active_span = ctx.get_active_span()
+            payload = LogPayload(
+                level=level.value,
+                message=message,
+                service_name=self._service_name,
+                timestamp=datetime.now(timezone.utc).isoformat(),
+                trace_id=active_span.trace_id if active_span else None,
+                span_id=active_span.id if active_span else None,
+                attributes=attributes or {},
+            )
+            self._transport.enqueue_log(payload.to_dict())
+        except Exception:
+            if self._debug:
+                sdk_logger.debug("log() internal error", exc_info=True)
+
+    # --- Metrics ---
+
+    def record_metric(
+        self,
+        metric_name: str,
+        value: float,
+        tags: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """
+        Record a custom infrastructure metric.
+
+        Args:
+            metric_name: Metric name using dot notation (e.g., "custom.queue_size").
+            value: Numeric value.
+            tags: Optional additional tags.
+
+        Example::
+
+            justanalytics.record_metric("custom.queue_size", 42, {"queue": "emails"})
+        """
+        if not self._initialized or not self._enabled or not self._transport:
+            return
+        try:
+            metric_tags: Dict[str, Any] = {
+                "hostname": platform.node(),
+            }
+            if self._environment:
+                metric_tags["environment"] = self._environment
+            if tags:
+                metric_tags.update(tags)
+            payload = MetricPayload(
+                metric_name=metric_name,
+                value=value,
+                service_name=self._service_name,
+                timestamp=datetime.now(timezone.utc).isoformat(),
+                tags=metric_tags,
+            )
+            self._transport.enqueue_metric(payload.to_dict())
+        except Exception:
+            if self._debug:
+                sdk_logger.debug("record_metric() internal error", exc_info=True)
+
+    # --- Lifecycle ---
+
+    def flush(self) -> None:
+        """
+        Manually flush all pending data to the server.
+
+        Useful before a serverless function completes or in tests.
+        """
+        if self._transport:
+            self._transport.flush()
+
+    def close(self) -> None:
+        """
+        Shut down the SDK: flush remaining data, stop timers.
+
+        After calling ``close()``, the SDK cannot be used again until
+        ``init()`` is called.
+        """
+        if not self._initialized:
+            return
+
+        if self._debug:
+            sdk_logger.debug("Closing SDK...")
+
+        if self._transport:
+            self._transport.flush()
+            self._transport.stop()
+            self._transport = None
+
+        self._initialized = False
+
+        if self._debug:
+            sdk_logger.debug("SDK closed.")
+
+    # --- Internal: Error Payload Construction ---
+
+    def _build_error_payload(
+        self,
+        error: Union[BaseException, str],
+        tags: Optional[Dict[str, str]],
+        extra: Optional[Dict[str, Any]],
+        user: Optional[UserContext],
+        level: ErrorLevel,
+        fingerprint: Optional[List[str]],
+        mechanism_type: MechanismType,
+        handled: bool,
+    ) -> ErrorPayload:
+        """Build a complete error payload from an error and options."""
+        # Serialize the error
+        if isinstance(error, BaseException):
+            error_dict = {
+                "message": str(error) or type(error).__name__,
+                "type": type(error).__name__,
+                "stack": "".join(traceback.format_exception(type(error), error, error.__traceback__)),
+            }
+            # Extract cause chain
+            cause_chain = self._extract_cause_chain(error, max_depth=3)
+        elif isinstance(error, str):
+            error_dict = {
+                "message": error,
+                "type": "Error",
+                "stack": None,
+            }
+            cause_chain = None
+        else:
+            error_dict = {
+                "message": str(error),
+                "type": "Error",
+                "stack": None,
+            }
+            cause_chain = None
+
+        # Get trace context
+        active_span = ctx.get_active_span()
+        trace_id = active_span.trace_id if active_span else None
+        span_id = active_span.id if active_span else None
+
+        # Get user context (parameter overrides context)
+        resolved_user = user or ctx.get_user()
+        user_dict = resolved_user.to_dict() if resolved_user else None
+
+        # Merge context tags with option tags
+        context_tags = ctx.get_tags()
+        merged_tags = {**context_tags, **(tags or {})}
+
+        # Merge extra with cause chain
+        merged_extra: Dict[str, Any] = dict(extra) if extra else {}
+        if cause_chain:
+            merged_extra["error.cause"] = cause_chain
+
+        event_id = os.urandom(12).hex()
+
+        return ErrorPayload(
+            event_id=event_id,
+            timestamp=datetime.now(timezone.utc).isoformat(),
+            error=error_dict,
+            level=level.value,
+            mechanism={
+                "type": mechanism_type.value,
+                "handled": handled,
+            },
+            context={
+                "serviceName": self._service_name,
+                "environment": self._environment,
+                "release": self._release,
+                "runtime": "python",
+                "nodeVersion": platform.python_version(),
+            },
+            trace={"traceId": trace_id, "spanId": span_id} if trace_id else None,
+            user=user_dict,
+            tags=merged_tags if merged_tags else None,
+            extra=merged_extra if merged_extra else None,
+            fingerprint=fingerprint,
+        )
+
+    @staticmethod
+    def _extract_cause_chain(
+        error: BaseException, max_depth: int = 3
+    ) -> Optional[List[Dict[str, Any]]]:
+        """Extract the __cause__ chain from an exception."""
+        chain: List[Dict[str, Any]] = []
+        seen: set = set()
+        seen.add(id(error))
+        current: Optional[BaseException] = error.__cause__
+
+        for _ in range(max_depth):
+            if current is None:
+                break
+            if id(current) in seen:
+                chain.append({"message": "[Circular]", "type": "Circular", "stack": None})
+                break
+            seen.add(id(current))
+            chain.append(
+                {
+                    "message": str(current),
+                    "type": type(current).__name__,
+                    "stack": "".join(
+                        traceback.format_exception(
+                            type(current), current, current.__traceback__
+                        )
+                    ),
+                }
+            )
+            current = current.__cause__
+
+        return chain if chain else None