merge main

Author: sohankshirsagar

.github/workflows/e2e.yml

@@ -22,7 +22,7 @@ jobs:
       fail-fast: false
       max-parallel: 6
       matrix:
-        library: [flask, fastapi, django, redis, requests, httpx, psycopg, psycopg2]
+        library: [flask, fastapi, django, redis, requests, httpx, psycopg, psycopg2, urllib3]
     steps:
       - name: Checkout
         uses: actions/checkout@v4

.tusk/config.yaml

@@ -45,21 +45,6 @@ uv run python -m unittest tests.unit.test_adapters -v
 timeout 30 uv run python -m unittest discover -s tests/integration -v
 ```
 
-### Database Integration Tests
-
-Requires Docker for running test databases.
-
-```bash
-# Start test databases
-docker compose -f docker-compose.test.yml up -d
-
-# Run database tests
-uv run python -m unittest tests.integration.test_database -v
-
-# Stop databases
-docker compose -f docker-compose.test.yml down
-```
-
 ### E2E Tests
 
 E2E tests validate full instrumentation workflows using Docker containers. They record real API interactions and verify replay behavior using the Tusk CLI.

CONTRIBUTING.md

@@ -69,19 +69,6 @@ After completing the CLI wizard, install the SDK:
 pip install tusk-drift-python-sdk
 ```
 
-#### With Framework Support
-
-```bash
-# With Flask support
-pip install tusk-drift-python-sdk[flask]
-
-# With FastAPI support
-pip install tusk-drift-python-sdk[fastapi]
-
-# With Django support
-pip install tusk-drift-python-sdk[django]
-```
-
 ### Step 3: Initialize the SDK for your service
 
 Refer to our [initialization guide](docs/initialization.md) to set up the SDK for your service.

MANIFEST.in

@@ -32,8 +32,10 @@
 from .instrumentation.fastapi import FastAPIInstrumentation
 from .instrumentation.flask import FlaskInstrumentation
 from .instrumentation.requests import RequestsInstrumentation
+from .instrumentation.urllib3 import Urllib3Instrumentation
+from .version import SDK_VERSION
 
