obsforge 0.1.1__py3-none-any.whl

obsforge/instrumentation/db/pool.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from threading import Lock
+
+from obsforge.config.settings import DatabaseInstrumentationSettings
+
+
+@dataclass(slots=True)
+class PoolSnapshot:
+    pool_name: str
+    pool_size: int | None = None
+    checked_out: int | None = None
+    available: int | None = None
+    waiters: int | None = None
+    saturated: bool = False
+    leak_suspected: bool = False
+    checkout_wait_ms: float | None = None
+
+
+@dataclass(slots=True)
+class CheckoutState:
+    checked_out_at: float
+    connection_id: str
+
+
+class PoolTracker:
+    def __init__(self, settings: DatabaseInstrumentationSettings) -> None:
+        self._settings = settings
+        self._states: dict[str, dict[str, CheckoutState]] = {}
+        self._lock = Lock()
+
+    def checkout(
+        self,
+        *,
+        pool_name: str,
+        connection_id: str,
+        pool_size: int | None = None,
+        checked_out: int | None = None,
+        available: int | None = None,
+        waiters: int | None = None,
+        checkout_wait_ms: float | None = None,
+    ) -> PoolSnapshot:
+        with self._lock:
+            pool = self._states.setdefault(pool_name, {})
+            pool[connection_id] = CheckoutState(checked_out_at=time.perf_counter(), connection_id=connection_id)
+            checked_out_now = checked_out if checked_out is not None else len(pool)
+            available_now = available if available is not None else (
+                max(pool_size - checked_out_now, 0) if pool_size is not None else None
+            )
+            waiters_now = waiters
+            return PoolSnapshot(
+                pool_name=pool_name,
+                pool_size=pool_size,
+                checked_out=checked_out_now,
+                available=available_now,
+                waiters=waiters_now,
+                saturated=(waiters_now or 0) >= self._settings.pool_exhaustion_waiters_threshold
+                or (available_now == 0 and checked_out_now > 0 if available_now is not None else False),
+                leak_suspected=self._leak_suspected(pool),
+                checkout_wait_ms=checkout_wait_ms,
+            )
+
+    def checkin(
+        self,
+        *,
+        pool_name: str,
+        connection_id: str,
+        pool_size: int | None = None,
+        checked_out: int | None = None,
+        available: int | None = None,
+        waiters: int | None = None,
+    ) -> PoolSnapshot:
+        with self._lock:
+            pool = self._states.setdefault(pool_name, {})
+            pool.pop(connection_id, None)
+            checked_out_now = checked_out if checked_out is not None else len(pool)
+            available_now = available if available is not None else (
+                max(pool_size - checked_out_now, 0) if pool_size is not None else None
+            )
+            return PoolSnapshot(
+                pool_name=pool_name,
+                pool_size=pool_size,
+                checked_out=checked_out_now,
+                available=available_now,
+                waiters=waiters,
+                saturated=(waiters or 0) >= self._settings.pool_exhaustion_waiters_threshold,
+                leak_suspected=self._leak_suspected(pool),
+            )
+
+    def snapshot(
+        self,
+        *,
+        pool_name: str,
+        pool_size: int | None = None,
+        checked_out: int | None = None,
+        available: int | None = None,
+        waiters: int | None = None,
+    ) -> PoolSnapshot:
+        with self._lock:
+            pool = self._states.setdefault(pool_name, {})
+            checked_out_now = checked_out if checked_out is not None else len(pool)
+            available_now = available if available is not None else (
+                max(pool_size - checked_out_now, 0) if pool_size is not None else None
+            )
+            return PoolSnapshot(
+                pool_name=pool_name,
+                pool_size=pool_size,
+                checked_out=checked_out_now,
+                available=available_now,
+                waiters=waiters,
+                saturated=(waiters or 0) >= self._settings.pool_exhaustion_waiters_threshold
+                or (available_now == 0 and checked_out_now > 0 if available_now is not None else False),
+                leak_suspected=self._leak_suspected(pool),
+            )
+
+    def _leak_suspected(self, pool: dict[str, CheckoutState]) -> bool:
+        now = time.perf_counter()
+        for state in pool.values():
+            if (now - state.checked_out_at) * 1000 >= self._settings.leak_checkout_threshold_ms:
+                return True
+        return False

