obsforge 0.1.1__py3-none-any.whl

obsforge/instrumentation/db/aiomysql.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from importlib.util import find_spec
+from typing import Any
+
+from obsforge.instrumentation.db.engine import DBObservabilityEngine
+
+_HAS_AIOMYSQL = find_spec("aiomysql") is not None
+
+
+class AioMySQLCursorProxy:
+    def __init__(self, cursor: Any, connection_proxy: AioMySQLConnectionProxy) -> None:
+        self._cursor = cursor
+        self._connection_proxy = connection_proxy
+
+    async def execute(self, query: str, args: Any = None) -> Any:
+        return await self._run("execute", query, args, many=False)
+
+    async def executemany(self, query: str, args: Any) -> Any:
+        return await self._run("executemany", query, args, many=True)
+
+    async def _run(self, method_name: str, query: str, args: Any, *, many: bool) -> Any:
+        capture = self._connection_proxy._db_engine.start_query(
+            system="mysql",
+            sql=query,
+            database_name=getattr(self._connection_proxy._connection, "db", None),
+            host=getattr(self._connection_proxy._connection, "host", None),
+            port=getattr(self._connection_proxy._connection, "port", None),
+            role="primary",
+            connection_id=str(id(self._connection_proxy._connection)),
+            pool_name="aiomysql",
+            many=many,
+            namespace=getattr(self._connection_proxy._connection, "db", None),
+        )
+        try:
+            result = await getattr(self._cursor, method_name)(query, args)
+        except BaseException as exc:
+            await self._connection_proxy._db_engine.emit_query(capture, exception=exc)
+            raise
+        await self._connection_proxy._db_engine.emit_query(
+            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)
+
+
+class AioMySQLConnectionProxy:
+    def __init__(self, connection: Any, db_engine: DBObservabilityEngine) -> None:
+        self._connection = connection
+        self._db_engine = db_engine
+
+    async def cursor(self, *args: Any, **kwargs: Any) -> AioMySQLCursorProxy:
+        cursor = await self._connection.cursor(*args, **kwargs)
+        return AioMySQLCursorProxy(cursor, self)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._connection, name)
+
+
+def instrument_aiomysql_connection(connection: Any, db_engine: DBObservabilityEngine) -> AioMySQLConnectionProxy:
+    return AioMySQLConnectionProxy(connection, db_engine)
+
+
+def instrument_aiomysql_pool(pool: Any, db_engine: DBObservabilityEngine) -> Any:
+    if not _HAS_AIOMYSQL:  # pragma: no cover
+        raise RuntimeError("aiomysql is not installed")
+
+    original_internal_acquire = getattr(pool, "_acquire", None)
+    original_acquire = pool.acquire
+    original_release = pool.release
+
+    async def _acquire(*args: Any, **kwargs: Any) -> AioMySQLConnectionProxy:
+        if callable(original_internal_acquire):
+            connection = await original_internal_acquire(*args, **kwargs)
+        else:
+            connection = await original_acquire(*args, **kwargs)
+        db_engine.pool_checkout(
+            pool_name="aiomysql",
+            connection_id=str(id(connection)),
+            pool_size=getattr(pool, "maxsize", None),
+            available=getattr(pool, "freesize", None),
+        )
+        return instrument_aiomysql_connection(connection, db_engine)
+
+    def release(connection: Any, *args: Any, **kwargs: Any) -> Any:
+        raw = getattr(connection, "_connection", connection)
+        db_engine.pool_checkin(
+            pool_name="aiomysql",
+            connection_id=str(id(raw)),
+            pool_size=getattr(pool, "maxsize", None),
+            available=getattr(pool, "freesize", None),
+        )
+        return original_release(raw, *args, **kwargs)
+
+    if callable(original_internal_acquire):
+        pool._acquire = _acquire
+    else:
+        pool.acquire = _acquire
+    pool.release = release
+    return pool

