timetracer 1.1.0__py3-none-any.whl

timetracer/plugins/requests_plugin.py

@@ -0,0 +1,333 @@
+"""
+requests plugin for Timetracer.
+
+Captures and replays requests library HTTP calls.
+This provides compatibility for codebases using requests instead of httpx.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, Any, Callable
+from urllib.parse import parse_qs, urlparse
+
+from timetracer.constants import EventType
+from timetracer.context import get_current_session, has_active_session
+from timetracer.policies import redact_headers_allowlist
+from timetracer.types import (
+    BodySnapshot,
+    DependencyEvent,
+    EventResult,
+    EventSignature,
+)
+from timetracer.utils.hashing import hash_body
+
+if TYPE_CHECKING:
+    import requests
+
+# Store original method for restoration
+_original_request: Callable | None = None
+_enabled = False
+
+
+def enable_requests() -> None:
+    """
+    Enable requests interception for recording and replay.
+
+    This patches requests.Session.request.
+    Call disable_requests() to restore original behavior.
+    """
+    global _original_request, _enabled
+
+    if _enabled:
+        return
+
+    try:
+        import requests
+    except ImportError:
+        raise ImportError(
+            "requests is required for the requests plugin. "
+            "Install it with: pip install requests"
+        )
+
+    # Patch Session.request (all requests go through this)
+    _original_request = requests.Session.request
+    requests.Session.request = _patched_request
+
+    _enabled = True
+
+
+def disable_requests() -> None:
+    """
+    Disable requests interception and restore original behavior.
+    """
+    global _original_request, _enabled
+
+    if not _enabled:
+        return
+
+    try:
+        import requests
+    except ImportError:
+        return
+
+    # Restore original
+    if _original_request is not None:
+        requests.Session.request = _original_request
+
+    _original_request = None
+    _enabled = False
+
+
+def _patched_request(
+    self: "requests.Session",
+    method: str,
+    url: str,
+    **kwargs: Any,
+) -> "requests.Response":
+    """Patched request method."""
+    # If no session, call original
+    if not has_active_session():
+        return _original_request(self, method, url, **kwargs)  # type: ignore
+
+    session = get_current_session()
+
+    # Handle based on session mode
+    if session.is_recording:
+        return _record_request(self, method, url, kwargs)
+    elif session.is_replaying:
+        # Check for hybrid replay - if plugin should stay live, make real call
+        from timetracer.session import ReplaySession
+        if isinstance(session, ReplaySession) and not session.should_mock_plugin("http"):
+            return _original_request(self, method, url, **kwargs)  # type: ignore
+        return _replay_request(method, url, kwargs)
+    else:
+        return _original_request(self, method, url, **kwargs)  # type: ignore
+
+
+def _record_request(
+    client: "requests.Session",
+    method: str,
+    url: str,
+    kwargs: dict[str, Any],
+) -> "requests.Response":
+    """Record a requests call."""
+    from timetracer.session import TraceSession
+
+    session = get_current_session()
+    if not isinstance(session, TraceSession):
+        return _original_request(client, method, url, **kwargs)  # type: ignore
+
+    start_offset = session.elapsed_ms
+    start_time = time.perf_counter()
+
+    # Make the actual request
+    error_info = None
+    response = None
+    try:
+        response = _original_request(client, method, url, **kwargs)  # type: ignore
+    except Exception as e:
+        error_info = (type(e).__name__, str(e))
+        raise
+    finally:
+        duration_ms = (time.perf_counter() - start_time) * 1000
+
+        # Build event
+        event = _build_event(
+            method=method,
+            url=url,
+            kwargs=kwargs,
+            response=response,
+            start_offset_ms=start_offset,
+            duration_ms=duration_ms,
+            error_info=error_info,
+        )
+
+        session.add_event(event)
+
+    return response
+
+
+def _replay_request(
+    method: str,
+    url: str,
+    kwargs: dict[str, Any],
+) -> "requests.Response":
+    """Replay a requests call from cassette."""
+    from timetracer.session import ReplaySession
+
+    session = get_current_session()
+    if not isinstance(session, ReplaySession):
+        raise RuntimeError("Expected ReplaySession for replay")
+
+    # Build signature for matching
+    actual_signature = _make_signature_dict(method, url, kwargs)
+
+    # Get expected event
+    event = session.get_next_event(EventType.HTTP_CLIENT, actual_signature)
+
+    # Build synthetic response
+    return _build_synthetic_response(event)
+
+
+def _build_event(
+    method: str,
+    url: str,
+    kwargs: dict[str, Any],
+    response: "requests.Response | None",
+    start_offset_ms: float,
+    duration_ms: float,
+    error_info: tuple[str, str] | None,
+) -> DependencyEvent:
+    """Build a DependencyEvent from requests call."""
+    # Build signature
+    signature = _make_signature(method, url, kwargs)
+
+    # Build result
+    result = _make_result(response, error_info)
+
+    return DependencyEvent(
+        eid=0,  # Will be set by session
+        event_type=EventType.HTTP_CLIENT,
+        start_offset_ms=start_offset_ms,
+        duration_ms=duration_ms,
+        signature=signature,
+        result=result,
+    )
+
+
+def _make_signature(method: str, url: str, kwargs: dict[str, Any]) -> EventSignature:
+    """Create EventSignature from requests call."""
+    # Normalize URL
+    parsed = urlparse(url)
+    normalized_url = f"{parsed.scheme}://{parsed.netloc}{parsed.path}"
+
+    # Parse query params (from URL or kwargs)
+    query = dict(parse_qs(parsed.query))
+    if "params" in kwargs and kwargs["params"]:
+        params = kwargs["params"]
+        if isinstance(params, dict):
+            query.update({k: [v] if not isinstance(v, list) else v for k, v in params.items()})
+    # Flatten single-value lists
+    query = {k: v[0] if len(v) == 1 else v for k, v in query.items()}
+
+    # Hash request body
+    body_hash = None
+    if "data" in kwargs and kwargs["data"]:
+        body_hash = hash_body(kwargs["data"])
+    elif "json" in kwargs and kwargs["json"]:
+        body_hash = hash_body(kwargs["json"])
+
+    # Hash relevant headers
+    headers_hash = None
+    if "headers" in kwargs and kwargs["headers"]:
+        allowed_headers = redact_headers_allowlist(kwargs["headers"])
+        if allowed_headers:
+            headers_hash = hash_body(str(sorted(allowed_headers.items())))
+
+    return EventSignature(
+        lib="requests",
+        method=method.upper(),
+        url=normalized_url,
+        query=query,
+        headers_hash=headers_hash,
+        body_hash=body_hash,
+    )
+
+
+def _make_signature_dict(method: str, url: str, kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Create signature dict for matching."""
+    sig = _make_signature(method, url, kwargs)
+    return {
+        "lib": sig.lib,
+        "method": sig.method,
+        "url": sig.url,
+        "query": sig.query,
+        "body_hash": sig.body_hash,
+    }
+
+
+def _make_result(
+    response: "requests.Response | None",
+    error_info: tuple[str, str] | None,
+) -> EventResult:
+    """Create EventResult from requests response."""
+    if error_info:
+        return EventResult(
+            error_type=error_info[0],
+            error=error_info[1],
+        )
+
+    if response is None:
+        return EventResult()
+
+    # Capture response headers (allow-list)
+    headers = dict(response.headers)
+    filtered_headers = redact_headers_allowlist(headers)
+
+    # Capture response body
+    body_snapshot = None
+    try:
+        content = response.content
+        if content:
+            body_snapshot = BodySnapshot(
+                captured=True,
+                encoding="bytes",
+                data=None,
+                hash=hash_body(content),
+                size_bytes=len(content),
+            )
+
+            # Try to parse as JSON for storage
+            try:
+                import json
+                data = json.loads(content.decode("utf-8"))
+                body_snapshot.encoding = "json"
+                body_snapshot.data = data
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                pass
+    except Exception:
+        pass
+
+    return EventResult(
+        status=response.status_code,
+        headers=filtered_headers,
+        body=body_snapshot,
+    )
+
+
+def _build_synthetic_response(event: DependencyEvent) -> "requests.Response":
+    """Build a synthetic requests.Response from recorded event."""
+    import requests
+    from requests.structures import CaseInsensitiveDict
+
+    result = event.result
+
+    # Handle error case
+    if result.error:
+        raise requests.RequestException(f"Recorded error: {result.error}")
+
+    # Build response content
+    content = b""
+    if result.body and result.body.captured and result.body.data is not None:
+        import json
+        if result.body.encoding == "json":
+            content = json.dumps(result.body.data).encode("utf-8")
+        elif isinstance(result.body.data, str):
+            content = result.body.data.encode("utf-8")
+        elif isinstance(result.body.data, bytes):
+            content = result.body.data
+
+    # Build headers
+    headers = CaseInsensitiveDict(result.headers or {})
+    if "content-length" not in headers:
+        headers["content-length"] = str(len(content))
+
+    # Create response object
+    response = requests.Response()
+    response.status_code = result.status or 200
+    response.headers = headers
+    response._content = content
+    response.encoding = "utf-8"
+
+    return response

