feat: urllib3 instrumentation (#30)

Author: jy-tan

drift/__init__.py

@@ -32,6 +32,7 @@
 from .instrumentation.fastapi import FastAPIInstrumentation
 from .instrumentation.flask import FlaskInstrumentation
 from .instrumentation.requests import RequestsInstrumentation
+from .instrumentation.urllib3 import Urllib3Instrumentation
 
 __version__ = "0.1.0"
 
@@ -61,6 +62,7 @@
     "FlaskInstrumentation",
     "FastAPIInstrumentation",
     "RequestsInstrumentation",
+    "Urllib3Instrumentation",
     # Adapters
     "SpanExportAdapter",
     "ExportResult",

drift/core/drift_sdk.py

@@ -402,6 +402,16 @@ def _init_auto_instrumentations(self) -> None:
         except ImportError:
             pass
 
+        try:
+            import urllib3  # type: ignore[unresolved-import]
+
+            from ..instrumentation.urllib3 import Urllib3Instrumentation
+
+            _ = Urllib3Instrumentation()
+            logger.debug("urllib3 instrumentation initialized")
+        except ImportError:
+            pass
+
         # Initialize PostgreSQL instrumentation before Django
         # Instrument BOTH psycopg2 and psycopg if available
         # This allows apps to use either or both

drift/instrumentation/urllib3/__init__.py

@@ -0,0 +1,5 @@
+"""urllib3 instrumentation module."""
+
+from .instrumentation import RequestDroppedByTransform, Urllib3Instrumentation
+
+__all__ = ["Urllib3Instrumentation", "RequestDroppedByTransform"]

drift/instrumentation/urllib3/e2e-tests/.tusk/config.yaml

@@ -0,0 +1,28 @@
+version: 1
+
+service:
+  id: "urllib3-e2e-test-id"
+  name: "urllib3-e2e-test"
+  port: 8000
+  start:
+    command: "python src/app.py"
+  readiness_check:
+    command: "curl -f http://localhost:8000/health"
+    timeout: 45s
+    interval: 5s
+
+tusk_api:
+  url: "http://localhost:8000"
+
+test_execution:
+  concurrent_limit: 10
+  batch_size: 10
+  timeout: 30s
+
+recording:
+  sampling_rate: 1.0
+  export_spans: false
+
+replay:
+  enable_telemetry: false
+

drift/instrumentation/urllib3/e2e-tests/Dockerfile

@@ -0,0 +1,22 @@
+FROM python-e2e-base:latest
+
+# Copy SDK source for editable install
+COPY . /sdk
+
+# Copy test files
+COPY drift/instrumentation/urllib3/e2e-tests /app
+
+WORKDIR /app
+
+# Install dependencies (requirements.txt uses -e /sdk for SDK)
+RUN pip install -q -r requirements.txt
+
+# Make entrypoint executable
+RUN chmod +x entrypoint.py
+
+# Create .tusk directories
+RUN mkdir -p /app/.tusk/traces /app/.tusk/logs
+
+# Run entrypoint
+ENTRYPOINT ["python", "entrypoint.py"]
+

drift/instrumentation/urllib3/e2e-tests/docker-compose.yml

@@ -0,0 +1,20 @@
+services:
+  app:
+    build:
+      context: ../../../..
+      dockerfile: drift/instrumentation/urllib3/e2e-tests/Dockerfile
+      args:
+        - TUSK_CLI_VERSION=${TUSK_CLI_VERSION:-latest}
+    environment:
+      - PORT=8000
+      - TUSK_ANALYTICS_DISABLED=1
+      - PYTHONUNBUFFERED=1
+    working_dir: /app
+    volumes:
+      # Mount SDK source for hot reload (no rebuild needed for SDK changes)
+      - ../../../..:/sdk
+      # Mount app source for development
+      - ./src:/app/src
+      # Mount .tusk folder to persist traces
+      - ./.tusk:/app/.tusk
+

drift/instrumentation/urllib3/e2e-tests/entrypoint.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+"""
+E2E Test Entrypoint for urllib3 Instrumentation
+
+This script orchestrates the full e2e test lifecycle:
+1. Setup: Install dependencies
+2. Record: Start app in RECORD mode, execute requests
+3. Test: Run Tusk CLI tests
+4. Teardown: Cleanup and return exit code
+"""
+
+import sys
+from pathlib import Path
+
+# Add SDK to path for imports
+sys.path.insert(0, "/sdk")
+
+from drift.instrumentation.e2e_common.base_runner import E2ETestRunnerBase
+
+
+class Urllib3E2ETestRunner(E2ETestRunnerBase):
+    """E2E test runner for urllib3 instrumentation."""
+
+    def __init__(self):
+        import os
+
+        port = int(os.getenv("PORT", "8000"))
+        super().__init__(app_port=port)
+
+
+if __name__ == "__main__":
+    runner = Urllib3E2ETestRunner()
+    exit_code = runner.run()
+    sys.exit(exit_code)

drift/instrumentation/urllib3/e2e-tests/requirements.txt

@@ -0,0 +1,5 @@
+-e /sdk
+Flask>=3.1.2
+urllib3>=2.0.0
+requests>=2.32.5
+

drift/instrumentation/urllib3/e2e-tests/run.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+
+# Exit on error
+set -e
+
+# Accept optional port parameter (default: 8000)
+APP_PORT=${1:-8000}
+export APP_PORT
+
+# Generate unique docker compose project name
+# Get the instrumentation name (parent directory of e2e-tests)
+TEST_NAME="$(basename "$(dirname "$(pwd)")")"
+PROJECT_NAME="python-${TEST_NAME}-${APP_PORT}"
+
+# Colors for output
+GREEN='\033[0;32m'
+RED='\033[0;31m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m'
+
+echo -e "${BLUE}========================================${NC}"
+echo -e "${BLUE}Running Python E2E Test: ${TEST_NAME}${NC}"
+echo -e "${BLUE}Port: ${APP_PORT}${NC}"
+echo -e "${BLUE}========================================${NC}"
+echo ""
+
+# Cleanup function
+cleanup() {
+    echo ""
+    echo -e "${YELLOW}Cleaning up containers...${NC}"
+    docker compose -p "$PROJECT_NAME" down -v 2>/dev/null || true
+}
+
+# Register cleanup on exit
+trap cleanup EXIT
+
+# Build containers
+echo -e "${BLUE}Building containers...${NC}"
+docker compose -p "$PROJECT_NAME" build --no-cache
+
+# Run the test container
+echo -e "${BLUE}Starting test...${NC}"
+echo ""
+
+# Run container and capture exit code (always use port 8000 inside container)
+# Disable set -e temporarily to capture exit code
+set +e
+docker compose -p "$PROJECT_NAME" run --rm app
+EXIT_CODE=$?
+set -e
+
+echo ""
+if [ $EXIT_CODE -eq 0 ]; then
+    echo -e "${GREEN}========================================${NC}"
+    echo -e "${GREEN}Test passed!${NC}"
+    echo -e "${GREEN}========================================${NC}"
+else
+    echo -e "${RED}========================================${NC}"
+    echo -e "${RED}Test failed with exit code ${EXIT_CODE}${NC}"
+    echo -e "${RED}========================================${NC}"
+fi
+
+exit $EXIT_CODE
+