obsforge/instrumentation/db/asyncpg.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from types import ModuleType
+from typing import Any
+
+from obsforge.instrumentation.db.engine import DBObservabilityEngine
+
+try:  # pragma: no cover - optional dependency
+    import asyncpg as _asyncpg
+
+    asyncpg: ModuleType | None = _asyncpg
+except ImportError:  # pragma: no cover - optional dependency
+    asyncpg = None
+
+
+class AsyncPGConnectionProxy:
+    def __init__(self, connection: Any, db_engine: DBObservabilityEngine) -> None:
+        self._connection = connection
+        self._db_engine = db_engine
+
+    async def execute(self, query: str, *args: Any) -> Any:
+        return await self._run_query("execute", query, *args)
+
+    async def executemany(self, query: str, args: Any) -> Any:
+        return await self._run_query("executemany", query, args, many=True)
+
+    async def fetch(self, query: str, *args: Any) -> Any:
+        return await self._run_query("fetch", query, *args)
+
+    async def fetchrow(self, query: str, *args: Any) -> Any:
+        return await self._run_query("fetchrow", query, *args)
+
+    async def fetchval(self, query: str, *args: Any) -> Any:
+        return await self._run_query("fetchval", query, *args)
+
+    async def _run_query(self, method_name: str, query: str, *args: Any, many: bool = False) -> Any:
+        capture = self._db_engine.start_query(
+            system="postgresql",
+            sql=query,
+            database_name=getattr(self._connection._params, "database", None),
+            host=getattr(self._connection._params, "host", None),
+            port=getattr(self._connection._params, "port", None),
+            role="primary",
+            connection_id=str(id(self._connection)),
+            pool_name="asyncpg",
+            many=many,
+            namespace=getattr(self._connection._params, "database", None),
+        )
+        try:
+            result = await getattr(self._connection, method_name)(query, *args)
+        except BaseException as exc:
+            await self._db_engine.emit_query(capture, exception=exc)
+            raise
+        rows_affected = len(result) if isinstance(result, list) else None
+        await self._db_engine.emit_query(
+            capture,
+            rows_affected=rows_affected,
+            transaction_state=self._db_engine.transaction_state(str(id(self._connection))),
+        )
+        return result
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._connection, name)
+
+
+def instrument_asyncpg_connection(connection: Any, db_engine: DBObservabilityEngine) -> AsyncPGConnectionProxy:
+    return AsyncPGConnectionProxy(connection, db_engine)
+
+
+if asyncpg is not None:
+    def instrument_asyncpg_pool(pool: Any, db_engine: DBObservabilityEngine) -> Any:
+        original_internal_acquire = getattr(pool, "_acquire", None)
+        original_acquire = pool.acquire
+        original_release = pool.release
+
+        async def _acquire(*args: Any, **kwargs: Any) -> Any:
+            if callable(original_internal_acquire):
+                connection = await original_internal_acquire(*args, **kwargs)
+            else:
+                connection = await original_acquire(*args, **kwargs)
+            db_engine.pool_checkout(
+                pool_name="asyncpg",
+                connection_id=str(id(connection)),
+                pool_size=getattr(pool, "_maxsize", None),
+                checked_out=len(getattr(pool, "_holders", [])) if hasattr(pool, "_holders") else None,
+            )
+            return instrument_asyncpg_connection(connection, db_engine)
+
+        async def release(connection: Any, *args: Any, **kwargs: Any) -> Any:
+            raw = getattr(connection, "_connection", connection)
+            db_engine.pool_checkin(
+                pool_name="asyncpg",
+                connection_id=str(id(raw)),
+                pool_size=getattr(pool, "_maxsize", None),
+            )
+            return await original_release(raw, *args, **kwargs)
+
+        if callable(original_internal_acquire):
+            pool._acquire = _acquire
+        else:
+            pool.acquire = _acquire
+        pool.release = release
+        return pool
+
+else:
+    def instrument_asyncpg_pool(pool: Any, db_engine: DBObservabilityEngine) -> Any:  # pragma: no cover
+        raise RuntimeError("asyncpg is not installed")