-__version__ = "0.1.0"
+__version__ = SDK_VERSION
 
 __all__ = [
     # Core
@@ -61,6 +63,7 @@
     "FlaskInstrumentation",
     "FastAPIInstrumentation",
     "RequestsInstrumentation",
+    "Urllib3Instrumentation",
     # Adapters
     "SpanExportAdapter",
     "ExportResult",

README.md

@@ -57,6 +57,7 @@ def __init__(
         self._config = config or BatchSpanProcessorConfig()
         self._queue: deque[CleanSpanData] = deque(maxlen=self._config.max_queue_size)
         self._lock = threading.Lock()
+        self._condition = threading.Condition(self._lock)
         self._shutdown_event = threading.Event()
         self._export_thread: threading.Thread | None = None
         self._started = False
@@ -88,6 +89,9 @@ def stop(self, timeout: float | None = None) -> None:
             return
 
         self._shutdown_event.set()
+        # Wake up the export thread so it can see the shutdown event
+        with self._condition:
+            self._condition.notify_all()
 
         if self._export_thread is not None:
             self._export_thread.join(timeout=timeout or self._config.export_timeout_seconds)
@@ -108,7 +112,7 @@ def add_span(self, span: CleanSpanData) -> bool:
         Returns:
             True if span was added, False if queue is full and span was dropped
         """
-        with self._lock:
+        with self._condition:
             if len(self._queue) >= self._config.max_queue_size:
                 self._dropped_spans += 1
                 logger.warning(
@@ -121,16 +125,17 @@ def add_span(self, span: CleanSpanData) -> bool:
 
             # 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
+                self._condition.notify()
 
             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)
+            # Wait for either: batch size reached, scheduled delay, or shutdown
+            with self._condition:
+                # Wait until batch is ready or timeout
+                self._condition.wait(timeout=self._config.scheduled_delay_seconds)
 
             if self._shutdown_event.is_set():
                 break
@@ -141,7 +146,7 @@ def _export_batch(self) -> None:
         """Export a batch of spans from the queue."""
         # Get batch of spans
         batch: list[CleanSpanData] = []
-        with self._lock:
+        with self._condition:
             while self._queue and len(batch) < self._config.max_export_batch_size:
                 batch.append(self._queue.popleft())
 
@@ -171,7 +176,7 @@ def _export_batch(self) -> None:
     def _force_flush(self) -> None:
         """Force export all remaining spans in the queue."""
         while True:
-            with self._lock:
+            with self._condition:
                 if not self._queue:
                     break
 
@@ -180,7 +185,7 @@ def _force_flush(self) -> None:
     @property
     def queue_size(self) -> int:
         """Get the current queue size."""
-        with self._lock:
+        with self._condition:
             return len(self._queue)
 
     @property

docker-compose.test.yml

@@ -253,7 +253,6 @@ def connect_sync(
             ConnectionError: If connection fails
             TimeoutError: If connection times out
         """
-        logger.info("[CONNECT_SYNC] Starting synchronous connection")
         # Determine address
         if connection_info:
             if "socketPath" in connection_info:
@@ -324,8 +323,6 @@ def connect_sync(
                 if response.success:
                     logger.debug("CLI acknowledged connection successfully")
                     self._connected = True
-                    logger.info(f"[CONNECT_SYNC] Connection successful! Socket is: {self._socket}")
-                    logger.info(f"[CONNECT_SYNC] _connected={self._connected}, is_connected={self.is_connected}")
 
                     # Start background reader for CLI-initiated messages (like SetTimeTravel)
                     self._start_background_reader()
@@ -344,8 +341,6 @@ def connect_sync(
         finally:
             calling_library_context.reset(context_token)
 
-        logger.info(f"[CONNECT_SYNC] Exiting connect_sync(). Socket still open: {self._socket is not None}")
-
     async def disconnect(self) -> None:
         """Disconnect from CLI."""
         self._cleanup()
@@ -517,8 +512,9 @@ async def send_unpatched_dependency_alert(
 
         try:
             await self._send_protobuf_message(sdk_message)
-        except Exception:
-            pass  # Alerts are non-critical
+        except Exception as e:
+            # Alerts are non-critical, just log at debug level
+            logger.debug(f"Failed to send unpatched dependency alert: {e}")
 
     async def _send_protobuf_message(self, message: SdkMessage) -> None:
         """Send a protobuf message to CLI."""
@@ -775,9 +771,6 @@ def _cleanup(self) -> None:
         """Clean up resources."""
         import traceback
 
-        logger.warning("[CLEANUP] _cleanup() called! Stack trace:")
-        logger.warning("".join(traceback.format_stack()))
-
         # Stop background reader thread
         self._stop_background_reader.set()
         if self._background_reader_thread and self._background_reader_thread.is_alive():
@@ -790,10 +783,10 @@ def _cleanup(self) -> None:
 
         if self._socket:
             try:
-                logger.warning("[CLEANUP] Closing socket")
                 self._socket.close()
-            except Exception:
-                pass
+            except OSError as e:
+                # Socket may already be closed, which is fine
+                logger.debug(f"Error closing socket during cleanup: {e}")
             self._socket = None
 
         self._pending_requests.clear()

drift/__init__.py

@@ -135,8 +135,17 @@ def _parse_service_config(data: dict[str, Any]) -> ServiceConfig:
 
 def _parse_recording_config(data: dict[str, Any]) -> RecordingConfig:
     """Parse recording configuration from raw dict."""
+    # Validate sampling_rate type
+    sampling_rate = data.get("sampling_rate")
+    if sampling_rate is not None and not isinstance(sampling_rate, (int, float)):
+        logger.warning(
+            f"Invalid 'sampling_rate' in config: expected number, got {type(sampling_rate).__name__}. "
+            "This value will be ignored."
+        )
+        sampling_rate = None
+
     return RecordingConfig(
-        sampling_rate=data.get("sampling_rate"),
+        sampling_rate=sampling_rate,
         export_spans=data.get("export_spans"),
         enable_env_var_recording=data.get("enable_env_var_recording"),
         enable_analytics=data.get("enable_analytics"),