obsforge/instrumentation/db/psycopg.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+from types import ModuleType
+from typing import Any
+
+from obsforge.instrumentation.db.engine import DBObservabilityEngine
+
+psycopg: ModuleType | None
+try:  # pragma: no cover - optional dependency
+    import psycopg
+except ImportError:  # pragma: no cover - optional dependency
+    psycopg = None
+
+
+class PsycopgCursorProxy:
+    def __init__(self, cursor: Any, connection_proxy: PsycopgConnectionProxy) -> None:
+        self._cursor = cursor
+        self._connection_proxy = connection_proxy
+
+    def execute(self, query: str, params: Any = None) -> Any:
+        capture = self._connection_proxy._db_engine.start_query(
+            system="postgresql",
+            sql=query,
+            database_name=getattr(self._connection_proxy._connection.info, "dbname", None),
+            host=getattr(self._connection_proxy._connection.info, "host", None),
+            port=getattr(self._connection_proxy._connection.info, "port", None),
+            role="primary",
+            connection_id=str(id(self._connection_proxy._connection)),
+            many=False,
+            namespace=getattr(self._connection_proxy._connection.info, "dbname", None),
+        )
+        try:
+            result = self._cursor.execute(query, params)
+        except BaseException as exc:
+            self._connection_proxy._db_engine.emit_query_sync(capture, exception=exc)
+            raise
+        self._connection_proxy._db_engine.emit_query_sync(
+            capture,
+            rows_affected=getattr(self._cursor, "rowcount", None),
+            transaction_state=self._connection_proxy._db_engine.transaction_state(
+                str(id(self._connection_proxy._connection))
+            ),
+        )
+        return result
+
+    def executemany(self, query: str, params_seq: Any) -> Any:
+        capture = self._connection_proxy._db_engine.start_query(
+            system="postgresql",
+            sql=query,
+            database_name=getattr(self._connection_proxy._connection.info, "dbname", None),
+            host=getattr(self._connection_proxy._connection.info, "host", None),
+            port=getattr(self._connection_proxy._connection.info, "port", None),
+            role="primary",
+            connection_id=str(id(self._connection_proxy._connection)),
+            many=True,
+            namespace=getattr(self._connection_proxy._connection.info, "dbname", None),
+        )
+        try:
+            result = self._cursor.executemany(query, params_seq)
+        except BaseException as exc:
+            self._connection_proxy._db_engine.emit_query_sync(capture, exception=exc)
+            raise
+        self._connection_proxy._db_engine.emit_query_sync(
+            capture,
+            rows_affected=getattr(self._cursor, "rowcount", None),
+            transaction_state=self._connection_proxy._db_engine.transaction_state(
+                str(id(self._connection_proxy._connection))
+            ),
+        )
+        return result
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._cursor, name)
+
+    def __enter__(self) -> PsycopgCursorProxy:
+        self._cursor.__enter__()
+        return self
+
+    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any:
+        return self._cursor.__exit__(exc_type, exc, tb)
+
+
+class PsycopgConnectionProxy:
+    def __init__(self, connection: Any, db_engine: DBObservabilityEngine) -> None:
+        self._connection = connection
+        self._db_engine = db_engine
+
+    def cursor(self, *args: Any, **kwargs: Any) -> PsycopgCursorProxy:
+        return PsycopgCursorProxy(self._connection.cursor(*args, **kwargs), self)
+
+    def execute(self, query: str, params: Any = None) -> Any:
+        with self.cursor() as cursor:
+            return cursor.execute(query, params)
+
+    def commit(self) -> Any:
+        self._db_engine.transaction_commit(str(id(self._connection)))
+        return self._connection.commit()
+
+    def rollback(self) -> Any:
+        self._db_engine.transaction_rollback(str(id(self._connection)))
+        return self._connection.rollback()
+
+    def transaction_begin(self) -> None:
+        self._db_engine.transaction_begin(
+            str(id(self._connection)),
+            isolation_level=getattr(self._connection, "isolation_level", None),
+            autocommit=getattr(self._connection, "autocommit", None),
+        )
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._connection, name)
+
+
+def instrument_psycopg_connection(connection: Any, db_engine: DBObservabilityEngine) -> PsycopgConnectionProxy:
+    return PsycopgConnectionProxy(connection, db_engine)
+
+
+def instrument_psycopg_pool(pool: Any, db_engine: DBObservabilityEngine) -> Any:
+    original_connection = pool.connection
+
+    def connection(*args: Any, **kwargs: Any) -> Any:
+        context_manager = original_connection(*args, **kwargs)
+
+        class ConnectionContextProxy:
+            def __enter__(self) -> PsycopgConnectionProxy:
+                raw = context_manager.__enter__()
+                stats = _pool_stats(pool)
+                db_engine.pool_checkout(
+                    pool_name=type(pool).__name__,
+                    connection_id=str(id(raw)),
+                    pool_size=stats.get("pool_size"),
+                    available=stats.get("pool_available"),
+                    waiters=stats.get("requests_waiting"),
+                )
+                return instrument_psycopg_connection(raw, db_engine)
+
+            def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any:
+                result = context_manager.__exit__(exc_type, exc, tb)
+                stats = _pool_stats(pool)
+                raw_id = None
+                if hasattr(context_manager, "__dict__"):
+                    raw = context_manager.__dict__.get("_conn")
+                    raw_id = str(id(raw)) if raw is not None else None
+                if raw_id is not None:
+                    db_engine.pool_checkin(
+                        pool_name=type(pool).__name__,
+                        connection_id=raw_id,
+                        pool_size=stats.get("pool_size"),
+                        available=stats.get("pool_available"),
+                        waiters=stats.get("requests_waiting"),
+                    )
+                return result
+
+        return ConnectionContextProxy()
+
+    pool.connection = connection
+    return pool
+
+
+def _pool_stats(pool: Any) -> dict[str, Any]:
+    getter = getattr(pool, "get_stats", None)
+    if callable(getter):
+        try:
+            stats: dict[str, Any] = getter()
+            return stats
+        except Exception:
+            return {}
+    return {}

