From 0e6a2ec37853e4d7e6ad66d4d72c66e3945f6f66 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 20 Mar 2026 21:11:16 -0700 Subject: [PATCH] feat(prometheus.py): emit db event metrics more accurate logging --- docs/my-website/docs/proxy/prometheus.md | 15 + litellm/integrations/prometheus_services.py | 70 +++- litellm/proxy/db/prisma_client.py | 202 +++++++++++- litellm/proxy/utils.py | 23 +- litellm/types/services.py | 13 + .../integrations/test_prometheus_services.py | 149 +++++++++ .../proxy/db/test_prisma_db_metrics.py | 298 ++++++++++++++++++ 7 files changed, 731 insertions(+), 39 deletions(-) create mode 100644 tests/test_litellm/proxy/db/test_prisma_db_metrics.py diff --git a/docs/my-website/docs/proxy/prometheus.md b/docs/my-website/docs/proxy/prometheus.md index d8f0d83b59d..30312d85404 100644 --- a/docs/my-website/docs/proxy/prometheus.md +++ b/docs/my-website/docs/proxy/prometheus.md @@ -549,6 +549,21 @@ litellm_settings: | `litellm_redis_fails` | Number of failed redis calls | | `litellm_self_latency` | Histogram latency for successful litellm api call | +#### DB Read/Write Metrics + +These metrics are instrumented at the Prisma query level, so they only fire when an actual database call is made (cache hits are not counted). + +All DB read/write metrics include a `call_type` label showing the table and operation (e.g. `litellm_usertable.find_unique`, `litellm_teamtable.find_many`, `query_raw`, `execute_raw`), so you can identify exactly which queries are driving DB load. + +| Metric Name | Description | +|----------------------------------|------------------------------------------------------| +| `litellm_db_read_latency` | Histogram latency for actual DB read operations, with `call_type` label (e.g. `find_unique`, `find_many`, `query_raw`) | +| `litellm_db_read_total_requests` | Total actual DB read operations, with `call_type` label | +| `litellm_db_read_failed_requests`| Failed DB read operations (with `call_type`, `error_class`, and `function_name` labels) | +| `litellm_db_write_latency` | Histogram latency for actual DB write operations, with `call_type` label (e.g. `create`, `update`, `delete`, `execute_raw`) | +| `litellm_db_write_total_requests`| Total actual DB write operations, with `call_type` label | +| `litellm_db_write_failed_requests`| Failed DB write operations (with `call_type`, `error_class`, and `function_name` labels) | + #### DB Transaction Queue Health Metrics Use these metrics to monitor the health of the DB Transaction Queue. Eg. Monitoring the size of the in-memory and redis buffers. diff --git a/litellm/integrations/prometheus_services.py b/litellm/integrations/prometheus_services.py index 55ce758ece6..803795f05fd 100644 --- a/litellm/integrations/prometheus_services.py +++ b/litellm/integrations/prometheus_services.py @@ -3,7 +3,7 @@ # On success + failure, log events to Prometheus for litellm / adjacent services (litellm, redis, postgres, llm api providers) -from typing import Dict, List, Optional, Union +from typing import Dict, FrozenSet, List, Optional, Union from litellm._logging import print_verbose, verbose_logger from litellm.types.integrations.prometheus import LATENCY_BUCKETS @@ -16,6 +16,10 @@ from litellm.types.services import ( FAILED_REQUESTS_LABELS = ["error_class", "function_name"] +_SERVICES_WITH_CALL_TYPE_LABEL: FrozenSet[str] = frozenset( + {ServiceTypes.DB_READ.value, ServiceTypes.DB_WRITE.value} +) + class PrometheusServicesLogger: # Class variables or attributes @@ -126,10 +130,13 @@ class PrometheusServicesLogger: is_registered = self.is_metric_registered(metric_name) if is_registered: return self._get_metric(metric_name) + labelnames = [service] + if service in _SERVICES_WITH_CALL_TYPE_LABEL: + labelnames.append("call_type") return self.Histogram( metric_name, "Latency for {} service".format(service), - labelnames=[service], + labelnames=labelnames, buckets=LATENCY_BUCKETS, ) @@ -152,10 +159,13 @@ class PrometheusServicesLogger: is_registered = self.is_metric_registered(metric_name) if is_registered: return self._get_metric(metric_name) + labelnames = [service] + if service in _SERVICES_WITH_CALL_TYPE_LABEL: + labelnames.append("call_type") return self.Counter( metric_name, "Total {} for {} service".format(type_of_request, service), - labelnames=[service] + (additional_labels or []), + labelnames=labelnames + (additional_labels or []), ) def observe_histogram( @@ -163,10 +173,14 @@ class PrometheusServicesLogger: histogram, labels: str, amount: float, + call_type: Optional[str] = None, ): assert isinstance(histogram, self.Histogram) - histogram.labels(labels).observe(amount) + if call_type is not None: + histogram.labels(labels, call_type).observe(amount) + else: + histogram.labels(labels).observe(amount) def update_gauge( self, @@ -183,19 +197,27 @@ class PrometheusServicesLogger: labels: str, amount: float, additional_labels: Optional[List[str]] = [], + call_type: Optional[str] = None, ): assert isinstance(counter, self.Counter) + label_values = [labels] + if call_type is not None: + label_values.append(call_type) if additional_labels: - counter.labels(labels, *additional_labels).inc(amount) - else: - counter.labels(labels).inc(amount) + label_values.extend(additional_labels) + counter.labels(*label_values).inc(amount) def service_success_hook(self, payload: ServiceLoggerPayload): if self.mock_testing: self.mock_testing_success_calls += 1 if payload.service.value in self.payload_to_prometheus_map: + _call_type = ( + payload.call_type + if payload.service.value in _SERVICES_WITH_CALL_TYPE_LABEL + else None + ) prom_objects = self.payload_to_prometheus_map[payload.service.value] for obj in prom_objects: if isinstance(obj, self.Histogram): @@ -203,12 +225,14 @@ class PrometheusServicesLogger: histogram=obj, labels=payload.service.value, amount=payload.duration, + call_type=_call_type, ) elif isinstance(obj, self.Counter) and "total_requests" in obj._name: self.increment_counter( counter=obj, labels=payload.service.value, - amount=1, # LOG TOTAL REQUESTS TO PROMETHEUS + amount=1, + call_type=_call_type, ) def service_failure_hook(self, payload: ServiceLoggerPayload): @@ -216,13 +240,19 @@ class PrometheusServicesLogger: self.mock_testing_failure_calls += 1 if payload.service.value in self.payload_to_prometheus_map: + _call_type = ( + payload.call_type + if payload.service.value in _SERVICES_WITH_CALL_TYPE_LABEL + else None + ) prom_objects = self.payload_to_prometheus_map[payload.service.value] for obj in prom_objects: if isinstance(obj, self.Counter): self.increment_counter( counter=obj, labels=payload.service.value, - amount=1, # LOG ERROR COUNT / TOTAL REQUESTS TO PROMETHEUS + amount=1, + call_type=_call_type, ) async def async_service_success_hook(self, payload: ServiceLoggerPayload): @@ -233,6 +263,11 @@ class PrometheusServicesLogger: self.mock_testing_success_calls += 1 if payload.service.value in self.payload_to_prometheus_map: + _call_type = ( + payload.call_type + if payload.service.value in _SERVICES_WITH_CALL_TYPE_LABEL + else None + ) prom_objects = self.payload_to_prometheus_map[payload.service.value] for obj in prom_objects: if isinstance(obj, self.Histogram): @@ -240,12 +275,14 @@ class PrometheusServicesLogger: histogram=obj, labels=payload.service.value, amount=payload.duration, + call_type=_call_type, ) elif isinstance(obj, self.Counter) and "total_requests" in obj._name: self.increment_counter( counter=obj, labels=payload.service.value, - amount=1, # LOG TOTAL REQUESTS TO PROMETHEUS + amount=1, + call_type=_call_type, ) elif isinstance(obj, self.Gauge): if payload.event_metadata: @@ -266,21 +303,26 @@ class PrometheusServicesLogger: function_name = payload.call_type if payload.service.value in self.payload_to_prometheus_map: + _call_type = ( + payload.call_type + if payload.service.value in _SERVICES_WITH_CALL_TYPE_LABEL + else None + ) prom_objects = self.payload_to_prometheus_map[payload.service.value] for obj in prom_objects: - # increment both failed and total requests if isinstance(obj, self.Counter): if "failed_requests" in obj._name: self.increment_counter( counter=obj, labels=payload.service.value, - # log additional_labels=["error_class", "function_name"], used for debugging what's going wrong with the DB additional_labels=[error_class, function_name], - amount=1, # LOG ERROR COUNT TO PROMETHEUS + amount=1, + call_type=_call_type, ) else: self.increment_counter( counter=obj, labels=payload.service.value, - amount=1, # LOG TOTAL REQUESTS TO PROMETHEUS + amount=1, + call_type=_call_type, ) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index c9c0cfe8f68..9c945a5f4a1 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -1,8 +1,10 @@ """ This file contains the PrismaWrapper class, which is used to wrap the Prisma client and handle the RDS IAM token. +It also contains PrismaModelProxy for instrumenting actual DB operations with metrics. """ import asyncio +import functools import os import random import subprocess @@ -10,10 +12,162 @@ import time import urllib import urllib.parse from datetime import datetime, timedelta -from typing import Any, Optional, Union +from typing import TYPE_CHECKING, Any, Callable, FrozenSet, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.secret_managers.main import str_to_bool +from litellm.types.services import ServiceTypes + +if TYPE_CHECKING: + from litellm._service_logger import ServiceLogging + +_DB_READ_METHODS: FrozenSet[str] = frozenset( + { + "find_unique", + "find_many", + "find_first", + "find_first_or_raise", + "find_unique_or_raise", + "count", + "group_by", + } +) + +_DB_WRITE_METHODS: FrozenSet[str] = frozenset( + { + "create", + "create_many", + "update", + "update_many", + "upsert", + "delete", + "delete_many", + } +) + +_INSTRUMENTED_METHODS: FrozenSet[str] = _DB_READ_METHODS | _DB_WRITE_METHODS + + +class PrismaModelProxy: + """ + Proxy around a Prisma model (e.g. ``prisma.litellm_usertable``) that + instruments every CRUD call with latency / count / error metrics via + :class:`ServiceLogging`. + + Only async methods are wrapped because all Prisma Python CRUD methods + are async. Non-CRUD attribute access is forwarded unchanged. + """ + + __slots__ = ("_model", "_table_name", "_service_logger_obj") + + def __init__( + self, + model: Any, + table_name: str, + service_logger_obj: "ServiceLogging", + ): + self._model = model + self._table_name = table_name + self._service_logger_obj = service_logger_obj + + def __getattr__(self, name: str) -> Any: + attr = getattr(self._model, name) + + if name not in _INSTRUMENTED_METHODS: + return attr + + service_type = ( + ServiceTypes.DB_READ if name in _DB_READ_METHODS else ServiceTypes.DB_WRITE + ) + call_type = f"{self._table_name}.{name}" + + @functools.wraps(attr) + async def _instrumented(*args: Any, **kwargs: Any) -> Any: + start_time = time.time() + try: + result = await attr(*args, **kwargs) + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self._service_logger_obj.async_service_success_hook( + service=service_type, + duration=_duration, + call_type=call_type, + start_time=start_time, + end_time=end_time, + ) + ) + return result + except Exception as e: + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self._service_logger_obj.async_service_failure_hook( + service=service_type, + duration=_duration, + error=e, + call_type=call_type, + start_time=start_time, + end_time=end_time, + ) + ) + raise + + return _instrumented + + +def _wrap_raw_query_method( + method: Callable, + service_type: ServiceTypes, + call_type: str, + service_logger_obj: "ServiceLogging", +) -> Callable: + """Wrap ``query_raw`` / ``query_first`` / ``execute_raw`` with metrics.""" + + @functools.wraps(method) + async def _instrumented(*args: Any, **kwargs: Any) -> Any: + start_time = time.time() + try: + result = await method(*args, **kwargs) + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + service_logger_obj.async_service_success_hook( + service=service_type, + duration=_duration, + call_type=call_type, + start_time=start_time, + end_time=end_time, + ) + ) + return result + except Exception as e: + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + service_logger_obj.async_service_failure_hook( + service=service_type, + duration=_duration, + error=e, + call_type=call_type, + start_time=start_time, + end_time=end_time, + ) + ) + raise + + return _instrumented + + +_RAW_METHOD_SERVICE_TYPES = { + "query_raw": ServiceTypes.DB_READ, + "query_first": ServiceTypes.DB_READ, + "execute_raw": ServiceTypes.DB_WRITE, +} + +_PASSTHROUGH_ATTRS: FrozenSet[str] = frozenset( + {"batch_", "tx", "connect", "disconnect", "is_connected"} +) class PrismaWrapper: @@ -36,9 +190,15 @@ class PrismaWrapper: # Fallback refresh interval if token parsing fails (10 minutes) FALLBACK_REFRESH_INTERVAL_SECONDS = 600 - def __init__(self, original_prisma: Any, iam_token_db_auth: bool): + def __init__( + self, + original_prisma: Any, + iam_token_db_auth: bool, + service_logger_obj: Optional["ServiceLogging"] = None, + ): self._original_prisma = original_prisma self.iam_token_db_auth = iam_token_db_auth + self._service_logger_obj = service_logger_obj # Background token refresh task management self._token_refresh_task: Optional[asyncio.Task] = None @@ -294,24 +454,22 @@ class PrismaWrapper: "Failed to generate new RDS IAM token during proactive refresh" ) - def __getattr__(self, name: str): + def __getattr__(self, name: str) -> Any: """ Proxy attribute access to the underlying Prisma client. - If IAM token auth is enabled and the token is expired, this method - provides a synchronous fallback to refresh the token. However, this - should rarely be needed since the background task proactively refreshes - tokens before they expire. - - FIXED: Now properly waits for reconnection to complete before returning, - instead of the previous fire-and-forget pattern that caused the bug. + Handles three concerns: + 1. IAM token refresh when tokens expire (RDS auth) + 2. Wrapping Prisma model objects with ``PrismaModelProxy`` for + per-query metrics (latency, count, errors) + 3. Wrapping ``query_raw`` / ``query_first`` / ``execute_raw`` + with equivalent metrics """ original_attr = getattr(self._original_prisma, name) if self.iam_token_db_auth: db_url = os.getenv("DATABASE_URL") - # Check if token is expired (should be rare if background task is running) if self.is_token_expired(db_url): verbose_proxy_logger.warning( "RDS IAM token expired in __getattr__ - proactive refresh may have failed. " @@ -323,13 +481,10 @@ class PrismaWrapper: loop = asyncio.get_event_loop() if loop.is_running(): - # FIXED: Actually wait for the reconnection to complete! - # The previous code used fire-and-forget which caused the bug. future = asyncio.run_coroutine_threadsafe( self.recreate_prisma_client(new_db_url), loop ) try: - # Wait up to 30 seconds for reconnection future.result(timeout=30) verbose_proxy_logger.info( "Synchronous token refresh completed successfully" @@ -342,11 +497,28 @@ class PrismaWrapper: else: asyncio.run(self.recreate_prisma_client(new_db_url)) - # Get the NEW attribute after reconnection original_attr = getattr(self._original_prisma, name) else: raise ValueError("Failed to get RDS IAM token") + if self._service_logger_obj is None or name in _PASSTHROUGH_ATTRS: + return original_attr + + if name in _RAW_METHOD_SERVICE_TYPES: + return _wrap_raw_query_method( + method=original_attr, + service_type=_RAW_METHOD_SERVICE_TYPES[name], + call_type=name, + service_logger_obj=self._service_logger_obj, + ) + + if hasattr(original_attr, "find_unique"): + return PrismaModelProxy( + model=original_attr, + table_name=name, + service_logger_obj=self._service_logger_obj, + ) + return original_attr diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6bbf0df74de..c605cd1d46b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2042,23 +2042,26 @@ class PrismaClient: raise Exception( "Unable to find Prisma binaries. Please run 'prisma generate' first." ) + from litellm._service_logger import ServiceLogging + + _iam_auth = ( + self.iam_token_db_auth + if self.iam_token_db_auth is not None + else False + ) + _service_logger_obj = ServiceLogging() + if http_client is not None: self.db = PrismaWrapper( original_prisma=Prisma(http=http_client), - iam_token_db_auth=( - self.iam_token_db_auth - if self.iam_token_db_auth is not None - else False - ), + iam_token_db_auth=_iam_auth, + service_logger_obj=_service_logger_obj, ) else: self.db = PrismaWrapper( original_prisma=Prisma(), - iam_token_db_auth=( - self.iam_token_db_auth - if self.iam_token_db_auth is not None - else False - ), + iam_token_db_auth=_iam_auth, + service_logger_obj=_service_logger_obj, ) # Client to connect to Prisma db verbose_proxy_logger.debug("Success - Created Prisma Client") diff --git a/litellm/types/services.py b/litellm/types/services.py index 580d5653dec..f275cfa478c 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -28,6 +28,12 @@ class ServiceTypes(str, enum.Enum): PROXY_PRE_CALL = "proxy_pre_call" POD_LOCK_MANAGER = "pod_lock_manager" + """ + Accurate DB operation metrics (instrumented at the Prisma query level) + """ + DB_READ = "db_read" + DB_WRITE = "db_write" + """ Operational metrics for DB Transaction Queues """ @@ -83,6 +89,13 @@ DEFAULT_SERVICE_CONFIGS = { ServiceTypes.PROXY_PRE_CALL.value: { "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] }, + # Accurate DB operation metrics (instrumented at Prisma query level) + ServiceTypes.DB_READ.value: { + "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] + }, + ServiceTypes.DB_WRITE.value: { + "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] + }, # Operational metrics for DB Transaction Queues ServiceTypes.POD_LOCK_MANAGER.value: {"metrics": [ServiceMetrics.GAUGE]}, ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE.value: { diff --git a/tests/test_litellm/integrations/test_prometheus_services.py b/tests/test_litellm/integrations/test_prometheus_services.py index ff80d7d9f8b..722ab15c63b 100644 --- a/tests/test_litellm/integrations/test_prometheus_services.py +++ b/tests/test_litellm/integrations/test_prometheus_services.py @@ -104,3 +104,152 @@ def test_update_gauge(): # Verify correct methods were called mock_labels.assert_called_once_with("test_label") mock_gauge.set.assert_called_once_with(42.5) + + +# --------------------------------------------------------------------------- +# DB read/write call_type label tests +# --------------------------------------------------------------------------- + +from litellm.types.services import ServiceLoggerPayload + + +def test_db_read_histogram_should_have_call_type_label(): + """DB_READ histograms should include a call_type label.""" + pl = PrometheusServicesLogger() + histogram = pl.create_histogram( + service=ServiceTypes.DB_READ.value, type_of_request="latency" + ) + assert "call_type" in histogram._labelnames + + +def test_db_write_histogram_should_have_call_type_label(): + """DB_WRITE histograms should include a call_type label.""" + pl = PrometheusServicesLogger() + histogram = pl.create_histogram( + service=ServiceTypes.DB_WRITE.value, type_of_request="latency" + ) + assert "call_type" in histogram._labelnames + + +def test_db_read_counter_should_have_call_type_label(): + """DB_READ counters should include a call_type label.""" + pl = PrometheusServicesLogger() + counter = pl.create_counter( + service=ServiceTypes.DB_READ.value, type_of_request="total_requests" + ) + assert "call_type" in counter._labelnames + + +def test_non_db_histogram_should_not_have_call_type_label(): + """Non-DB service histograms should NOT get a call_type label.""" + pl = PrometheusServicesLogger() + histogram = pl.create_histogram( + service=ServiceTypes.REDIS.value, type_of_request="latency" + ) + assert "call_type" not in histogram._labelnames + + +def test_service_success_hook_should_pass_call_type_for_db_read(): + """service_success_hook should pass call_type to histogram and total_requests counter for DB_READ.""" + pl = PrometheusServicesLogger() + payload = ServiceLoggerPayload( + is_error=False, + error=None, + service=ServiceTypes.DB_READ, + duration=0.05, + call_type="litellm_usertable.find_unique", + event_metadata=None, + ) + + pl.service_success_hook(payload) + + prom_objects = pl.payload_to_prometheus_map[ServiceTypes.DB_READ.value] + for obj in prom_objects: + if isinstance(obj, pl.Histogram): + assert "call_type" in obj._labelnames + child = obj.labels( + ServiceTypes.DB_READ.value, "litellm_usertable.find_unique" + ) + assert child is not None + elif isinstance(obj, pl.Counter) and "total_requests" in obj._name: + assert "call_type" in obj._labelnames + child = obj.labels( + ServiceTypes.DB_READ.value, "litellm_usertable.find_unique" + ) + assert child is not None + + +@pytest.mark.asyncio +async def test_async_service_success_hook_should_pass_call_type_for_db_write(): + """async_service_success_hook should pass call_type for DB_WRITE.""" + pl = PrometheusServicesLogger() + payload = ServiceLoggerPayload( + is_error=False, + error=None, + service=ServiceTypes.DB_WRITE, + duration=0.1, + call_type="litellm_teamtable.create", + event_metadata=None, + ) + + await pl.async_service_success_hook(payload) + + prom_objects = pl.payload_to_prometheus_map[ServiceTypes.DB_WRITE.value] + for obj in prom_objects: + if isinstance(obj, pl.Histogram): + assert "call_type" in obj._labelnames + child = obj.labels( + ServiceTypes.DB_WRITE.value, "litellm_teamtable.create" + ) + assert child is not None + + +@pytest.mark.asyncio +async def test_async_service_failure_hook_should_pass_call_type_for_db_read(): + """async_service_failure_hook should pass call_type for DB_READ failures.""" + pl = PrometheusServicesLogger() + payload = ServiceLoggerPayload( + is_error=True, + error="connection lost", + service=ServiceTypes.DB_READ, + duration=0.01, + call_type="litellm_usertable.find_unique", + event_metadata=None, + ) + + await pl.async_service_failure_hook( + payload=payload, + error=RuntimeError("connection lost"), + ) + + prom_objects = pl.payload_to_prometheus_map[ServiceTypes.DB_READ.value] + for obj in prom_objects: + if isinstance(obj, pl.Counter) and "failed_requests" in obj._name: + assert "call_type" in obj._labelnames + child = obj.labels( + ServiceTypes.DB_READ.value, + "litellm_usertable.find_unique", + "RuntimeError", + "litellm_usertable.find_unique", + ) + assert child is not None + + +def test_service_success_hook_should_not_pass_call_type_for_redis(): + """service_success_hook should NOT pass call_type for non-DB services like REDIS.""" + pl = PrometheusServicesLogger() + payload = ServiceLoggerPayload( + is_error=False, + error=None, + service=ServiceTypes.REDIS, + duration=0.01, + call_type="async_get_cache", + event_metadata=None, + ) + + pl.service_success_hook(payload) + + prom_objects = pl.payload_to_prometheus_map[ServiceTypes.REDIS.value] + for obj in prom_objects: + if hasattr(obj, "_labelnames"): + assert "call_type" not in obj._labelnames diff --git a/tests/test_litellm/proxy/db/test_prisma_db_metrics.py b/tests/test_litellm/proxy/db/test_prisma_db_metrics.py new file mode 100644 index 00000000000..91fca8dd1a5 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_prisma_db_metrics.py @@ -0,0 +1,298 @@ +""" +Tests for PrismaModelProxy and PrismaWrapper DB metrics instrumentation. + +Verifies that actual Prisma CRUD operations are instrumented with +ServiceTypes.DB_READ / DB_WRITE via ServiceLogging, and that +non-CRUD / passthrough attributes are forwarded unchanged. +""" + +import asyncio +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.db.prisma_client import ( + _DB_READ_METHODS, + _DB_WRITE_METHODS, + _PASSTHROUGH_ATTRS, + _RAW_METHOD_SERVICE_TYPES, + PrismaModelProxy, + PrismaWrapper, + _wrap_raw_query_method, +) +from litellm.types.services import ServiceTypes + + +def _make_mock_service_logger(): + """Create a mock ServiceLogging with async hooks.""" + logger = MagicMock() + logger.async_service_success_hook = AsyncMock() + logger.async_service_failure_hook = AsyncMock() + return logger + + +def _make_mock_model(): + """Create a mock Prisma model with all CRUD methods as AsyncMocks.""" + model = MagicMock() + for method_name in _DB_READ_METHODS | _DB_WRITE_METHODS: + setattr(model, method_name, AsyncMock(return_value={"id": "test"})) + return model + + +# --------------------------------------------------------------------------- +# PrismaModelProxy tests +# --------------------------------------------------------------------------- + + +class TestPrismaModelProxy: + @pytest.mark.parametrize("method_name", sorted(_DB_READ_METHODS)) + @pytest.mark.asyncio + async def test_should_log_read_methods_as_db_read(self, method_name): + model = _make_mock_model() + logger = _make_mock_service_logger() + proxy = PrismaModelProxy(model, "litellm_usertable", logger) + + wrapped = getattr(proxy, method_name) + result = await wrapped(where={"id": "x"}) + + assert result == {"id": "test"} + logger.async_service_success_hook.assert_called_once() + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_READ + assert call_kwargs["call_type"] == f"litellm_usertable.{method_name}" + assert call_kwargs["duration"] >= 0 + + @pytest.mark.parametrize("method_name", sorted(_DB_WRITE_METHODS)) + @pytest.mark.asyncio + async def test_should_log_write_methods_as_db_write(self, method_name): + model = _make_mock_model() + logger = _make_mock_service_logger() + proxy = PrismaModelProxy(model, "litellm_teamtable", logger) + + wrapped = getattr(proxy, method_name) + result = await wrapped(data={"name": "test"}) + + assert result == {"id": "test"} + logger.async_service_success_hook.assert_called_once() + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_WRITE + assert call_kwargs["call_type"] == f"litellm_teamtable.{method_name}" + + @pytest.mark.asyncio + async def test_should_log_failure_on_exception(self): + model = _make_mock_model() + model.find_unique = AsyncMock(side_effect=RuntimeError("connection lost")) + logger = _make_mock_service_logger() + proxy = PrismaModelProxy(model, "litellm_usertable", logger) + + with pytest.raises(RuntimeError, match="connection lost"): + await proxy.find_unique(where={"id": "x"}) + + logger.async_service_failure_hook.assert_called_once() + call_kwargs = logger.async_service_failure_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_READ + assert call_kwargs["call_type"] == "litellm_usertable.find_unique" + assert isinstance(call_kwargs["error"], RuntimeError) + + def test_should_passthrough_non_crud_attributes(self): + model = _make_mock_model() + model.some_property = "hello" + logger = _make_mock_service_logger() + proxy = PrismaModelProxy(model, "litellm_usertable", logger) + + assert proxy.some_property == "hello" + + @pytest.mark.asyncio + async def test_should_capture_duration(self): + async def slow_find(*args, **kwargs): + await asyncio.sleep(0.05) + return {"id": "slow"} + + model = _make_mock_model() + model.find_unique = slow_find + logger = _make_mock_service_logger() + proxy = PrismaModelProxy(model, "litellm_usertable", logger) + + await proxy.find_unique(where={"id": "x"}) + + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["duration"] >= 0.04 + + +# --------------------------------------------------------------------------- +# _wrap_raw_query_method tests +# --------------------------------------------------------------------------- + + +class TestWrapRawQueryMethod: + @pytest.mark.asyncio + async def test_should_log_query_raw_as_db_read(self): + logger = _make_mock_service_logger() + raw_method = AsyncMock(return_value=[{"count": 1}]) + + wrapped = _wrap_raw_query_method( + method=raw_method, + service_type=ServiceTypes.DB_READ, + call_type="query_raw", + service_logger_obj=logger, + ) + + result = await wrapped("SELECT 1") + assert result == [{"count": 1}] + logger.async_service_success_hook.assert_called_once() + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_READ + assert call_kwargs["call_type"] == "query_raw" + + @pytest.mark.asyncio + async def test_should_log_execute_raw_as_db_write(self): + logger = _make_mock_service_logger() + raw_method = AsyncMock(return_value=5) + + wrapped = _wrap_raw_query_method( + method=raw_method, + service_type=ServiceTypes.DB_WRITE, + call_type="execute_raw", + service_logger_obj=logger, + ) + + result = await wrapped("UPDATE foo SET bar = 1") + assert result == 5 + logger.async_service_success_hook.assert_called_once() + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_WRITE + + @pytest.mark.asyncio + async def test_should_log_failure_for_raw_method(self): + logger = _make_mock_service_logger() + raw_method = AsyncMock(side_effect=ConnectionError("db unreachable")) + + wrapped = _wrap_raw_query_method( + method=raw_method, + service_type=ServiceTypes.DB_READ, + call_type="query_raw", + service_logger_obj=logger, + ) + + with pytest.raises(ConnectionError, match="db unreachable"): + await wrapped("SELECT 1") + + logger.async_service_failure_hook.assert_called_once() + call_kwargs = logger.async_service_failure_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_READ + assert isinstance(call_kwargs["error"], ConnectionError) + + +# --------------------------------------------------------------------------- +# PrismaWrapper instrumentation tests +# --------------------------------------------------------------------------- + + +class TestPrismaWrapperInstrumentation: + def _make_wrapper_with_mock_prisma(self): + """Build a PrismaWrapper with a mock Prisma client and service logger.""" + mock_prisma = MagicMock() + + mock_user_model = MagicMock() + mock_user_model.find_unique = AsyncMock(return_value={"id": "u1"}) + mock_user_model.create = AsyncMock(return_value={"id": "u1"}) + mock_prisma.litellm_usertable = mock_user_model + + mock_prisma.query_raw = AsyncMock(return_value=[{"count": 42}]) + mock_prisma.query_first = AsyncMock(return_value={"id": "first"}) + mock_prisma.execute_raw = AsyncMock(return_value=3) + + mock_prisma.connect = AsyncMock() + mock_prisma.disconnect = AsyncMock() + mock_prisma.batch_ = MagicMock() + mock_prisma.tx = MagicMock() + + logger = _make_mock_service_logger() + + wrapper = PrismaWrapper( + original_prisma=mock_prisma, + iam_token_db_auth=False, + service_logger_obj=logger, + ) + + return wrapper, mock_prisma, logger + + def test_should_wrap_model_in_proxy(self): + wrapper, _, _ = self._make_wrapper_with_mock_prisma() + model_proxy = wrapper.litellm_usertable + assert isinstance(model_proxy, PrismaModelProxy) + + @pytest.mark.asyncio + async def test_should_instrument_model_find_unique(self): + wrapper, mock_prisma, logger = self._make_wrapper_with_mock_prisma() + result = await wrapper.litellm_usertable.find_unique(where={"id": "u1"}) + + assert result == {"id": "u1"} + mock_prisma.litellm_usertable.find_unique.assert_called_once_with( + where={"id": "u1"} + ) + logger.async_service_success_hook.assert_called_once() + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_READ + + @pytest.mark.asyncio + async def test_should_instrument_model_create(self): + wrapper, mock_prisma, logger = self._make_wrapper_with_mock_prisma() + result = await wrapper.litellm_usertable.create(data={"name": "new"}) + + assert result == {"id": "u1"} + logger.async_service_success_hook.assert_called_once() + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_WRITE + + @pytest.mark.asyncio + async def test_should_wrap_query_raw_as_read(self): + wrapper, _, logger = self._make_wrapper_with_mock_prisma() + wrapped = wrapper.query_raw + result = await wrapped("SELECT 1") + + assert result == [{"count": 42}] + logger.async_service_success_hook.assert_called_once() + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_READ + assert call_kwargs["call_type"] == "query_raw" + + @pytest.mark.asyncio + async def test_should_wrap_execute_raw_as_write(self): + wrapper, _, logger = self._make_wrapper_with_mock_prisma() + wrapped = wrapper.execute_raw + result = await wrapped("DELETE FROM foo") + + assert result == 3 + logger.async_service_success_hook.assert_called_once() + call_kwargs = logger.async_service_success_hook.call_args[1] + assert call_kwargs["service"] == ServiceTypes.DB_WRITE + assert call_kwargs["call_type"] == "execute_raw" + + @pytest.mark.parametrize("attr_name", sorted(_PASSTHROUGH_ATTRS)) + def test_should_passthrough_control_attributes(self, attr_name): + wrapper, mock_prisma, logger = self._make_wrapper_with_mock_prisma() + result = getattr(wrapper, attr_name) + + assert result is getattr(mock_prisma, attr_name) + logger.async_service_success_hook.assert_not_called() + + def test_should_skip_instrumentation_when_no_logger(self): + mock_prisma = MagicMock() + mock_model = MagicMock() + mock_model.find_unique = AsyncMock() + mock_prisma.litellm_usertable = mock_model + + wrapper = PrismaWrapper( + original_prisma=mock_prisma, + iam_token_db_auth=False, + service_logger_obj=None, + ) + + result = wrapper.litellm_usertable + assert not isinstance(result, PrismaModelProxy) + assert result is mock_model