scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

Merge pull request #24299 from BerriAI/litellm_1_81_6_prometheus

Browse source 
feat(prometheus.py): emit db event metrics
This commit is contained in:
Krish Dholakia 2026-03-21 15:27:22 -07:00 • committed by GitHub
commit 67afe827cc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 731 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
docs/my-website/docs/proxy/prometheus.md
View file

@ -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.

70
litellm/integrations/prometheus_services.py
View file

@ -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,
)

202
litellm/proxy/db/prisma_client.py
View file

@ -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

23
litellm/proxy/utils.py
View file

@ -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")

13
litellm/types/services.py
View file

@ -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: {

149
tests/test_litellm/integrations/test_prometheus_services.py
View file

@ -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

298
tests/test_litellm/proxy/db/test_prisma_db_metrics.py Normal file
View file

@ -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