obsforge/instrumentation/db/sql.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import re
+from hashlib import sha1
+
+from obsforge.config.settings import DatabaseInstrumentationSettings
+
+STRING_RE = re.compile(r"'(?:''|[^'])*'")
+NUMBER_RE = re.compile(r"\b\d+(?:\.\d+)?\b")
+UUID_RE = re.compile(
+    r"\b[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\b",
+    re.IGNORECASE,
+)
+WS_RE = re.compile(r"\s+")
+IN_RE = re.compile(r"\bin\s*\(([^)]+)\)", re.IGNORECASE)
+
+
+class SQLNormalizer:
+    def __init__(self, settings: DatabaseInstrumentationSettings) -> None:
+        self._settings = settings
+
+    def normalize(self, sql: str) -> str:
+        normalized = UUID_RE.sub(self._settings.redacted_literals_placeholder, sql)
+        normalized = STRING_RE.sub(self._settings.redacted_literals_placeholder, normalized)
+        normalized = NUMBER_RE.sub(self._settings.redacted_literals_placeholder, normalized)
+        normalized = IN_RE.sub("IN (...)", normalized)
+        normalized = WS_RE.sub(" ", normalized).strip()
+        if len(normalized) > self._settings.max_normalized_query_chars:
+            normalized = normalized[: self._settings.max_normalized_query_chars - 3] + "..."
+        return normalized
+
+    def fingerprint(self, sql: str) -> str:
+        return sha1(self.normalize(sql).lower().encode("utf-8")).hexdigest()
+
+    def summary(self, sql: str) -> str:
+        normalized = self.normalize(sql)
+        parts = normalized.split(" ", 3)
+        if not parts:
+            return "unknown query"
+        operation = parts[0].upper()
+        if operation == "SELECT":
+            return self._select_summary(normalized)
+        if operation in {"UPDATE", "INSERT", "DELETE"}:
+            target = parts[1] if len(parts) > 1 else "unknown"
+            return f"{operation} {target}"
+        if operation in {"BEGIN", "COMMIT", "ROLLBACK", "SAVEPOINT", "RELEASE"}:
+            return operation
+        return " ".join(parts[:2])
+
+    @staticmethod
+    def operation(sql: str) -> str:
+        stripped = sql.lstrip()
+        if not stripped:
+            return "UNKNOWN"
+        return stripped.split(" ", 1)[0].upper()
+
+    @staticmethod
+    def _select_summary(sql: str) -> str:
+        lower = sql.lower()
+        from_index = lower.find(" from ")
+        if from_index == -1:
+            return "SELECT"
+        after_from = sql[from_index + 6 :]
+        table = after_from.split(" ", 1)[0]
+        return f"SELECT {table}"

