tusk-drift-python-sdk 0.1.0__tar.gz

tusk_drift_python_sdk-0.1.0/MANIFEST.in

@@ -0,0 +1,5 @@
+include README.md
+include LICENSE
+include py.typed
+recursive-include drift *.py
+recursive-include drift *.typed

tusk_drift_python_sdk-0.1.0/PKG-INFO

@@ -0,0 +1,83 @@
+Metadata-Version: 2.4
+Name: tusk-drift-python-sdk
+Version: 0.1.0
+Summary: Python SDK for Tusk Drift instrumentation and replay
+Author-email: Tusk <support@usetusk.ai>
+License: MIT
+Project-URL: Homepage, https://usetusk.ai
+Project-URL: Documentation, https://docs.usetusk.ai
+Project-URL: Repository, https://github.com/Use-Tusk/drift-node-sdk
+Project-URL: Issues, https://github.com/Use-Tusk/drift-node-sdk/issues
+Keywords: tusk,drift,testing,instrumentation,tracing,replay
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Software Development :: Testing
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: protobuf>=6.0
+Requires-Dist: PyYAML>=6.0
+Requires-Dist: requests>=2.32.5
+Requires-Dist: tusk-drift-schemas>=0.1.9
+Requires-Dist: aiohttp>=3.9.0
+Requires-Dist: aiofiles>=23.0.0
+Provides-Extra: flask
+Requires-Dist: Flask>=3.1.2; extra == "flask"
+Provides-Extra: fastapi
+Requires-Dist: fastapi>=0.115.6; extra == "fastapi"
+Requires-Dist: uvicorn>=0.34.2; extra == "fastapi"
+Requires-Dist: starlette<0.42.0; extra == "fastapi"
+Provides-Extra: dev
+Requires-Dist: Flask>=3.1.2; extra == "dev"
+Requires-Dist: fastapi>=0.115.6; extra == "dev"
+Requires-Dist: uvicorn>=0.34.2; extra == "dev"
+Requires-Dist: python-jsonpath>=0.10; extra == "dev"
+
+# Drift Python SDK
+
+## Installation
+
+```bash
+python -m venv venv
+. venv/bin/activate
+pip install -r requirements.txt
+```
+
+Alternatively, if you have [direnv](https://direnv.net/), just allow this folder and run `pip install -r requirements.txt`.
+
+## Running Tests
+
+```bash
+# Run all unit tests
+python -m unittest discover -s tests/unit -v
+
+# Run all integration tests (Flask/FastAPI)
+timeout 30 python -m unittest discover -s tests/integration -v
+
+# Run specific test file
+python -m unittest tests.unit.test_json_schema_helper -v
+python -m unittest tests.unit.test_adapters -v
+```
+
+### Database Integration Tests
+
+```bash
+# Start test databases
+docker compose -f docker-compose.test.yml up -d
+
+# Run database tests
+python -m unittest tests.integration.test_database -v
+
+# Stop databases
+docker compose -f docker-compose.test.yml down
+```
+
+### Demo Scripts
+
+```bash
+timeout 10 python tests/test_flask_demo.py
+timeout 10 python tests/test_fastapi_demo.py
+```

tusk_drift_python_sdk-0.1.0/README.md

@@ -0,0 +1,45 @@
+# Drift Python SDK
+
+## Installation
+
+```bash
+python -m venv venv
+. venv/bin/activate
+pip install -r requirements.txt
+```
+
+Alternatively, if you have [direnv](https://direnv.net/), just allow this folder and run `pip install -r requirements.txt`.
+
+## Running Tests
+
+```bash
+# Run all unit tests
+python -m unittest discover -s tests/unit -v
+
+# Run all integration tests (Flask/FastAPI)
+timeout 30 python -m unittest discover -s tests/integration -v
+
+# Run specific test file
+python -m unittest tests.unit.test_json_schema_helper -v
+python -m unittest tests.unit.test_adapters -v
+```
+
+### Database Integration Tests
+
+```bash
+# Start test databases
+docker compose -f docker-compose.test.yml up -d
+
+# Run database tests
+python -m unittest tests.integration.test_database -v
+
+# Stop databases
+docker compose -f docker-compose.test.yml down
+```
+
+### Demo Scripts
+
+```bash
+timeout 10 python tests/test_flask_demo.py
+timeout 10 python tests/test_fastapi_demo.py
+```

tusk_drift_python_sdk-0.1.0/drift/__init__.py

@@ -0,0 +1,68 @@
+"""Drift Python SDK for distributed tracing and instrumentation."""
+
+from .core import (
+    TuskDrift,
+    CleanSpanData,
+    PackageType,
+    SpanKind,
+    StatusCode,
+    DriftMode,
+    BatchSpanProcessorConfig,
+    # Config
+    TuskConfig,
+    TuskFileConfig,
+    ServiceConfig,
+    RecordingConfig,
+    TracesConfig,
+    TuskApiConfig,
+    load_tusk_config,
+    find_project_root,
+)
+from .instrumentation.flask import FlaskInstrumentation
+from .instrumentation.fastapi import FastAPIInstrumentation
+from .instrumentation.requests import RequestsInstrumentation
+from .tracing.adapters import (
+    SpanExportAdapter,
+    ExportResult,
+    ExportResultCode,
+    InMemorySpanAdapter,
+    FilesystemSpanAdapter,
+    ApiSpanAdapter,
+    ApiSpanAdapterConfig,
+    create_api_adapter,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Core
+    "TuskDrift",
+    "CleanSpanData",
+    "PackageType",
+    "SpanKind",
+    "StatusCode",
+    "DriftMode",
+    "BatchSpanProcessorConfig",
+    # Config
+    "TuskConfig",
+    "TuskFileConfig",
+    "ServiceConfig",
+    "RecordingConfig",
+    "TracesConfig",
+    "TuskApiConfig",
+    "load_tusk_config",
+    "find_project_root",
+    # Instrumentations
+    "FlaskInstrumentation",
+    "FastAPIInstrumentation",
+    "RequestsInstrumentation",
+    # Adapters
+    "SpanExportAdapter",
+    "ExportResult",
+    "ExportResultCode",
+    "InMemorySpanAdapter",
+    "FilesystemSpanAdapter",
+    "ApiSpanAdapter",
+    "ApiSpanAdapterConfig",
+    "create_api_adapter",
+]

tusk_drift_python_sdk-0.1.0/drift/core/__init__.py

@@ -0,0 +1,66 @@
+"""Core module for the Drift SDK."""
+
+from .drift_sdk import TuskDrift
+from .types import DriftMode, CleanSpanData, PackageType, SpanKind, StatusCode
+from .config import (
+    TuskConfig,
+    TuskFileConfig,
+    ServiceConfig,
+    RecordingConfig,
+    TracesConfig,
+    TuskApiConfig,
+    load_tusk_config,
+    find_project_root,
+)
+from .batch_processor import BatchSpanProcessor, BatchSpanProcessorConfig
+from .sampling import should_sample, validate_sampling_rate
+from .data_normalization import (
+    normalize_input_data,
+    remove_none_values,
+    create_span_input_value,
+    create_mock_input_value,
+)
+from .trace_blocking_manager import (
+    TraceBlockingManager,
+    estimate_span_size,
+    should_block_span,
+    MAX_SPAN_SIZE_MB,
+    MAX_SPAN_SIZE_BYTES,
+)
+
+__all__ = [
+    # Main SDK
+    "TuskDrift",
+    # Config
+    "TuskConfig",
+    "TuskFileConfig",
+    "ServiceConfig",
+    "RecordingConfig",
+    "TracesConfig",
+    "TuskApiConfig",
+    "load_tusk_config",
+    "find_project_root",
+    # Types
+    "DriftMode",
+    "CleanSpanData",
+    "PackageType",
+    "SpanKind",
+    "StatusCode",
+    # Batching
+    "BatchSpanProcessor",
+    "BatchSpanProcessorConfig",
+    # Sampling
+    "should_sample",
+    "validate_sampling_rate",
+    # Data normalization
+    "normalize_input_data",
+    "remove_none_values",
+    "create_span_input_value",
+    "create_mock_input_value",
+    # Trace blocking
+    "TraceBlockingManager",
+    "estimate_span_size",
+    "should_block_span",
+    "MAX_SPAN_SIZE_MB",
+    "MAX_SPAN_SIZE_BYTES",
+]

tusk_drift_python_sdk-0.1.0/drift/core/batch_processor.py

@@ -0,0 +1,189 @@
+"""Batch span processor for efficient span export."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+import time
+from collections import deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .types import CleanSpanData
+    from ..tracing.adapters.base import SpanExportAdapter
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BatchSpanProcessorConfig:
+    """Configuration for the batch span processor."""
+
+    # Maximum queue size before spans are dropped
+    max_queue_size: int = 2048
+    # Maximum batch size per export
+    max_export_batch_size: int = 512
+    # Interval between scheduled exports (in seconds)
+    scheduled_delay_seconds: float = 2.0
+    # Maximum time to wait for export (in seconds)
+    export_timeout_seconds: float = 30.0
+
+
+class BatchSpanProcessor:
+    """
+    Batches spans and exports them periodically or when batch size is reached.
+
+    Matches the behavior of OpenTelemetry's BatchSpanProcessor:
+    - Queues spans in memory
+    - Exports in batches at regular intervals or when max batch size is reached
+    - Drops spans if queue is full
+    - Handles graceful shutdown with final export
+    """
+
+    def __init__(
+        self,
+        adapters: list["SpanExportAdapter"],
+        config: BatchSpanProcessorConfig | None = None,
+    ) -> None:
+        """
+        Initialize the batch processor.
+
+        Args:
+            adapters: List of adapters to export spans to
+            config: Optional configuration (uses defaults if not provided)
+        """
+        self._adapters = adapters
+        self._config = config or BatchSpanProcessorConfig()
+        self._queue: deque[CleanSpanData] = deque(maxlen=self._config.max_queue_size)
+        self._lock = threading.Lock()
+        self._shutdown_event = threading.Event()
+        self._export_thread: threading.Thread | None = None
+        self._started = False
+        self._dropped_spans = 0
+
+    def start(self) -> None:
+        """Start the background export thread."""
+        if self._started:
+            return
+
+        self._started = True
+        self._shutdown_event.clear()
+        self._export_thread = threading.Thread(
+            target=self._export_loop,
+            daemon=True,
+            name="drift-batch-exporter",
+        )
+        self._export_thread.start()
+        logger.debug("BatchSpanProcessor started")
+
+    def stop(self, timeout: float | None = None) -> None:
+        """
+        Stop the processor and export remaining spans.
+
+        Args:
+            timeout: Maximum time to wait for final export
+        """
+        if not self._started:
+            return
+
+        self._shutdown_event.set()
+
+        if self._export_thread is not None:
+            self._export_thread.join(timeout=timeout or self._config.export_timeout_seconds)
+
+        # Final export of remaining spans
+        self._force_flush()
+
+        self._started = False
+        logger.debug(f"BatchSpanProcessor stopped. Dropped {self._dropped_spans} spans total.")
+
+    def add_span(self, span: "CleanSpanData") -> bool:
+        """
+        Add a span to the queue for export.
+
+        Args:
+            span: The span to add
+
+        Returns:
+            True if span was added, False if queue is full and span was dropped
+        """
+        with self._lock:
+            if len(self._queue) >= self._config.max_queue_size:
+                self._dropped_spans += 1
+                logger.warning(
+                    f"Span queue full ({self._config.max_queue_size}), dropping span. "
+                    f"Total dropped: {self._dropped_spans}"
+                )
+                return False
+
+            self._queue.append(span)
+
+            # Trigger immediate export if batch size reached
+            if len(self._queue) >= self._config.max_export_batch_size:
+                # Signal export thread to wake up (if using condition variable)
+                pass
+
+            return True
+
+    def _export_loop(self) -> None:
+        """Background thread that periodically exports spans."""
+        while not self._shutdown_event.is_set():
+            # Wait for scheduled delay or shutdown
+            self._shutdown_event.wait(timeout=self._config.scheduled_delay_seconds)
+
+            if self._shutdown_event.is_set():
+                break
+
+            self._export_batch()
+
+    def _export_batch(self) -> None:
+        """Export a batch of spans from the queue."""
+        # Get batch of spans
+        batch: list[CleanSpanData] = []
+        with self._lock:
+            while self._queue and len(batch) < self._config.max_export_batch_size:
+                batch.append(self._queue.popleft())
+
+        if not batch:
+            return
+
+        # Export to all adapters
+        for adapter in self._adapters:
+            try:
+                # Handle async adapters
+                if asyncio.iscoroutinefunction(adapter.export_spans):
+                    try:
+                        loop = asyncio.get_running_loop()
+                        asyncio.create_task(adapter.export_spans(batch))
+                    except RuntimeError:
+                        # No running loop, create one
+                        asyncio.run(adapter.export_spans(batch))
+                else:
+                    adapter.export_spans(batch)  # type: ignore
+
+                logger.debug(f"Exported {len(batch)} spans via {adapter.name}")
+
+            except Exception as e:
+                logger.error(f"Failed to export batch via {adapter.name}: {e}")
+
+    def _force_flush(self) -> None:
+        """Force export all remaining spans in the queue."""
+        while True:
+            with self._lock:
+                if not self._queue:
+                    break
+
+            self._export_batch()
+
+    @property
+    def queue_size(self) -> int:
+        """Get the current queue size."""
+        with self._lock:
+            return len(self._queue)
+
+    @property
+    def dropped_span_count(self) -> int:
+        """Get the number of dropped spans."""
+        return self._dropped_spans

tusk_drift_python_sdk-0.1.0/drift/core/communication/__init__.py

@@ -0,0 +1,52 @@
+"""CLI communication module for Drift SDK.
+
+This module handles bidirectional communication between the SDK and the Tusk CLI
+for replay testing. Communication uses Protocol Buffers over Unix sockets or TCP.
+"""
+
+from .types import (
+    MessageType,
+    SDKMessageType,
+    CLIMessageType,
+    ConnectRequest,
+    ConnectResponse,
+    GetMockRequest,
+    GetMockResponse,
+    EnvVarRequest,
+    EnvVarResponse,
+    MockRequestInput,
+    MockResponseOutput,
+    # Protobuf types (re-exported)
+    SdkMessage,
+    CliMessage,
+    span_to_proto,
+    dict_to_span,
+    extract_response_data,
+)
+from .communicator import ProtobufCommunicator, CommunicatorConfig
+
+__all__ = [
+    # Message types
+    "MessageType",
+    "SDKMessageType",
+    "CLIMessageType",
+    # Request/Response types
+    "ConnectRequest",
+    "ConnectResponse",
+    "GetMockRequest",
+    "GetMockResponse",
+    "EnvVarRequest",
+    "EnvVarResponse",
+    "MockRequestInput",
+    "MockResponseOutput",
+    # Protobuf types
+    "SdkMessage",
+    "CliMessage",
+    # Utilities
+    "span_to_proto",
+    "dict_to_span",
+    "extract_response_data",
+    # Communicator
+    "ProtobufCommunicator",
+    "CommunicatorConfig",
+]