obsforge/instrumentation/db/django.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from contextlib import ExitStack, contextmanager
+from typing import Any
+
+from obsforge.instrumentation.db.engine import DBObservabilityEngine
+
+
+class DjangoQueryWrapper:
+    def __init__(self, db_engine: DBObservabilityEngine, alias: str) -> None:
+        self._engine = db_engine
+        self._alias = alias
+
+    def __call__(
+        self,
+        execute: Callable[[str, Any, bool, dict[str, Any]], Any],
+        sql: str,
+        params: Any,
+        many: bool,
+        context: dict[str, Any],
+    ) -> Any:
+        connection = context["connection"]
+        capture = self._engine.start_query(
+            system=connection.vendor,
+            sql=sql,
+            alias=self._alias,
+            database_name=_settings_dict(connection, "NAME"),
+            host=_settings_dict(connection, "HOST"),
+            port=_settings_dict(connection, "PORT"),
+            role="primary",
+            connection_id=str(id(connection.connection)),
+            pool_name=f"django:{self._alias}",
+            many=many,
+            read_only=sql.lstrip().upper().startswith("SELECT"),
+            namespace=_settings_dict(connection, "NAME"),
+        )
+        try:
+            result = execute(sql, params, many, context)
+        except BaseException as exc:
+            self._engine.emit_query_sync(
+                capture,
+                exception=exc,
+                transaction_state=self._engine.transaction_state(str(id(connection.connection))),
+            )
+            raise
+
+        cursor = context.get("cursor")
+        self._engine.emit_query_sync(
+            capture,
+            rows_affected=getattr(cursor, "rowcount", None),
+            transaction_state=self._engine.transaction_state(str(id(connection.connection))),
+        )
+        return result
+
+
+@contextmanager
+def instrument_django_connections(
+    connections: Any, db_engine: DBObservabilityEngine
+) -> Iterator[None]:
+    with ExitStack() as stack:
+        for alias in connections:
+            connection = connections[alias]
+            stack.enter_context(connection.execute_wrapper(DjangoQueryWrapper(db_engine, alias)))
+        yield
+
+
+def _settings_dict(connection: Any, key: str) -> Any:
+    settings_dict = getattr(connection, "settings_dict", {})
+    return settings_dict.get(key)