obsforge/instrumentation/db/sqlalchemy.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, TypeVar, cast
+
+from obsforge.instrumentation.db.engine import DBObservabilityEngine
+
+event: Any
+try:  # pragma: no cover - optional dependency
+    from sqlalchemy import event
+except ImportError:  # pragma: no cover - optional dependency
+    event = None
+
+
+_F = TypeVar("_F", bound=Callable[..., Any])
+
+
+def _listens_for(target: Any, identifier: str) -> Callable[[_F], _F]:
+    """Typed wrapper around ``sqlalchemy.event.listens_for``.
+
+    SQLAlchemy returns the decorated function unchanged, so this preserves
+    runtime behavior while giving mypy a typed decorator.
+    """
+
+    return cast("Callable[[_F], _F]", event.listens_for(target, identifier))
+
+
+if event is not None:
+    def instrument_sqlalchemy_engine(engine: Any, db_engine: DBObservabilityEngine) -> Any:
+        target = getattr(engine, "sync_engine", engine)
+        sa_pool = getattr(target, "pool", None)
+
+        @_listens_for(target, "before_cursor_execute")
+        def before_cursor_execute(
+            conn: Any,
+            cursor: Any,
+            statement: Any,
+            parameters: Any,
+            context: Any,
+            executemany: Any,
+        ) -> None:
+            conn.info["_obsforge_db_capture"] = db_engine.start_query(
+                system=target.dialect.name,
+                sql=statement,
+                alias=getattr(target.url, "database", None),
+                database_name=getattr(target.url, "database", None),
+                host=getattr(target.url, "host", None),
+                port=getattr(target.url, "port", None),
+                role="primary",
+                connection_id=str(id(conn.connection)),
+                pool_name=type(sa_pool).__name__ if sa_pool is not None else None,
+                many=executemany,
+                read_only=getattr(context, "isinsert", False) is False and getattr(context, "isupdate", False) is False,
+                namespace=getattr(target.url, "database", None),
+            )
+
+        @_listens_for(target, "after_cursor_execute")
+        def after_cursor_execute(
+            conn: Any,
+            cursor: Any,
+            statement: Any,
+            parameters: Any,
+            context: Any,
+            executemany: Any,
+        ) -> None:
+            capture = conn.info.pop("_obsforge_db_capture", None)
+            if capture is None:
+                return
+            db_engine.emit_query_sync(
+                capture,
+                rows_affected=getattr(cursor, "rowcount", None),
+                pool_snapshot=_pool_snapshot(db_engine, sa_pool),
+                transaction_state=db_engine.transaction_state(str(id(conn.connection))),
+            )
+
+        @_listens_for(target, "handle_error")
+        def handle_error(exception_context: Any) -> None:
+            conn = exception_context.connection
+            capture = conn.info.pop("_obsforge_db_capture", None) if conn is not None else None
+            if capture is None:
+                return
+            db_engine.emit_query_sync(
+                capture,
+                exception=exception_context.original_exception,
+                pool_snapshot=_pool_snapshot(db_engine, sa_pool),
+                transaction_state=(
+                    db_engine.transaction_state(str(id(conn.connection))) if conn is not None else None
+                ),
+            )
+
+        @_listens_for(target, "begin")
+        def begin(conn: Any) -> None:
+            db_engine.transaction_begin(
+                str(id(conn.connection)),
+                isolation_level=_safe_call(conn, "get_isolation_level"),
+                autocommit=_safe_getattr(conn.connection, "autocommit"),
+            )
+
+        @_listens_for(target, "commit")
+        def commit(conn: Any) -> None:
+            db_engine.transaction_commit(str(id(conn.connection)))
+
+        @_listens_for(target, "rollback")
+        def rollback(conn: Any) -> None:
+            db_engine.transaction_rollback(str(id(conn.connection)))
+
+        if sa_pool is not None:
+            @_listens_for(sa_pool, "checkout")
+            def checkout(
+                dbapi_connection: Any,
+                connection_record: Any,
+                connection_proxy: Any,
+            ) -> None:
+                db_engine.pool_checkout(
+                    pool_name=type(sa_pool).__name__,
+                    connection_id=str(id(dbapi_connection)),
+                    pool_size=_safe_call(sa_pool, "size"),
+                    checked_out=_safe_call(sa_pool, "checkedout"),
+                )
+
+            @_listens_for(sa_pool, "checkin")
+            def checkin(dbapi_connection: Any, connection_record: Any) -> None:
+                db_engine.pool_checkin(
+                    pool_name=type(sa_pool).__name__,
+                    connection_id=str(id(dbapi_connection)),
+                    pool_size=_safe_call(sa_pool, "size"),
+                    checked_out=_safe_call(sa_pool, "checkedout"),
+                )
+
+        return engine
+
+
+    def _pool_snapshot(db_engine: DBObservabilityEngine, pool: Any) -> Any:
+        if pool is None:
+            return None
+        return db_engine.pool_snapshot(
+            pool_name=type(pool).__name__,
+            pool_size=_safe_call(pool, "size"),
+            checked_out=_safe_call(pool, "checkedout"),
+        )
+
+
+    def _safe_call(obj: Any, name: str) -> Any:
+        method = getattr(obj, name, None)
+        if callable(method):
+            try:
+                return method()
+            except Exception:
+                return None
+        return None
+
+
+    def _safe_getattr(obj: Any, name: str) -> Any:
+        try:
+            return getattr(obj, name)
+        except Exception:
+            return None
+
+else:
+    def instrument_sqlalchemy_engine(engine: Any, db_engine: DBObservabilityEngine) -> Any:  # pragma: no cover
+        raise RuntimeError("sqlalchemy is not installed")

