From dcacbbc505696be8905b2a223380af48ac05abaa Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 14 Aug 2026 09:22:33 +0530 Subject: [PATCH 1/5] feat(telemetry): Implement database instrumentation suppression for LLM job execution and add SQLAlchemy instrumentation dependency --- backend/app/celery/tasks/job_execution.py | 66 +++++++------ backend/app/core/telemetry.py | 113 +++++++++++++++------- backend/pyproject.toml | 1 + backend/uv.lock | 20 +++- 4 files changed, 133 insertions(+), 67 deletions(-) diff --git a/backend/app/celery/tasks/job_execution.py b/backend/app/celery/tasks/job_execution.py index 1a84f9eb5..1a8892f5d 100644 --- a/backend/app/celery/tasks/job_execution.py +++ b/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) diff --git a/backend/app/core/telemetry.py b/backend/app/core/telemetry.py index 99d2fc959..cc2ddab6f 100644 --- a/backend/app/core/telemetry.py +++ b/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 @@ -13,7 +12,10 @@ from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor from opentelemetry.instrumentation.logging import LoggingInstrumentor from opentelemetry.instrumentation.requests import RequestsInstrumentor -from opentelemetry.instrumentation.utils import _SUPPRESS_HTTP_INSTRUMENTATION_KEY +from opentelemetry.instrumentation.utils import ( + _SUPPRESS_HTTP_INSTRUMENTATION_KEY, + _SUPPRESS_INSTRUMENTATION_KEY, +) from opentelemetry.sdk.resources import SERVICE_NAME, Resource from opentelemetry.sdk.trace import TracerProvider @@ -24,6 +26,17 @@ 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 ) @@ -358,30 +371,57 @@ def suppress_http_instrumentation() -> Iterator[None]: otel_context.detach(token) -def record_db_query_finished( +@contextmanager +def suppress_db_instrumentation() -> Iterator[None]: + """Suppress OTel auto-instrumentation (incl. SQLAlchemy DB spans) for the wrapped block. + + Wrap LLM job execution so its DB reads/writes do not clutter the LLM waterfall. + Trade-off: those queries also drop from the Sentry Queries page. Manually-created + spans (e.g. the gen_ai span) are unaffected — only auto-instrumentors honor this key. + """ + token = otel_context.attach( + otel_context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) + ) + try: + yield + finally: + otel_context.detach(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 +563,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 +602,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 +612,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: diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 8d6219f76..0f0d94572 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "opentelemetry-instrumentation-httpx>=0.51b0", "opentelemetry-instrumentation-requests>=0.51b0", "opentelemetry-instrumentation-logging>=0.51b0", + "opentelemetry-instrumentation-sqlalchemy>=0.51b0", "pyjwt>=2.13.0,<3.0.0", "boto3>=1.37.20", "moto[s3]>=5.1.1", diff --git a/backend/uv.lock b/backend/uv.lock index 2ab854e81..e6d41531f 100644 --- a/backend/uv.lock +++ b/backend/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.12, <4.0" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -285,6 +285,7 @@ dependencies = [ { name = "opentelemetry-instrumentation-httpx" }, { name = "opentelemetry-instrumentation-logging" }, { name = "opentelemetry-instrumentation-requests" }, + { name = "opentelemetry-instrumentation-sqlalchemy" }, { name = "opentelemetry-sdk" }, { name = "pandas" }, { name = "passlib", extra = ["bcrypt"] }, @@ -352,6 +353,7 @@ requires-dist = [ { name = "opentelemetry-instrumentation-httpx", specifier = ">=0.51b0" }, { name = "opentelemetry-instrumentation-logging", specifier = ">=0.51b0" }, { name = "opentelemetry-instrumentation-requests", specifier = ">=0.51b0" }, + { name = "opentelemetry-instrumentation-sqlalchemy", specifier = ">=0.51b0" }, { name = "opentelemetry-sdk", specifier = ">=1.30.0" }, { name = "pandas", specifier = ">=2.3.2" }, { name = "passlib", extras = ["bcrypt"], specifier = ">=1.7.4,<2.0.0" }, @@ -2590,6 +2592,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/97/54/a73d688d969283f344de96c43366912addb94bdbef413bcebac22979545e/opentelemetry_instrumentation_requests-0.62b0-py3-none-any.whl", hash = "sha256:edf61785ecb3ec6923e33c24074c82067f286a418f817b2b82546956d120e6d6", size = 14209, upload-time = "2026-04-09T14:40:02.987Z" }, ] +[[package]] +name = "opentelemetry-instrumentation-sqlalchemy" +version = "0.62b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-instrumentation" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "packaging" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2a/3d/40adc8c38e5be017ceb230a28ca57ca81981d4dc0c4b902cc930c77fd14f/opentelemetry_instrumentation_sqlalchemy-0.62b0.tar.gz", hash = "sha256:d02f85b83f349e9ef70a34cb3f4c3a3481fa15b11747f09209818663e161cac4", size = 18539, upload-time = "2026-04-09T14:40:50.251Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/e0/77954ac593f34740dc32e28a15fe7170e90f6ba6398eaaa5c88b34c05ed1/opentelemetry_instrumentation_sqlalchemy-0.62b0-py3-none-any.whl", hash = "sha256:ec576e0660080d9d15ce4fa44d2a07fff8cb4b796a84344cb0f2c9e5d6e26f79", size = 15534, upload-time = "2026-04-09T14:40:03.957Z" }, +] + [[package]] name = "opentelemetry-proto" version = "1.41.0" From d020fa5e36a0661097ff0c20990c27318cbc983c Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:17:27 +0530 Subject: [PATCH 2/5] feat(telemetry): Implement DB span filtering for Sentry in SQLAlchemy instrumentation --- backend/app/core/telemetry.py | 49 ++++-- backend/app/tests/core/test_telemetry.py | 185 +++++++++++++++++++++++ 2 files changed, 222 insertions(+), 12 deletions(-) create mode 100644 backend/app/tests/core/test_telemetry.py diff --git a/backend/app/core/telemetry.py b/backend/app/core/telemetry.py index cc2ddab6f..2a1cd3afb 100644 --- a/backend/app/core/telemetry.py +++ b/backend/app/core/telemetry.py @@ -12,10 +12,7 @@ from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor from opentelemetry.instrumentation.logging import LoggingInstrumentor from opentelemetry.instrumentation.requests import RequestsInstrumentor -from opentelemetry.instrumentation.utils import ( - _SUPPRESS_HTTP_INSTRUMENTATION_KEY, - _SUPPRESS_INSTRUMENTATION_KEY, -) +from opentelemetry.instrumentation.utils import _SUPPRESS_HTTP_INSTRUMENTATION_KEY from opentelemetry.sdk.resources import SERVICE_NAME, Resource from opentelemetry.sdk.trace import TracerProvider @@ -41,6 +38,22 @@ "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, @@ -207,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) @@ -373,19 +399,18 @@ def suppress_http_instrumentation() -> Iterator[None]: @contextmanager def suppress_db_instrumentation() -> Iterator[None]: - """Suppress OTel auto-instrumentation (incl. SQLAlchemy DB spans) for the wrapped block. + """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. - Trade-off: those queries also drop from the Sentry Queries page. Manually-created - spans (e.g. the gen_ai span) are unaffected — only auto-instrumentors honor this key. + 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 = otel_context.attach( - otel_context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) + token = _suppress_db_spans_var.set(True) try: yield finally: - otel_context.detach(token) + _suppress_db_spans_var.reset(token) def record_db_query_failed( diff --git a/backend/app/tests/core/test_telemetry.py b/backend/app/tests/core/test_telemetry.py new file mode 100644 index 000000000..938475d38 --- /dev/null +++ b/backend/app/tests/core/test_telemetry.py @@ -0,0 +1,185 @@ +"""Tests for the DB-observability helpers in telemetry.py. + +All Sentry and OTel calls are mocked; no real Sentry, OTel provider, or DB +connection is used. Most emitters gate on settings.OTEL_ENABLED, so the +enabled cases patch it True. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from app.core import telemetry + + +def _active_sentry() -> MagicMock: + """A sentry_sdk stand-in whose client reports active.""" + fake = MagicMock() + fake.get_client.return_value.is_active.return_value = True + return fake + + +def _inactive_sentry() -> MagicMock: + fake = MagicMock() + fake.get_client.return_value.is_active.return_value = False + return fake + + +class TestNotableSqlstates: + def test_maps_known_postgres_codes(self): + assert telemetry.NOTABLE_SQLSTATES["40P01"] == "deadlock_detected" + assert telemetry.NOTABLE_SQLSTATES["57014"] == "query_canceled" + assert telemetry.NOTABLE_SQLSTATES["40001"] == "serialization_failure" + + def test_unknown_code_absent(self): + assert "99999" not in telemetry.NOTABLE_SQLSTATES + + +class TestRecordDbQueryFailed: + def test_emits_count_with_operation_and_sqlstate(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_query_failed(operation="SELECT", sqlstate="40P01") + + fake.metrics.count.assert_called_once() + kwargs = fake.metrics.count.call_args.kwargs + assert kwargs["name"] == "db.query.failed" + assert kwargs["value"] == 1 + assert kwargs["attributes"]["db.operation"] == "SELECT" + assert kwargs["attributes"]["db.sqlstate"] == "40P01" + + def test_omits_missing_attributes(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_query_failed() + + assert fake.metrics.count.call_args.kwargs["attributes"] == {} + + def test_noop_when_otel_disabled(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", False), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_query_failed(operation="SELECT", sqlstate="40P01") + + fake.metrics.count.assert_not_called() + + def test_noop_when_sentry_inactive(self): + fake = _inactive_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_query_failed(operation="SELECT", sqlstate="40P01") + + fake.metrics.count.assert_not_called() + + +class TestTagDbError: + def _recording_span(self) -> MagicMock: + span = MagicMock() + span.is_recording.return_value = True + return span + + def test_known_code_sets_named_tag_and_span_attribute(self): + fake = _active_sentry() + span = self._recording_span() + with ( + patch.object(telemetry, "sentry_sdk", fake), + patch.object(telemetry.trace, "get_current_span", return_value=span), + ): + telemetry._tag_db_error("40P01") + + fake.set_tag.assert_any_call("db.system", "postgresql") + fake.set_tag.assert_any_call("db.sqlstate", "40P01") + fake.set_tag.assert_any_call("db.error.name", "deadlock_detected") + span.set_attribute.assert_called_once_with("db.sqlstate", "40P01") + + def test_unknown_code_skips_error_name_tag(self): + fake = _active_sentry() + span = self._recording_span() + with ( + patch.object(telemetry, "sentry_sdk", fake), + patch.object(telemetry.trace, "get_current_span", return_value=span), + ): + telemetry._tag_db_error("99999") + + fake.set_tag.assert_any_call("db.system", "postgresql") + fake.set_tag.assert_any_call("db.sqlstate", "99999") + tag_names = [c.args[0] for c in fake.set_tag.call_args_list] + assert "db.error.name" not in tag_names + + def test_none_sqlstate_is_noop(self): + fake = _active_sentry() + span = self._recording_span() + with ( + patch.object(telemetry, "sentry_sdk", fake), + patch.object(telemetry.trace, "get_current_span", return_value=span), + ): + telemetry._tag_db_error(None) + + fake.set_tag.assert_not_called() + span.set_attribute.assert_not_called() + + def test_swallows_exceptions(self): + fake = MagicMock() + fake.get_client.side_effect = RuntimeError("sentry exploded") + span = self._recording_span() + with ( + patch.object(telemetry, "sentry_sdk", fake), + patch.object(telemetry.trace, "get_current_span", return_value=span), + ): + # Must not raise. + telemetry._tag_db_error("40P01") + + +class TestSuppressDbInstrumentation: + def test_scope_constant(self): + assert telemetry._SQLALCHEMY_SCOPE == "opentelemetry.instrumentation.sqlalchemy" + + def _sqlalchemy_span(self) -> SimpleNamespace: + return SimpleNamespace( + instrumentation_scope=SimpleNamespace( + name="opentelemetry.instrumentation.sqlalchemy" + ) + ) + + def _httpx_span(self) -> SimpleNamespace: + return SimpleNamespace( + instrumentation_scope=SimpleNamespace( + name="opentelemetry.instrumentation.httpx" + ) + ) + + def test_outside_cm_never_drops(self): + assert telemetry._suppress_db_spans_var.get() is False + assert telemetry._should_drop_db_span(self._sqlalchemy_span()) is False + + def test_inside_cm_drops_only_sqlalchemy_spans(self): + with telemetry.suppress_db_instrumentation(): + assert telemetry._suppress_db_spans_var.get() is True + assert telemetry._should_drop_db_span(self._sqlalchemy_span()) is True + assert telemetry._should_drop_db_span(self._httpx_span()) is False + + def test_span_without_scope_not_dropped(self): + with telemetry.suppress_db_instrumentation(): + assert telemetry._should_drop_db_span(SimpleNamespace()) is False + + def test_contextvar_resets_after_exit(self): + with telemetry.suppress_db_instrumentation(): + assert telemetry._suppress_db_spans_var.get() is True + assert telemetry._suppress_db_spans_var.get() is False + + def test_contextvar_resets_even_when_body_raises(self): + try: + with telemetry.suppress_db_instrumentation(): + raise ValueError("boom") + except ValueError: + pass + assert telemetry._suppress_db_spans_var.get() is False From 48f47e4a067cb2f6c9e8fbc5e251f73becc83320 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:27:14 +0530 Subject: [PATCH 3/5] feat(tests): Add unit tests for LLM/response Celery task wrappers and telemetry observability --- .../app/tests/celery/test_job_execution.py | 69 ++++++++++++++ backend/app/tests/core/test_telemetry.py | 94 ++++++++++++++++++- 2 files changed, 161 insertions(+), 2 deletions(-) create mode 100644 backend/app/tests/celery/test_job_execution.py diff --git a/backend/app/tests/celery/test_job_execution.py b/backend/app/tests/celery/test_job_execution.py new file mode 100644 index 000000000..783e23222 --- /dev/null +++ b/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 diff --git a/backend/app/tests/core/test_telemetry.py b/backend/app/tests/core/test_telemetry.py index 938475d38..9eed0e680 100644 --- a/backend/app/tests/core/test_telemetry.py +++ b/backend/app/tests/core/test_telemetry.py @@ -1,13 +1,19 @@ """Tests for the DB-observability helpers in telemetry.py. -All Sentry and OTel calls are mocked; no real Sentry, OTel provider, or DB -connection is used. Most emitters gate on settings.OTEL_ENABLED, so the +Sentry and OTel emission are mocked; no real Sentry connection or OTel provider +is used. The instrument_db_engine tests DO use a real in-memory SQLite engine to +drive the SQLAlchemy event hooks, but stub out the span instrumentor and the +Sentry-backed emit helpers. Most emitters gate on settings.OTEL_ENABLED, so the enabled cases patch it True. """ from types import SimpleNamespace from unittest.mock import MagicMock, patch +import pytest +from sqlalchemy import create_engine, text +from sqlalchemy.pool import QueuePool + from app.core import telemetry @@ -183,3 +189,87 @@ def test_contextvar_resets_even_when_body_raises(self): except ValueError: pass assert telemetry._suppress_db_spans_var.get() is False + + +class TestInstrumentDbEngine: + """Drive the SQLAlchemy event hooks with a real in-memory SQLite engine. + + The span instrumentor is stubbed (no OTel provider needed) and the + Sentry-backed emit helpers are patched so we assert on the hooks alone. + """ + + def _engine(self): + return create_engine("sqlite://", poolclass=QueuePool) + + def test_successful_query_emits_pool_stats(self): + engine = self._engine() + pool_stats = MagicMock() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch( + "opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor" + ), + patch.object(telemetry, "record_db_pool_stats", pool_stats), + ): + telemetry.instrument_db_engine(engine) + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + + pool_stats.assert_called() + kwargs = pool_stats.call_args.kwargs + assert set(kwargs) == {"active", "idle", "total", "overflow"} + + def test_failing_query_fires_error_hook(self): + engine = self._engine() + query_failed = MagicMock() + tag_error = MagicMock() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch( + "opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor" + ), + patch.object(telemetry, "record_db_pool_stats", MagicMock()), + patch.object(telemetry, "record_db_query_failed", query_failed), + patch.object(telemetry, "_tag_db_error", tag_error), + ): + telemetry.instrument_db_engine(engine) + with engine.connect() as conn: + with pytest.raises(Exception): + conn.execute(text("SELECT * FROM does_not_exist")) + + query_failed.assert_called_once() + # SQLite driver errors carry no sqlstate, so the operation is SELECT and + # sqlstate is None — the hook still runs end to end. + assert query_failed.call_args.kwargs["operation"] == "SELECT" + assert query_failed.call_args.kwargs["sqlstate"] is None + tag_error.assert_called_once_with(None) + + def test_second_call_is_idempotent(self): + engine = self._engine() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch( + "opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor" + ) as instrumentor, + patch.object(telemetry, "record_db_pool_stats", MagicMock()), + ): + telemetry.instrument_db_engine(engine) + assert engine._kaapi_db_telemetry_instrumented is True + instrumentor.return_value.instrument.assert_called_once() + + telemetry.instrument_db_engine(engine) + # Guard short-circuits: the instrumentor is not invoked a second time. + instrumentor.return_value.instrument.assert_called_once() + + def test_noop_when_otel_disabled(self): + engine = self._engine() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", False), + patch( + "opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor" + ) as instrumentor, + ): + telemetry.instrument_db_engine(engine) + + instrumentor.assert_not_called() + assert not getattr(engine, "_kaapi_db_telemetry_instrumented", False) From b8ee39d360d7e8476dcadc75f321885b8908f4b8 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:27:28 +0530 Subject: [PATCH 4/5] feat(tests): Simplify patching of SQLAlchemyInstrumentor in DB engine instrumentation tests --- backend/app/tests/core/test_telemetry.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/backend/app/tests/core/test_telemetry.py b/backend/app/tests/core/test_telemetry.py index 9eed0e680..00a6bf36c 100644 --- a/backend/app/tests/core/test_telemetry.py +++ b/backend/app/tests/core/test_telemetry.py @@ -206,9 +206,7 @@ def test_successful_query_emits_pool_stats(self): pool_stats = MagicMock() with ( patch.object(telemetry.settings, "OTEL_ENABLED", True), - patch( - "opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor" - ), + patch("opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor"), patch.object(telemetry, "record_db_pool_stats", pool_stats), ): telemetry.instrument_db_engine(engine) @@ -225,9 +223,7 @@ def test_failing_query_fires_error_hook(self): tag_error = MagicMock() with ( patch.object(telemetry.settings, "OTEL_ENABLED", True), - patch( - "opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor" - ), + patch("opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor"), patch.object(telemetry, "record_db_pool_stats", MagicMock()), patch.object(telemetry, "record_db_query_failed", query_failed), patch.object(telemetry, "_tag_db_error", tag_error), From c0d1171e813774f54b9357f80264965fc00376fd Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Mon, 24 Aug 2026 01:19:55 +0530 Subject: [PATCH 5/5] feat(telemetry): Enhance database observability with slow query and connection metrics --- backend/app/core/config.py | 2 + backend/app/core/middleware.py | 80 ++++--- backend/app/core/telemetry.py | 167 ++++++++++++-- backend/app/tests/core/test_middleware.py | 124 +++++++++++ backend/app/tests/core/test_telemetry.py | 258 +++++++++++++++++++++- 5 files changed, 589 insertions(+), 42 deletions(-) create mode 100644 backend/app/tests/core/test_middleware.py diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 4d1544da3..9609db31a 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -153,6 +153,8 @@ def AWS_S3_BUCKET(self) -> str: OTEL_SERVICE_NAME: str = "kaapi-backend" BACKEND_SERVICE_NAME: str = "kaapi-backend" CRON_SERVICE_NAME: str = "kaapi-cron" + # Queries at/above this duration increment the db.query.slow counter. + DB_SLOW_QUERY_MS: int = 500 # Threshold Request Rate per minute THRESHOLD_LLM_CALL_RATE: int = 15 diff --git a/backend/app/core/middleware.py b/backend/app/core/middleware.py index bc8664259..8081e2018 100644 --- a/backend/app/core/middleware.py +++ b/backend/app/core/middleware.py @@ -19,6 +19,11 @@ } ) +CRON_PATH_PREFIX: str = f"{settings.API_V1_STR}/cron/" + +# Excluded from traces only; logs/metrics and spans inside the handler still emit. +TRACE_EXCLUDED_PATH_PREFIXES: frozenset[str] = frozenset({CRON_PATH_PREFIX}) + class StripTrailingSlashMiddleware: """ @@ -50,8 +55,33 @@ def _resolve_http_route(request: Request) -> str: return templated or "unmatched" +def _emit_http_metrics( + *, method: str, http_route: str, status: int, duration_ms: float +) -> None: + """Emit HTTP traffic/latency/error counters to Sentry. No-op if the SDK is inactive.""" + try: + if not sentry_sdk.get_client().is_active(): + return + attrs = { + "http.method": method, + "http.route": http_route, + "http.status_code": str(status), + } + sentry_sdk.metrics.count("http.server.request.count", 1, attributes=attrs) + sentry_sdk.metrics.distribution( + "http.server.request.duration", + duration_ms, + unit="millisecond", + attributes=attrs, + ) + if status >= 400: + sentry_sdk.metrics.count("http.server.request.error", 1, attributes=attrs) + except Exception: + logger.debug("[_emit_http_metrics] Sentry metric emit failed") + + async def http_request_logger(request: Request, call_next) -> Response: - if request.url.path.startswith(f"{settings.API_V1_STR}/cron/"): + if request.url.path.startswith(CRON_PATH_PREFIX): with log_service_name(settings.CRON_SERVICE_NAME): return await _log_http_request(request, call_next) @@ -62,6 +92,8 @@ async def _log_http_request(request: Request, call_next) -> Response: start_time = time.time() method = request.method raw_path = request.url.path + # Health/utility paths excluded so they don't skew platform traffic metrics. + metrics_enabled = raw_path not in SILENT_LOG_PATHS span = trace.get_current_span() if span.is_recording(): @@ -78,6 +110,7 @@ async def _log_http_request(request: Request, call_next) -> Response: try: response = await call_next(request) except Exception: + duration_ms = (time.time() - start_time) * 1000 status = 500 http_route = _resolve_http_route(request) if span.is_recording(): @@ -88,6 +121,13 @@ async def _log_http_request(request: Request, call_next) -> Response: sentry_sdk.set_tag("http.route", http_route) sentry_sdk.set_tag("http.status_code", str(status)) sentry_sdk.set_tag("http.response.status_code", str(status)) + if metrics_enabled: + _emit_http_metrics( + method=method, + http_route=http_route, + status=status, + duration_ms=duration_ms, + ) logger.exception("Unhandled exception during request") raise @@ -101,32 +141,18 @@ async def _log_http_request(request: Request, call_next) -> Response: span.set_attribute("http.response.status_code", status) span.set_attribute("http.request.duration_ms", round(duration_ms, 2)) - if raw_path not in SILENT_LOG_PATHS: - logger.info(f"{method} {raw_path} - {status} [{duration_ms:.2f}ms]") - - try: - if sentry_sdk.get_client().is_active(): - sentry_sdk.set_tag("http.route", http_route) - sentry_sdk.set_tag("http.status_code", str(status)) - sentry_sdk.set_tag("http.response.status_code", str(status)) + if sentry_sdk.get_client().is_active(): + sentry_sdk.set_tag("http.route", http_route) + sentry_sdk.set_tag("http.status_code", str(status)) + sentry_sdk.set_tag("http.response.status_code", str(status)) - attrs = { - "http.method": method, - "http.route": http_route, - "http.status_code": str(status), - } - sentry_sdk.metrics.count("http.server.request.count", 1, attributes=attrs) - sentry_sdk.metrics.distribution( - "http.server.request.duration", - duration_ms, - unit="millisecond", - attributes=attrs, - ) - if status >= 400: - sentry_sdk.metrics.count( - "http.server.request.error", 1, attributes=attrs - ) - except Exception: - logger.debug("[http_request_logger] Sentry metric emit failed") + if metrics_enabled: + logger.info(f"{method} {raw_path} - {status} [{duration_ms:.2f}ms]") + _emit_http_metrics( + method=method, + http_route=http_route, + status=status, + duration_ms=duration_ms, + ) return response diff --git a/backend/app/core/telemetry.py b/backend/app/core/telemetry.py index 2a1cd3afb..1cb1ee246 100644 --- a/backend/app/core/telemetry.py +++ b/backend/app/core/telemetry.py @@ -1,5 +1,7 @@ import json import logging +import re +import time from collections.abc import Iterator from contextlib import contextmanager from contextvars import ContextVar @@ -15,6 +17,7 @@ from opentelemetry.instrumentation.utils import _SUPPRESS_HTTP_INSTRUMENTATION_KEY from opentelemetry.sdk.resources import SERVICE_NAME, Resource from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.trace import SpanKind, StatusCode, format_span_id from app.core.config import settings @@ -34,6 +37,17 @@ "53300": "too_many_connections", } +_DB_CONNECTION_EVENT_METRICS: dict[str, str] = { + "opened": "db.connection.opened", + "closed": "db.connection.closed", + "invalidated": "db.connection.invalidated", +} + +_DB_TRANSACTION_METRICS: dict[str, str] = { + "commit": "db.transaction.commit", + "rollback": "db.transaction.rollback", +} + _log_context_var: ContextVar[dict[str, str] | None] = ContextVar( "kaapi_log_context", default=None ) @@ -55,6 +69,22 @@ def _should_drop_db_span(otel_span: object) -> bool: return scope is not None and getattr(scope, "name", None) == _SQLALCHEMY_SCOPE +def _should_drop_bare_http_trace( + *, is_root: bool, kind: SpanKind, status_code: StatusCode, had_children: bool +) -> bool: + """True for a root HTTP server span whose trace has no child spans and no error. + + These single-span transactions carry no work worth a trace; errors are kept so + failures stay visible (5xx also surface in http.server.request.error metrics). + """ + return ( + is_root + and not had_children + and kind == SpanKind.SERVER + and status_code != StatusCode.ERROR + ) + + def _emit_sentry_metric( metric_type: str, name: str, @@ -220,20 +250,47 @@ def setup_telemetry(service_name: str | None = None) -> None: if settings.SENTRY_DSN: from sentry_sdk.integrations.opentelemetry import SentrySpanProcessor - class _DbSpanFilteringProcessor(SentrySpanProcessor): - """Drop SQLAlchemy DB spans from Sentry while suppress_db_instrumentation() is active. + class _NoiseFilteringSpanProcessor(SentrySpanProcessor): + """Keep Sentry traces meaningful: drop suppressed DB spans and empty HTTP traces. - 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. + - DB: while suppress_db_instrumentation() is active, SQLAlchemy spans are + skipped by scope (HTTP/Requests keep flowing). + - HTTP: a root server span whose trace gathered no child spans (and no error) + is dropped at on_end — it never reaches Sentry as a bare single-span trace. """ + def __init__(self) -> None: + super().__init__() + self._traces_with_children: set[int] = set() + def on_start(self, otel_span, parent_context=None): # type: ignore[override] if _should_drop_db_span(otel_span): return + parent = otel_span.parent + if parent is not None and not parent.is_remote: + self._traces_with_children.add( + otel_span.get_span_context().trace_id + ) super().on_start(otel_span, parent_context) - tracer_provider.add_span_processor(_DbSpanFilteringProcessor()) + def on_end(self, otel_span) -> None: # type: ignore[override] + span_context = otel_span.get_span_context() + parent = otel_span.parent + is_root = parent is None or parent.is_remote + had_children = span_context.trace_id in self._traces_with_children + if is_root: + self._traces_with_children.discard(span_context.trace_id) + if _should_drop_bare_http_trace( + is_root=is_root, + kind=otel_span.kind, + status_code=otel_span.status.status_code, + had_children=had_children, + ): + self.otel_span_map.pop(format_span_id(span_context.span_id), None) + return + super().on_end(otel_span) + + tracer_provider.add_span_processor(_NoiseFilteringSpanProcessor()) trace.set_tracer_provider(tracer_provider) @@ -402,7 +459,7 @@ 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 + Only DB spans are filtered (by _NoiseFilteringSpanProcessor via instrumentation scope) — HTTP/Requests instrumentation stays active. Trade-off: the wrapped DB queries also drop from the Sentry Queries page. """ @@ -431,6 +488,40 @@ def record_db_query_failed( _emit_sentry_metric("count", "db.query.failed", 1, attributes=attrs) +def record_db_slow_query(operation: str | None = None) -> None: + """Emit a slow-query counter to Sentry (queries at/above settings.DB_SLOW_QUERY_MS).""" + if not settings.OTEL_ENABLED: + return + + attrs: dict[str, str | int | float] = {} + if operation: + attrs["db.operation"] = operation + + _emit_sentry_metric("count", "db.query.slow", 1, attributes=attrs) + + +def record_db_connection_event(event: str) -> None: + """Emit a connection-lifecycle counter (opened/closed/invalidated) to Sentry.""" + if not settings.OTEL_ENABLED: + return + + metric = _DB_CONNECTION_EVENT_METRICS.get(event) + if metric is None: + return + _emit_sentry_metric("count", metric, 1) + + +def record_db_transaction(outcome: str) -> None: + """Emit a transaction-outcome counter (commit/rollback) to Sentry for ratio tracking.""" + if not settings.OTEL_ENABLED: + return + + metric = _DB_TRANSACTION_METRICS.get(outcome) + if metric is None: + return + _emit_sentry_metric("count", metric, 1) + + def _tag_db_error(sqlstate: str | None) -> None: """Tag the active Sentry scope + span with a DB error's Postgres SQLSTATE.""" if not sqlstate: @@ -573,16 +664,33 @@ def instrument_app(app: object) -> None: """Instrument the FastAPI app. Call after the app is created.""" if not settings.OTEL_ENABLED: return - # Health checks are high-volume/noise and should not generate traces. + from app.core.middleware import SILENT_LOG_PATHS, TRACE_EXCLUDED_PATH_PREFIXES + + # Non-business request spans are noise: health/utility, FastAPI's own doc/schema + # endpoints (read off the app so they track config), and cron polling (prefix). + exact_paths = set(SILENT_LOG_PATHS) + for attr in ( + "docs_url", + "redoc_url", + "openapi_url", + "swagger_ui_oauth2_redirect_url", + ): + path = getattr(app, attr, None) + if path: + exact_paths.add(path) + + patterns = [rf"^{re.escape(p)}/?$" for p in exact_paths] + patterns += [rf"^{re.escape(p)}" for p in TRACE_EXCLUDED_PATH_PREFIXES] + excluded_urls = ",".join(sorted(patterns)) FastAPIInstrumentor.instrument_app( # type: ignore[arg-type] app, - excluded_urls=r"^/health/?$", + excluded_urls=excluded_urls, ) logger.debug("[instrument_app] FastAPI instrumented with OpenTelemetry") def instrument_db_engine(engine: object) -> None: - """Instrument SQLAlchemy engine with DB query spans + pool gauges.""" + """Instrument a SQLAlchemy engine: query spans, slow-query/pool/connection/transaction metrics.""" if not settings.OTEL_ENABLED: return if getattr(engine, "_kaapi_db_telemetry_instrumented", False): @@ -600,7 +708,7 @@ def instrument_db_engine(engine: object) -> None: ) try: - from sqlalchemy import event + from sqlalchemy import ExceptionContext, event except Exception: logger.exception("[instrument_db_engine] Failed importing SQLAlchemy events") return @@ -627,7 +735,8 @@ def _before_cursor_execute( conn, cursor, statement, parameters, context, executemany ) -> None: del cursor, parameters, executemany - # Operation kept for error attribution; per-query timing now lives on the span. + # Timing only flags slow queries; per-query duration lives on the span. + context._kaapi_db_started_at = time.perf_counter() context._kaapi_db_operation = ( str(statement).split(None, 1)[0].upper() if statement else "UNKNOWN" ) @@ -637,11 +746,16 @@ def _before_cursor_execute( def _after_cursor_execute( conn, cursor, statement, parameters, context, executemany ) -> None: - del cursor, statement, parameters, context, executemany + del cursor, statement, parameters, executemany + started_at = getattr(context, "_kaapi_db_started_at", None) + if started_at is not None: + duration_ms = (time.perf_counter() - started_at) * 1000 + if duration_ms >= settings.DB_SLOW_QUERY_MS: + record_db_slow_query(getattr(context, "_kaapi_db_operation", None)) _emit_pool_metrics(conn.engine.pool) @event.listens_for(engine, "handle_error") - def _handle_error(exception_context) -> None: + def _handle_error(exception_context: ExceptionContext) -> None: context = exception_context.execution_context operation = ( getattr(context, "_kaapi_db_operation", None) @@ -662,5 +776,30 @@ def _on_checkin(dbapi_connection, connection_record) -> None: del dbapi_connection, connection_record _emit_pool_metrics(engine.pool) + @event.listens_for(engine.pool, "connect") + def _on_connect(dbapi_connection, connection_record) -> None: + del dbapi_connection, connection_record + record_db_connection_event("opened") + + @event.listens_for(engine.pool, "close") + def _on_close(dbapi_connection, connection_record) -> None: + del dbapi_connection, connection_record + record_db_connection_event("closed") + + @event.listens_for(engine.pool, "invalidate") + def _on_invalidate(dbapi_connection, connection_record, exception) -> None: + del dbapi_connection, connection_record, exception + record_db_connection_event("invalidated") + + @event.listens_for(engine, "commit") + def _on_commit(conn) -> None: + del conn + record_db_transaction("commit") + + @event.listens_for(engine, "rollback") + def _on_rollback(conn) -> None: + del conn + record_db_transaction("rollback") + engine._kaapi_db_telemetry_instrumented = True logger.debug("[instrument_db_engine] SQLAlchemy DB telemetry enabled") diff --git a/backend/app/tests/core/test_middleware.py b/backend/app/tests/core/test_middleware.py new file mode 100644 index 000000000..2987737a9 --- /dev/null +++ b/backend/app/tests/core/test_middleware.py @@ -0,0 +1,124 @@ +"""Tests for HTTP request metrics in core/middleware.py. + +sentry_sdk and the OTel span are mocked; no real Sentry connection is used. The +middleware is driven directly with a fake Request + call_next so traffic, latency +and error counters can be asserted without an ASGI server. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from app.core import middleware + + +def _active_sentry() -> MagicMock: + fake = MagicMock() + fake.get_client.return_value.is_active.return_value = True + return fake + + +def _request(path: str, method: str = "GET", route: str | None = None) -> MagicMock: + req = MagicMock() + req.url.path = path + req.method = method + req.scope = {"route": SimpleNamespace(path=route)} if route else {} + return req + + +def _metric_names(fake: MagicMock) -> list[str]: + calls = list(fake.metrics.count.call_args_list) + list( + fake.metrics.distribution.call_args_list + ) + return [c.args[0] for c in calls] + + +@pytest.fixture +def non_recording_span(): + span = MagicMock() + span.is_recording.return_value = False + with patch.object(middleware.trace, "get_current_span", return_value=span): + yield span + + +class TestHttpRequestMetrics: + async def _run(self, request, call_next, fake): + with patch.object(middleware, "sentry_sdk", fake): + return await middleware._log_http_request(request, call_next) + + @pytest.mark.asyncio + async def test_success_emits_traffic_and_latency(self, non_recording_span): + fake = _active_sentry() + request = _request("/api/v1/items", route="/api/v1/items") + call_next = AsyncMock(return_value=SimpleNamespace(status_code=200)) + + await self._run(request, call_next, fake) + + names = _metric_names(fake) + assert "http.server.request.count" in names + assert "http.server.request.duration" in names + assert "http.server.request.error" not in names + + @pytest.mark.asyncio + async def test_client_error_status_emits_error_metric(self, non_recording_span): + fake = _active_sentry() + request = _request("/api/v1/items", route="/api/v1/items") + call_next = AsyncMock(return_value=SimpleNamespace(status_code=404)) + + await self._run(request, call_next, fake) + + error_calls = [ + c + for c in fake.metrics.count.call_args_list + if c.args[0] == "http.server.request.error" + ] + assert len(error_calls) == 1 + assert error_calls[0].kwargs["attributes"]["http.status_code"] == "404" + + @pytest.mark.asyncio + async def test_unhandled_exception_counts_500_and_reraises( + self, non_recording_span + ): + fake = _active_sentry() + request = _request("/api/v1/items", route="/api/v1/items") + call_next = AsyncMock(side_effect=ValueError("boom")) + + with pytest.raises(ValueError): + await self._run(request, call_next, fake) + + error_calls = [ + c + for c in fake.metrics.count.call_args_list + if c.args[0] == "http.server.request.error" + ] + assert len(error_calls) == 1 + assert error_calls[0].kwargs["attributes"]["http.status_code"] == "500" + + @pytest.mark.asyncio + async def test_health_path_excluded_from_metrics(self, non_recording_span): + fake = _active_sentry() + request = _request("/health", route="/health") + call_next = AsyncMock(return_value=SimpleNamespace(status_code=200)) + + await self._run(request, call_next, fake) + + fake.metrics.count.assert_not_called() + fake.metrics.distribution.assert_not_called() + + @pytest.mark.asyncio + async def test_route_template_used_as_metric_tag(self, non_recording_span): + fake = _active_sentry() + request = _request("/api/v1/items/42", route="/api/v1/items/{item_id}") + call_next = AsyncMock(return_value=SimpleNamespace(status_code=200)) + + await self._run(request, call_next, fake) + + count_call = next( + c + for c in fake.metrics.count.call_args_list + if c.args[0] == "http.server.request.count" + ) + assert ( + count_call.kwargs["attributes"]["http.route"] == "/api/v1/items/{item_id}" + ) diff --git a/backend/app/tests/core/test_telemetry.py b/backend/app/tests/core/test_telemetry.py index 00a6bf36c..4fcb9962c 100644 --- a/backend/app/tests/core/test_telemetry.py +++ b/backend/app/tests/core/test_telemetry.py @@ -11,7 +11,11 @@ from unittest.mock import MagicMock, patch import pytest +from fastapi import FastAPI +from opentelemetry.trace import SpanKind, StatusCode +from opentelemetry.util.http import ExcludeList from sqlalchemy import create_engine, text +from sqlalchemy.exc import OperationalError from sqlalchemy.pool import QueuePool from app.core import telemetry @@ -230,7 +234,7 @@ def test_failing_query_fires_error_hook(self): ): telemetry.instrument_db_engine(engine) with engine.connect() as conn: - with pytest.raises(Exception): + with pytest.raises(OperationalError): conn.execute(text("SELECT * FROM does_not_exist")) query_failed.assert_called_once() @@ -269,3 +273,255 @@ def test_noop_when_otel_disabled(self): instrumentor.assert_not_called() assert not getattr(engine, "_kaapi_db_telemetry_instrumented", False) + + +class TestShouldDropBareHttpTrace: + def test_drops_root_http_span_with_no_children(self): + assert telemetry._should_drop_bare_http_trace( + is_root=True, + kind=SpanKind.SERVER, + status_code=StatusCode.UNSET, + had_children=False, + ) + + def test_keeps_when_trace_has_children(self): + assert not telemetry._should_drop_bare_http_trace( + is_root=True, + kind=SpanKind.SERVER, + status_code=StatusCode.UNSET, + had_children=True, + ) + + def test_keeps_error_trace_even_without_children(self): + assert not telemetry._should_drop_bare_http_trace( + is_root=True, + kind=SpanKind.SERVER, + status_code=StatusCode.ERROR, + had_children=False, + ) + + def test_keeps_non_root_span(self): + assert not telemetry._should_drop_bare_http_trace( + is_root=False, + kind=SpanKind.SERVER, + status_code=StatusCode.UNSET, + had_children=False, + ) + + def test_keeps_non_server_root(self): + assert not telemetry._should_drop_bare_http_trace( + is_root=True, + kind=SpanKind.INTERNAL, + status_code=StatusCode.UNSET, + had_children=False, + ) + + +class TestInstrumentApp: + def _app(self) -> FastAPI: + return FastAPI(openapi_url="/api/v1/openapi.json") + + def _excluded(self) -> ExcludeList: + instrument = MagicMock() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry.FastAPIInstrumentor, "instrument_app", instrument), + ): + telemetry.instrument_app(self._app()) + + excluded_urls = instrument.call_args.kwargs["excluded_urls"] + return ExcludeList(excluded_urls.split(",")) + + def test_health_and_cron_paths_excluded_from_traces(self): + el = self._excluded() + assert el.url_disabled("/health") + assert el.url_disabled("/api/v1/utils/health") + assert el.url_disabled("/api/v1/cron/run_batches") + + def test_framework_doc_paths_excluded_from_traces(self): + el = self._excluded() + assert el.url_disabled("/docs") + assert el.url_disabled("/redoc") + assert el.url_disabled("/api/v1/openapi.json") + assert el.url_disabled("/docs/oauth2-redirect") + + def test_real_endpoints_still_traced(self): + el = self._excluded() + assert not el.url_disabled("/api/v1/llm/generate") + assert not el.url_disabled("/api/v1/cronies") + assert not el.url_disabled("/api/v1/documents") + + def test_noop_when_otel_disabled(self): + instrument = MagicMock() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", False), + patch.object(telemetry.FastAPIInstrumentor, "instrument_app", instrument), + ): + telemetry.instrument_app(self._app()) + + instrument.assert_not_called() + + +class TestRecordDbSlowQuery: + def test_emits_count_with_operation(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_slow_query("SELECT") + + fake.metrics.count.assert_called_once() + kwargs = fake.metrics.count.call_args.kwargs + assert kwargs["name"] == "db.query.slow" + assert kwargs["value"] == 1 + assert kwargs["attributes"]["db.operation"] == "SELECT" + + def test_noop_when_otel_disabled(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", False), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_slow_query("SELECT") + + fake.metrics.count.assert_not_called() + + +class TestRecordDbConnectionEvent: + def test_known_event_emits_named_metric(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_connection_event("invalidated") + + assert ( + fake.metrics.count.call_args.kwargs["name"] == "db.connection.invalidated" + ) + + def test_unknown_event_is_noop(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_connection_event("teleported") + + fake.metrics.count.assert_not_called() + + def test_noop_when_otel_disabled(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", False), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_connection_event("opened") + + fake.metrics.count.assert_not_called() + + +class TestRecordDbTransaction: + @pytest.mark.parametrize( + ("outcome", "metric"), + [ + ("commit", "db.transaction.commit"), + ("rollback", "db.transaction.rollback"), + ], + ) + def test_known_outcome_emits_named_metric(self, outcome, metric): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_transaction(outcome) + + assert fake.metrics.count.call_args.kwargs["name"] == metric + + def test_unknown_outcome_is_noop(self): + fake = _active_sentry() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry, "sentry_sdk", fake), + ): + telemetry.record_db_transaction("savepoint") + + fake.metrics.count.assert_not_called() + + +class TestInstrumentDbEngineMetrics: + """Drive the new slow-query/connection/transaction hooks with a real SQLite engine.""" + + def _engine(self): + return create_engine("sqlite://", poolclass=QueuePool) + + def test_slow_query_counted_when_over_threshold(self): + engine = self._engine() + slow = MagicMock() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry.settings, "DB_SLOW_QUERY_MS", 0), + patch("opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor"), + patch.object(telemetry, "record_db_pool_stats", MagicMock()), + patch.object(telemetry, "record_db_slow_query", slow), + ): + telemetry.instrument_db_engine(engine) + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + + slow.assert_called() + assert "SELECT" in [c.args[0] for c in slow.call_args_list] + + def test_fast_query_not_counted_when_under_threshold(self): + engine = self._engine() + slow = MagicMock() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch.object(telemetry.settings, "DB_SLOW_QUERY_MS", 10_000), + patch("opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor"), + patch.object(telemetry, "record_db_pool_stats", MagicMock()), + patch.object(telemetry, "record_db_slow_query", slow), + ): + telemetry.instrument_db_engine(engine) + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + + slow.assert_not_called() + + def test_connection_open_event_recorded(self): + engine = self._engine() + conn_event = MagicMock() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch("opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor"), + patch.object(telemetry, "record_db_pool_stats", MagicMock()), + patch.object(telemetry, "record_db_connection_event", conn_event), + ): + telemetry.instrument_db_engine(engine) + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + + assert "opened" in [c.args[0] for c in conn_event.call_args_list] + + def test_commit_and_rollback_recorded(self): + engine = self._engine() + txn = MagicMock() + with ( + patch.object(telemetry.settings, "OTEL_ENABLED", True), + patch("opentelemetry.instrumentation.sqlalchemy.SQLAlchemyInstrumentor"), + patch.object(telemetry, "record_db_pool_stats", MagicMock()), + patch.object(telemetry, "record_db_transaction", txn), + ): + telemetry.instrument_db_engine(engine) + with engine.connect() as conn: + conn.execute(text("CREATE TABLE t (x)")) + conn.commit() + with engine.connect() as conn: + conn.execute(text("INSERT INTO t VALUES (1)")) + conn.rollback() + + outcomes = [c.args[0] for c in txn.call_args_list] + assert "commit" in outcomes + assert "rollback" in outcomes