obsforge/instrumentation/db/engine.py

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any
+
+from obsforge.api.events import EventClient
+from obsforge.config.settings import DatabaseInstrumentationSettings, ObsforgeSettings
+from obsforge.core.models import (
+    BusinessContext,
+    CanonicalEvent,
+    CorrelationContext,
+    DatabaseContext,
+    DatabaseRole,
+    DependencyContext,
+    Event,
+    EventKind,
+    Outcome,
+    PerformanceContext,
+    Severity,
+    SpanKind,
+    TelemetryContext,
+)
+from obsforge.instrumentation.db.pool import PoolSnapshot, PoolTracker
+from obsforge.instrumentation.db.sql import SQLNormalizer
+from obsforge.instrumentation.db.transactions import TransactionState, TransactionTracker
+from obsforge.instrumentation.exceptions.engine import ExceptionIntelligenceEngine
+from obsforge.instrumentation.exceptions.state import ExceptionContextSnapshot
+from obsforge.propagation.correlation import CorrelationManager
+
+
+@dataclass(slots=True)
+class QueryCapture:
+    started_at: float
+    database: DatabaseContext
+    normalized_sql: str
+    prior_exception_snapshot: ExceptionContextSnapshot | None
+
+
+class DBObservabilityEngine:
+    def __init__(
+        self,
+        *,
+        event_client: EventClient,
+        correlation_manager: CorrelationManager,
+        settings: ObsforgeSettings,
+        exception_engine: ExceptionIntelligenceEngine | None = None,
+    ) -> None:
+        self._event_client = event_client
+        self._correlation_manager = correlation_manager
+        self._settings = settings
+        self._db_settings: DatabaseInstrumentationSettings = settings.database
+        self._normalizer = SQLNormalizer(settings.database)
+        self._pool = PoolTracker(settings.database)
+        self._transactions = TransactionTracker()
+        self._exception_engine = exception_engine
+
+    def start_query(
+        self,
+        *,
+        system: str,
+        sql: str,
+        alias: str | None = None,
+        database_name: str | None = None,
+        host: str | None = None,
+        port: int | None = None,
+        role: str | None = None,
+        connection_id: str | None = None,
+        pool_name: str | None = None,
+        timeout_ms: float | None = None,
+        many: bool | None = None,
+        read_only: bool | None = None,
+        namespace: str | None = None,
+    ) -> QueryCapture:
+        normalized = self._normalizer.normalize(sql)
+        database = DatabaseContext(
+            system=system,
+            namespace=namespace,
+            database_name=database_name,
+            host=host,
+            port=port,
+            role=(
+                DatabaseRole(role)
+                if role is not None and role in DatabaseRole._value2member_map_
+                else DatabaseRole.UNKNOWN
+            ),
+            alias=alias,
+            operation_name=self._normalizer.operation(sql),
+            query_summary=self._normalizer.summary(sql),
+            query_fingerprint=self._normalizer.fingerprint(sql),
+            query_text_redacted=normalized,
+            connection_id=connection_id,
+            pool_name=pool_name,
+            timeout_ms=timeout_ms,
+            many=many,
+            read_only=read_only,
+        )
+        transaction = self._transactions.state(connection_id or "default")
+        if transaction is not None:
+            database = self._merge_transaction(database, transaction)
+
+        prior_snapshot: ExceptionContextSnapshot | None = None
+        if self._exception_engine is not None:
+            prior_snapshot = self._exception_engine.snapshot_context()
+            self._exception_engine.bind_database_context(database)
+        return QueryCapture(
+            started_at=time.perf_counter(),
+            database=database,
+            normalized_sql=normalized,
+            prior_exception_snapshot=prior_snapshot,
+        )
+
+    def finish_query(
+        self,
+        capture: QueryCapture,
+        *,
+        rows_affected: int | None = None,
+        retry_count: int | None = None,
+        exception: BaseException | None = None,
+        pool_snapshot: PoolSnapshot | None = None,
+        transaction_state: TransactionState | None = None,
+        lock_contended: bool = False,
+        deadlock_detected: bool = False,
+    ) -> Event:
+        duration_ms = (time.perf_counter() - capture.started_at) * 1000
+        database = capture.database.model_copy()
+        database.rows_affected = rows_affected
+        database.retry_count = retry_count
+        database.retry_storm = (retry_count or 0) >= self._db_settings.retry_storm_threshold
+        database.deadlock_detected = deadlock_detected or _detect_deadlock(exception)
+        database.lock_contended = lock_contended or _detect_lock_contention(exception)
+        if pool_snapshot is not None:
+            database.pool_name = pool_snapshot.pool_name
+            database.pool_size = pool_snapshot.pool_size
+            database.pool_checked_out = pool_snapshot.checked_out
+            database.pool_available = pool_snapshot.available
+            database.pool_waiters = pool_snapshot.waiters
+            database.pool_saturated = pool_snapshot.saturated
+            database.connection_leak_suspected = pool_snapshot.leak_suspected
+            database.checkout_wait_ms = pool_snapshot.checkout_wait_ms
+        if transaction_state is not None:
+            database = self._merge_transaction(database, transaction_state)
+
+        performance = PerformanceContext(
+            duration_ms=duration_ms,
+            slow_threshold_ms=self._db_settings.slow_query_threshold_ms,
+            sampled=True,
+        )
+        outcome = Outcome.SUCCEEDED
+        severity = Severity.INFO
+        ai_hints: list[str] = []
+        if duration_ms >= self._db_settings.slow_query_threshold_ms:
+            ai_hints.append("slow_query")
+            severity = Severity.WARNING
+            outcome = Outcome.DEGRADED
+        if database.deadlock_detected:
+            ai_hints.append("deadlock")
+            severity = Severity.ERROR
+            outcome = Outcome.FAILED
+        if database.pool_saturated:
+            ai_hints.append("pool_exhaustion")
+            severity = Severity.WARNING if severity == Severity.INFO else severity
+            if outcome is Outcome.SUCCEEDED:
+                outcome = Outcome.DEGRADED
+        if database.connection_leak_suspected:
+            ai_hints.append("leaked_connection")
+            severity = Severity.WARNING if severity == Severity.INFO else severity
+        if database.retry_storm:
+            ai_hints.append("retry_storm")
+            severity = Severity.WARNING if severity == Severity.INFO else severity
+        if database.lock_contended:
+            ai_hints.append("lock_contention")
+            severity = Severity.WARNING if severity == Severity.INFO else severity
+        exception_context = None
+        if exception is not None:
+            outcome = Outcome.TIMEOUT if "timeout" in str(exception).lower() else Outcome.FAILED
+            severity = Severity.ERROR
+            if self._exception_engine is not None:
+                exception_context = self._exception_engine.capture_exception_context(
+                    exception,
+                    handled=True,
+                    escaped=False,
+                    database=database,
+                )
+                for hint in exception_context.ai_hints:
+                    if hint not in ai_hints:
+                        ai_hints.append(hint)
+
+        event = CanonicalEvent(
+            severity=severity,
+            event_name="db.query.completed",
+            event_kind=EventKind.DB,
+            event_domain="db",
+            operation=database.operation_name.lower() if database.operation_name else "query",
+            outcome=outcome,
+            message=f"{database.operation_name or 'QUERY'} {database.query_summary or database.system or 'db'}",
+            correlation=self._correlation_manager.get_correlation()
+            or CorrelationContext(correlation_id=str(time.time_ns())),
+            telemetry=self._telemetry_context(outcome),
+            database=database,
+            business=self._bound_business_context(),
+            dependency=self._bound_dependency_context(),
+            exception=exception_context,
+            performance=performance,
+            tags={"source_kind": "database"},
+            attributes={
+                "source.service.name": self._settings.service_name,
+                "db.host": database.host,
+                "db.role": database.role.value if database.role is not None else DatabaseRole.UNKNOWN.value,
+            },
+            investigation_keys=_investigation_keys(database),
+            ai_hints=ai_hints,
+        )
+
+        if self._exception_engine is not None and capture.prior_exception_snapshot is not None:
+            self._exception_engine.restore_context(capture.prior_exception_snapshot)
+        return event
+
+    async def emit_query(self, capture: QueryCapture, **kwargs: Any) -> Event:
+        event = self.finish_query(capture, **kwargs)
+        await self._event_client.emit(event)
+        return event
+
+    def emit_query_sync(self, capture: QueryCapture, **kwargs: Any) -> Event:
+        event = self.finish_query(capture, **kwargs)
+        self._event_client.emit_sync(event)
+        return event
+
+    def pool_checkout(self, *, pool_name: str, connection_id: str, **kwargs: Any) -> PoolSnapshot:
+        return self._pool.checkout(pool_name=pool_name, connection_id=connection_id, **kwargs)
+
+    def pool_checkin(self, *, pool_name: str, connection_id: str, **kwargs: Any) -> PoolSnapshot:
+        return self._pool.checkin(pool_name=pool_name, connection_id=connection_id, **kwargs)
+
+    def pool_snapshot(self, *, pool_name: str, **kwargs: Any) -> PoolSnapshot:
+        return self._pool.snapshot(pool_name=pool_name, **kwargs)
+
+    def transaction_begin(
+        self,
+        connection_id: str,
+        *,
+        isolation_level: str | None = None,
+        autocommit: object = None,
+    ) -> TransactionState:
+        return self._transactions.begin(
+            connection_id,
+            isolation_level=isolation_level,
+            autocommit=_coerce_autocommit(autocommit),
+        )
+
+    def transaction_commit(self, connection_id: str) -> TransactionState | None:
+        return self._transactions.commit(connection_id)
+
+    def transaction_rollback(self, connection_id: str) -> TransactionState | None:
+        return self._transactions.rollback(connection_id)
+
+    def transaction_state(self, connection_id: str) -> TransactionState | None:
+        return self._transactions.state(connection_id)
+
+    def bind_business_context(self, business: BusinessContext | None) -> None:
+        if self._exception_engine is not None:
+            self._exception_engine.bind_business_context(business)
+
+    def _merge_transaction(self, database: DatabaseContext, tx: TransactionState) -> DatabaseContext:
+        database.transaction_id = tx.transaction_id
+        database.transaction_depth = tx.depth
+        database.transaction_active = tx.depth > 0
+        database.transaction_status = tx.status
+        database.transaction_isolation_level = tx.isolation_level
+        database.autocommit = tx.autocommit
+        return database
+
+    def _telemetry_context(self, outcome: Outcome) -> TelemetryContext:
+        return TelemetryContext(
+            trace=self._correlation_manager.get_trace(),
+            span_kind=SpanKind.INTERNAL,
+            otel_status_code="ERROR" if outcome in {Outcome.FAILED, Outcome.TIMEOUT} else "OK",
+            telemetry_sdk_version=self._settings.telemetry_sdk_version,
+        )
+
+    def _bound_business_context(self) -> BusinessContext | None:
+        if self._exception_engine is None:
+            return None
+        return self._exception_engine.state.get_business_context()
+
+    def _bound_dependency_context(self) -> DependencyContext | None:
+        if self._exception_engine is None:
+            return None
+        return self._exception_engine.state.get_dependency_context()
+
+
+def _detect_deadlock(exc: BaseException | None) -> bool:
+    if exc is None:
+        return False
+    text = str(exc).lower()
+    return "deadlock" in text
+
+
+def _detect_lock_contention(exc: BaseException | None) -> bool:
+    if exc is None:
+        return False
+    text = str(exc).lower()
+    return "lock wait" in text or "could not obtain lock" in text or "database is locked" in text
+
+
+def _investigation_keys(database: DatabaseContext) -> list[str]:
+    keys = [database.system or "unknown-db"]
+    if database.operation_name is not None:
+        keys.append(database.operation_name)
+    if database.query_fingerprint is not None:
+        keys.append(database.query_fingerprint)
+    if database.host is not None:
+        keys.append(database.host)
+    return keys
+
+
+def _coerce_autocommit(value: object) -> bool | None:
+    # Drivers may report legacy/sentinel ints (e.g. sqlite's -1). Normalize so
+    # the model never serializes a non-bool into the bool-typed field.
+    if value is None or isinstance(value, bool):
+        return value
+    if isinstance(value, int) and value in (0, 1):
+        return bool(value)
+    return None