obsforge/instrumentation/db/state.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from obsforge.core.errors import ConfigurationError
+
+if TYPE_CHECKING:
+    from obsforge.instrumentation.db.engine import DBObservabilityEngine
+
+_default_db_engine: DBObservabilityEngine | None = None
+
+
+def set_default_db_engine(engine: DBObservabilityEngine) -> None:
+    global _default_db_engine
+    _default_db_engine = engine
+
+
+def get_default_db_engine() -> DBObservabilityEngine | None:
+    return _default_db_engine
+
+
+def require_default_db_engine() -> DBObservabilityEngine:
+    if _default_db_engine is None:
+        raise ConfigurationError(
+            "obsforge is not initialized: call obsforge.bootstrap() before using the "
+            "database instrumentation, or pass db_engine= explicitly."
+        )
+    return _default_db_engine

obsforge/instrumentation/db/transactions.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from threading import Lock
+from uuid import uuid4
+
+
+@dataclass(slots=True)
+class TransactionState:
+    transaction_id: str
+    depth: int
+    isolation_level: str | None
+    autocommit: bool | None
+    status: str
+
+
+class TransactionTracker:
+    def __init__(self) -> None:
+        self._states: dict[str, TransactionState] = {}
+        self._lock = Lock()
+
+    def begin(
+        self,
+        connection_id: str,
+        *,
+        isolation_level: str | None = None,
+        autocommit: bool | None = None,
+    ) -> TransactionState:
+        with self._lock:
+            current = self._states.get(connection_id)
+            if current is None:
+                current = TransactionState(
+                    transaction_id=str(uuid4()),
+                    depth=1,
+                    isolation_level=isolation_level,
+                    autocommit=autocommit,
+                    status="active",
+                )
+            else:
+                current.depth += 1
+                current.status = "active"
+                if isolation_level is not None:
+                    current.isolation_level = isolation_level
+                if autocommit is not None:
+                    current.autocommit = autocommit
+            self._states[connection_id] = current
+            return current
+
+    def commit(self, connection_id: str) -> TransactionState | None:
+        with self._lock:
+            current = self._states.get(connection_id)
+            if current is None:
+                return None
+            current.depth = max(current.depth - 1, 0)
+            current.status = "committed"
+            if current.depth == 0:
+                self._states.pop(connection_id, None)
+            return current
+
+    def rollback(self, connection_id: str) -> TransactionState | None:
+        with self._lock:
+            current = self._states.get(connection_id)
+            if current is None:
+                return None
+            current.depth = max(current.depth - 1, 0)
+            current.status = "rolled_back"
+            if current.depth == 0:
+                self._states.pop(connection_id, None)
+            return current
+
+    def state(self, connection_id: str) -> TransactionState | None:
+        with self._lock:
+            return self._states.get(connection_id)

obsforge/instrumentation/exceptions/__init__.py

@@ -0,0 +1,15 @@
+from obsforge.instrumentation.exceptions.engine import ExceptionIntelligenceEngine
+from obsforge.instrumentation.exceptions.state import (
+    ExceptionContextSnapshot,
+    ExceptionStateManager,
+    get_default_exception_engine,
+    set_default_exception_engine,
+)
+
+__all__ = [
+    "ExceptionContextSnapshot",
+    "ExceptionIntelligenceEngine",
+    "ExceptionStateManager",
+    "get_default_exception_engine",
+    "set_default_exception_engine",
+]