feat(postgres): Implement Database Instrumentation Suppression

backend/app/celery/tasks/job_execution.py

@@ -22,6 +22,7 @@
 from app.celery.celery_app import celery_app
 from app.celery.utils import gevent_timeout
 from app.core.config import settings
+from app.core.telemetry import suppress_db_instrumentation
 
 logger = logging.getLogger(__name__)
 
@@ -82,16 +83,19 @@ def run_llm_job(self, project_id: int, job_id: str, trace_id: str, **kwargs):
     from app.services.llm.jobs import execute_job
 
     _set_trace(trace_id)
-    return _run_with_otel_parent(
-        self,
-        lambda: execute_job(
-            project_id=project_id,
-            job_id=job_id,
-            task_id=current_task.request.id,
-            task_instance=self,
-            **kwargs,
-        ),
-    )
+    # DB spans suppressed job-wide so LLM waterfalls stay clean (drops these queries
+    # from the Sentry Queries page too — accepted trade-off).
+    with suppress_db_instrumentation():
+        return _run_with_otel_parent(
+            self,
+            lambda: execute_job(
+                project_id=project_id,
+                job_id=job_id,
+                task_id=current_task.request.id,
+                task_instance=self,
+                **kwargs,
+            ),
+        )
 
 
 @celery_app.task(bind=True, queue="default", priority=9)