timetracer/plugins/sqlalchemy_plugin.py

@@ -0,0 +1,280 @@
+"""
+SQLAlchemy plugin for Timetracer.
+
+Captures and replays SQLAlchemy database queries.
+This enables recording database interactions for deterministic replay.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from typing import TYPE_CHECKING, Any
+
+from timetracer.constants import EventType
+from timetracer.context import get_current_session, has_active_session
+from timetracer.types import (
+    DependencyEvent,
+    EventResult,
+    EventSignature,
+)
+from timetracer.utils.hashing import hash_body
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Connection, Engine
+
+# Store original state for restoration
+_enabled = False
+_listeners_attached = False
+
+
+def enable_sqlalchemy(engine: "Engine | None" = None) -> None:
+    """
+    Enable SQLAlchemy interception for recording and replay.
+
+    This attaches event listeners to capture query execution.
+    Call disable_sqlalchemy() to restore original behavior.
+
+    Args:
+        engine: Optional specific engine. If None, uses global event.
+    """
+    global _enabled, _listeners_attached
+
+    if _enabled:
+        return
+
+    try:
+        from sqlalchemy import event as sa_event
+        from sqlalchemy.engine import Engine
+    except ImportError:
+        raise ImportError(
+            "sqlalchemy is required for the sqlalchemy plugin. "
+            "Install it with: pip install sqlalchemy"
+        )
+
+    # Attach listeners
+    if engine is not None:
+        sa_event.listen(engine, "before_cursor_execute", _before_cursor_execute)
+        sa_event.listen(engine, "after_cursor_execute", _after_cursor_execute)
+    else:
+        # Global listener for all engines
+        sa_event.listen(Engine, "before_cursor_execute", _before_cursor_execute)
+        sa_event.listen(Engine, "after_cursor_execute", _after_cursor_execute)
+
+    _listeners_attached = True
+    _enabled = True
+
+
+def disable_sqlalchemy(engine: "Engine | None" = None) -> None:
+    """
+    Disable SQLAlchemy interception and restore original behavior.
+    """
+    global _enabled, _listeners_attached
+
+    if not _enabled:
+        return
+
+    try:
+        from sqlalchemy import event as sa_event
+        from sqlalchemy.engine import Engine
+    except ImportError:
+        return
+
+    # Remove listeners
+    if engine is not None:
+        sa_event.remove(engine, "before_cursor_execute", _before_cursor_execute)
+        sa_event.remove(engine, "after_cursor_execute", _after_cursor_execute)
+    else:
+        sa_event.remove(Engine, "before_cursor_execute", _before_cursor_execute)
+        sa_event.remove(Engine, "after_cursor_execute", _after_cursor_execute)
+
+    _listeners_attached = False
+    _enabled = False
+
+
+# Thread-local storage for timing
+_query_timing = threading.local()
+
+
+def _before_cursor_execute(
+    conn: "Connection",
+    cursor: Any,
+    statement: str,
+    parameters: Any,
+    context: Any,
+    executemany: bool,
+) -> None:
+    """Event listener called before query execution."""
+    if not has_active_session():
+        return
+
+    session = get_current_session()
+
+    # Store timing info
+    _query_timing.start_time = time.perf_counter()
+    _query_timing.start_offset = session.elapsed_ms if hasattr(session, 'elapsed_ms') else 0
+    _query_timing.statement = statement
+    _query_timing.parameters = parameters
+
+
+def _after_cursor_execute(
+    conn: "Connection",
+    cursor: Any,
+    statement: str,
+    parameters: Any,
+    context: Any,
+    executemany: bool,
+) -> None:
+    """Event listener called after query execution."""
+    if not has_active_session():
+        return
+
+    session = get_current_session()
+
+    # Check if we have timing info
+    if not hasattr(_query_timing, 'start_time'):
+        return
+
+    # Handle based on session mode
+    if session.is_recording:
+        _record_query(
+            conn, cursor, statement, parameters, context, executemany
+        )
+    # Replay is handled differently for DB - we can't easily mock cursor results
+    # For now, just record and skip replay mocking
+
+
+def _record_query(
+    conn: "Connection",
+    cursor: Any,
+    statement: str,
+    parameters: Any,
+    context: Any,
+    executemany: bool,
+) -> None:
+    """Record a database query."""
+    from timetracer.session import TraceSession
+
+    session = get_current_session()
+    if not isinstance(session, TraceSession):
+        return
+
+    # Calculate duration
+    duration_ms = (time.perf_counter() - _query_timing.start_time) * 1000
+    start_offset = getattr(_query_timing, 'start_offset', 0)
+
+    # Build event
+    event = _build_event(
+        statement=statement,
+        parameters=parameters,
+        cursor=cursor,
+        start_offset_ms=start_offset,
+        duration_ms=duration_ms,
+        executemany=executemany,
+    )
+
+    session.add_event(event)
+
+
+def _build_event(
+    statement: str,
+    parameters: Any,
+    cursor: Any,
+    start_offset_ms: float,
+    duration_ms: float,
+    executemany: bool,
+) -> DependencyEvent:
+    """Build a DependencyEvent from SQLAlchemy query."""
+    # Build signature
+    signature = _make_signature(statement, parameters)
+
+    # Build result (basic - just rowcount)
+    result = _make_result(cursor)
+
+    return DependencyEvent(
+        eid=0,  # Will be set by session
+        event_type=EventType.DB_QUERY,
+        start_offset_ms=start_offset_ms,
+        duration_ms=duration_ms,
+        signature=signature,
+        result=result,
+    )
+
+
+def _make_signature(statement: str, parameters: Any) -> EventSignature:
+    """Create EventSignature from SQLAlchemy query."""
+    # Normalize statement (strip whitespace)
+    normalized = " ".join(statement.split())
+
+    # Extract operation type
+    operation = normalized.split()[0].upper() if normalized else "UNKNOWN"
+
+    # Hash parameters
+    params_hash = None
+    if parameters:
+        try:
+            params_hash = hash_body(str(parameters))
+        except Exception:
+            params_hash = None
+
+    # Extract table name (basic parsing)
+    table = _extract_table_name(normalized, operation)
+
+    return EventSignature(
+        lib="sqlalchemy",
+        method=operation,  # SELECT, INSERT, UPDATE, DELETE
+        url=table,  # Use url field for table name
+        query={},
+        body_hash=params_hash,
+    )
+
+
+def _extract_table_name(statement: str, operation: str) -> str:
+    """Extract table name from SQL statement."""
+    statement_upper = statement.upper()
+
+    try:
+        if operation == "SELECT":
+            # SELECT ... FROM table_name
+            from_idx = statement_upper.find(" FROM ")
+            if from_idx != -1:
+                after_from = statement[from_idx + 6:].strip()
+                return after_from.split()[0].strip("();")
+
+        elif operation == "INSERT":
+            # INSERT INTO table_name
+            into_idx = statement_upper.find(" INTO ")
+            if into_idx != -1:
+                after_into = statement[into_idx + 6:].strip()
+                return after_into.split()[0].strip("();")
+
+        elif operation == "UPDATE":
+            # UPDATE table_name
+            parts = statement.split()
+            if len(parts) > 1:
+                return parts[1].strip("();")
+
+        elif operation == "DELETE":
+            # DELETE FROM table_name
+            from_idx = statement_upper.find(" FROM ")
+            if from_idx != -1:
+                after_from = statement[from_idx + 6:].strip()
+                return after_from.split()[0].strip("();")
+    except Exception:
+        pass
+
+    return "unknown"
+
+
+def _make_result(cursor: Any) -> EventResult:
+    """Create EventResult from cursor."""
+    rowcount = -1
+    try:
+        rowcount = cursor.rowcount
+    except Exception:
+        pass
+
+    return EventResult(
+        status=rowcount,  # Use status for rowcount
+        headers={"rowcount": str(rowcount)},
+    )

timetracer/policies/__init__.py

@@ -0,0 +1,16 @@
+"""Policies module for redaction and capture control."""
+
+from timetracer.policies.capture import should_store_body, truncate_body
+from timetracer.policies.redaction import (
+    redact_body,
+    redact_headers,
+    redact_headers_allowlist,
+)
+
+__all__ = [
+    "redact_headers",
+    "redact_headers_allowlist",
+    "redact_body",
+    "should_store_body",
+    "truncate_body",
+]

timetracer/policies/capture.py

@@ -0,0 +1,64 @@
+"""
+Capture policies for body data.
+
+Controls when and how much body data is captured.
+"""
+
+from __future__ import annotations
+
+from timetracer.constants import CapturePolicy
+
+
+def should_store_body(
+    policy: CapturePolicy | str,
+    is_error: bool = False,
+) -> bool:
+    """
+    Determine if body should be stored based on policy.
+
+    Args:
+        policy: The capture policy (never, on_error, always).
+        is_error: Whether the request resulted in an error.
+
+    Returns:
+        True if body should be stored.
+    """
+    # Handle string policy
+    if isinstance(policy, str):
+        policy = CapturePolicy(policy.lower())
+
+    if policy == CapturePolicy.ALWAYS:
+        return True
+    elif policy == CapturePolicy.NEVER:
+        return False
+    elif policy == CapturePolicy.ON_ERROR:
+        return is_error
+
+    return False
+
+
+def truncate_body(
+    data: bytes,
+    max_kb: int,
+) -> tuple[bytes, bool]:
+    """
+    Truncate body data if it exceeds size limit.
+
+    Args:
+        data: The body bytes.
+        max_kb: Maximum size in kilobytes.
+
+    Returns:
+        Tuple of (truncated_data, was_truncated).
+    """
+    max_bytes = max_kb * 1024
+
+    if len(data) <= max_bytes:
+        return data, False
+
+    return data[:max_bytes], True
+
+
+def get_body_size_kb(data: bytes) -> float:
+    """Get body size in kilobytes."""
+    return len(data) / 1024