@@ -100,16 +104,17 @@ def run_llm_chain_job(self, project_id: int, job_id: str, trace_id: str, **kwarg
     from app.services.llm.jobs import execute_chain_job
 
     _set_trace(trace_id)
-    return _run_with_otel_parent(
-        self,
-        lambda: execute_chain_job(
-            project_id=project_id,
-            job_id=job_id,
-            task_id=current_task.request.id,
-            task_instance=self,
-            **kwargs,
-        ),
-    )
+    with suppress_db_instrumentation():
+        return _run_with_otel_parent(
+            self,
+            lambda: execute_chain_job(
+                project_id=project_id,
+                job_id=job_id,
+                task_id=current_task.request.id,
+                task_instance=self,
+                **kwargs,
+            ),
+        )
 
 
 @celery_app.task(bind=True, queue="default", priority=9)
@@ -118,16 +123,17 @@ def run_response_job(self, project_id: int, job_id: str, trace_id: str, **kwargs
     from app.services.response.jobs import execute_job
 
     _set_trace(trace_id)
-    return _run_with_otel_parent(
-        self,
-        lambda: execute_job(
-            project_id=project_id,
-            job_id=job_id,
-            task_id=current_task.request.id,
-            task_instance=self,
-            **kwargs,
-        ),
-    )
+    with suppress_db_instrumentation():
+        return _run_with_otel_parent(
+            self,
+            lambda: execute_job(
+                project_id=project_id,
+                job_id=job_id,
+                task_id=current_task.request.id,
+                task_instance=self,
+                **kwargs,
+            ),
+        )
 
 
 @celery_app.task(bind=True, queue="default", priority=9)

backend/app/core/telemetry.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import time
 from collections.abc import Iterator
 from contextlib import contextmanager
 from contextvars import ContextVar
@@ -24,10 +23,37 @@
 
 logger = logging.getLogger(__name__)
 
+# Postgres SQLSTATE codes surfaced as named Sentry tags; others pass through as the raw code.
+NOTABLE_SQLSTATES: dict[str, str] = {
+    "40P01": "deadlock_detected",
+    "57014": "query_canceled",  # statement timeout
+    "40001": "serialization_failure",
+    "55P03": "lock_not_available",  # lock timeout
+    "08006": "connection_failure",
+    "08003": "connection_does_not_exist",
+    "53300": "too_many_connections",
+}
+
 _log_context_var: ContextVar[dict[str, str] | None] = ContextVar(
     "kaapi_log_context", default=None
 )
 
+# OTel instrumentation scope emitted by SQLAlchemyInstrumentor; used to filter its spans.
+_SQLALCHEMY_SCOPE = "opentelemetry.instrumentation.sqlalchemy"
+
+# When True in the current context, SQLAlchemy DB spans are dropped before reaching Sentry.
+_suppress_db_spans_var: ContextVar[bool] = ContextVar(
+    "kaapi_suppress_db_spans", default=False
+)
+
+
+def _should_drop_db_span(otel_span: object) -> bool:
+    """True when DB-span suppression is active and `otel_span` is a SQLAlchemy span."""
+    if not _suppress_db_spans_var.get():
+        return False
+    scope = getattr(otel_span, "instrumentation_scope", None)
+    return scope is not None and getattr(scope, "name", None) == _SQLALCHEMY_SCOPE
+
 
 def _emit_sentry_metric(
     metric_type: str,
@@ -194,7 +220,20 @@ def setup_telemetry(service_name: str | None = None) -> None:
     if settings.SENTRY_DSN:
         from sentry_sdk.integrations.opentelemetry import SentrySpanProcessor
 
-        tracer_provider.add_span_processor(SentrySpanProcessor())
+        class _DbSpanFilteringProcessor(SentrySpanProcessor):
+            """Drop SQLAlchemy DB spans from Sentry while suppress_db_instrumentation() is active.
+
+            Filters by instrumentation scope so only DB spans are skipped — HTTP/Requests
+            spans keep flowing. Skipping on_start leaves the span unmapped, and the parent
+            on_end safely no-ops on unmapped spans.
+            """
+
+            def on_start(self, otel_span, parent_context=None):  # type: ignore[override]
+                if _should_drop_db_span(otel_span):
+                    return
+                super().on_start(otel_span, parent_context)
+
+        tracer_provider.add_span_processor(_DbSpanFilteringProcessor())
 
     trace.set_tracer_provider(tracer_provider)
 
@@ -358,30 +397,56 @@ def suppress_http_instrumentation() -> Iterator[None]:
         otel_context.detach(token)
 
 
-def record_db_query_finished(
+@contextmanager
+def suppress_db_instrumentation() -> Iterator[None]:
+    """Drop SQLAlchemy DB spans from the Sentry trace for the wrapped block.
+
+    Wrap LLM job execution so its DB reads/writes do not clutter the LLM waterfall.
+    Only DB spans are filtered (by _DbSpanFilteringProcessor via instrumentation
+    scope) — HTTP/Requests instrumentation stays active. Trade-off: the wrapped
+    DB queries also drop from the Sentry Queries page.
+    """
+    token = _suppress_db_spans_var.set(True)
+    try:
+        yield
+    finally:
+        _suppress_db_spans_var.reset(token)
+
+
+def record_db_query_failed(
     *,
-    duration_ms: float,
     operation: str | None = None,
-    error: bool = False,
+    sqlstate: str | None = None,
 ) -> None:
-    """Emit DB query metrics to Sentry."""
+    """Emit a DB query-failure counter to Sentry. Per-query duration/throughput come from spans."""
     if not settings.OTEL_ENABLED:
         return
 
-    attrs: dict[str, str] = {}
+    attrs: dict[str, str | int | float] = {}
     if operation:
         attrs["db.operation"] = operation
+    if sqlstate:
+        attrs["db.sqlstate"] = sqlstate
 
-    _emit_sentry_metric("count", "db.query.total", 1, attributes=attrs)
-    _emit_sentry_metric(
-        "distribution",
-        "db.query.duration",
-        duration_ms,
-        unit="millisecond",
-        attributes=attrs,
-    )
-    if error:
-        _emit_sentry_metric("count", "db.query.failed", 1, attributes=attrs)
+    _emit_sentry_metric("count", "db.query.failed", 1, attributes=attrs)
+
+
+def _tag_db_error(sqlstate: str | None) -> None:
+    """Tag the active Sentry scope + span with a DB error's Postgres SQLSTATE."""
+    if not sqlstate:
+        return
+    try:
+        span = trace.get_current_span()
+        if span.is_recording():
+            span.set_attribute("db.sqlstate", sqlstate)
+        if sentry_sdk.get_client().is_active():
+            sentry_sdk.set_tag("db.system", "postgresql")
+            sentry_sdk.set_tag("db.sqlstate", sqlstate)
+            name = NOTABLE_SQLSTATES.get(sqlstate)
+            if name:
+                sentry_sdk.set_tag("db.error.name", name)
+    except Exception:
+        logger.debug("[_tag_db_error] Failed to tag DB error | sqlstate: %s", sqlstate)
 
 
 def record_db_pool_stats(
@@ -523,6 +588,17 @@ def instrument_db_engine(engine: object) -> None:
     if getattr(engine, "_kaapi_db_telemetry_instrumented", False):
         return
 
+    # DB query spans -> SentrySpanProcessor -> Sentry Insights/Queries. ProxyTracer defers
+    # to the real provider set later in setup_telemetry(), so load order here is safe.
+    try:
+        from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+
+        SQLAlchemyInstrumentor().instrument(engine=engine)
+    except Exception:
+        logger.exception(
+            "[instrument_db_engine] Failed to load SQLAlchemy span instrumentation"
+        )
+
     try:
         from sqlalchemy import event
     except Exception:
@@ -551,7 +627,7 @@ def _before_cursor_execute(
         conn, cursor, statement, parameters, context, executemany
     ) -> None:
         del cursor, parameters, executemany
-        context._kaapi_db_started_at = time.perf_counter()
+        # Operation kept for error attribution; per-query timing now lives on the span.
         context._kaapi_db_operation = (
             str(statement).split(None, 1)[0].upper() if statement else "UNKNOWN"
         )
@@ -561,30 +637,20 @@ def _before_cursor_execute(
     def _after_cursor_execute(
         conn, cursor, statement, parameters, context, executemany
     ) -> None:
-        del cursor, statement, parameters, executemany
-        started_at = getattr(context, "_kaapi_db_started_at", None)
-        duration_ms = (
-            (time.perf_counter() - started_at) * 1000 if started_at is not None else 0.0
-        )
-        operation = getattr(context, "_kaapi_db_operation", None)
-        record_db_query_finished(
-            duration_ms=duration_ms, operation=operation, error=False
-        )
+        del cursor, statement, parameters, context, executemany
         _emit_pool_metrics(conn.engine.pool)
 
     @event.listens_for(engine, "handle_error")
     def _handle_error(exception_context) -> None:
         context = exception_context.execution_context
-        if context is None:
-            return
-        started_at = getattr(context, "_kaapi_db_started_at", None)
-        duration_ms = (
-            (time.perf_counter() - started_at) * 1000 if started_at is not None else 0.0
-        )
-        operation = getattr(context, "_kaapi_db_operation", None)
-        record_db_query_finished(
-            duration_ms=duration_ms, operation=operation, error=True
+        operation = (
+            getattr(context, "_kaapi_db_operation", None)
+            if context is not None
+            else None
         )
+        sqlstate = getattr(exception_context.original_exception, "sqlstate", None)
+        _tag_db_error(sqlstate)
+        record_db_query_failed(operation=operation, sqlstate=sqlstate)
 
     @event.listens_for(engine.pool, "checkout")
     def _on_checkout(dbapi_connection, connection_record, connection_proxy) -> None:

backend/app/tests/celery/test_job_execution.py

@@ -0,0 +1,69 @@
+"""Tests for the LLM/response Celery task wrappers in tasks/job_execution.py.
+
+The service entrypoints are mocked; no DB, OTel provider, or broker is used.
+Celery tasks are callable and run synchronously when invoked directly (self is
+bound), so we drive the real wrapper — including the gevent_timeout decorator
+and suppress_db_instrumentation — end to end.
+"""
+
+from unittest.mock import patch
+
+import pytest
+
+from app.celery.tasks import job_execution
+from app.core import telemetry
+
+SENTINEL = object()
+
+
+def _capture_suppression(captured: dict):
+    """execute_* stand-in that records whether DB spans are suppressed at call time."""
+
+    def _exec(**kwargs):
+        captured["suppressed"] = telemetry._suppress_db_spans_var.get()
+        return SENTINEL
+
+    return _exec
+
+
+@pytest.mark.parametrize(
+    ("task_name", "service_target"),
+    [
+        ("run_llm_job", "app.services.llm.jobs.execute_job"),
+        ("run_llm_chain_job", "app.services.llm.jobs.execute_chain_job"),
+        ("run_response_job", "app.services.response.jobs.execute_job"),
+    ],
+)
+class TestJobWrappers:
+    def test_returns_service_result(self, task_name, service_target):
+        task = getattr(job_execution, task_name)
+        captured: dict = {}
+        with (
+            patch.object(job_execution, "_set_trace", lambda trace_id: None),
+            patch(service_target, _capture_suppression(captured)),
+        ):
+            result = task(project_id=1, job_id="job-1", trace_id="trace-1")
+
+        assert result is SENTINEL
+
+    def test_db_spans_suppressed_during_execution(self, task_name, service_target):
+        task = getattr(job_execution, task_name)
+        captured: dict = {}
+        with (
+            patch.object(job_execution, "_set_trace", lambda trace_id: None),
+            patch(service_target, _capture_suppression(captured)),
+        ):
+            task(project_id=1, job_id="job-1", trace_id="trace-1")
+
+        assert captured["suppressed"] is True
+
+    def test_suppression_resets_after_return(self, task_name, service_target):
+        task = getattr(job_execution, task_name)
+        captured: dict = {}
+        with (
+            patch.object(job_execution, "_set_trace", lambda trace_id: None),
+            patch(service_target, _capture_suppression(captured)),
+        ):
+            task(project_id=1, job_id="job-1", trace_id="trace-1")
+
+        assert telemetry._suppress_db_spans_var.get() is False