Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/log_db_metrics.py: 94%

1""" 

2Handles logging DB success/failure to ServiceLogger() 

3 

4ServiceLogger() then sends DB logs to Prometheus, OTEL, Datadog etc 

5""" 

6 

7import asyncio 

8from collections.abc import Callable 

9from datetime import datetime 

10from functools import wraps 

11from typing import Final 

12 

13from litellm._service_logger import ServiceTypes 

14from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs 

15 

16 

17def _safe_db_event_metadata(kwargs: dict) -> dict[str, str] | None: 

18 """Minimal, non-sensitive ``event_metadata`` for a DB service log. 

19 

20 The raw ``kwargs``/``args`` carry live objects (Prisma client, OTel spans) 

21 and secrets (tokens), none of which belongs on a span — so we surface only 

22 the table name when present. Everything else is dropped. 

23 """ 

24 table_name: Final = kwargs.get("table_name") 

25 return {"table_name": table_name} if isinstance(table_name, str) else None 

26 

27 

28def log_db_metrics(func): 

29 """ 

30 Decorator to log the duration of a DB related function to ServiceLogger() 

31 

32 Handles logging DB success/failure to ServiceLogger(), which logs to Prometheus, OTEL, Datadog 

33 

34 When logging Failure it checks if the Exception is a PrismaError, httpx.ConnectError or httpx.TimeoutException and then logs that as a DB Service Failure 

35 

36 Args: 

37 func: The function to be decorated 

38 

39 Returns: 

40 Result from the decorated function 

41 

42 Raises: 

43 Exception: If the decorated function raises an exception 

44 """ 

45 

46 @wraps(func) 

47 async def wrapper(*args, **kwargs): 

48 start_time: Final[datetime] = datetime.now() 

49 

50 try: 

51 result: Final = await func(*args, **kwargs) 

52 end_time: datetime = datetime.now() 

53 from litellm.proxy.proxy_server import proxy_logging_obj 

54 

55 if "PROXY" not in func.__name__: 

56 asyncio.create_task( 

57 proxy_logging_obj.service_logging_obj.async_service_success_hook( 

58 service=ServiceTypes.DB, 

59 call_type=func.__name__, 

60 parent_otel_span=kwargs.get("parent_otel_span", None), 

61 duration=(end_time - start_time).total_seconds(), 

62 start_time=start_time, 

63 end_time=end_time, 

64 event_metadata=_safe_db_event_metadata(kwargs), 

65 ) 

66 ) 

67 elif ( 67 ↛ 90line 67 didn't jump to line 90 because the condition on line 67 was always true

68 # in litellm custom callbacks kwargs is passed as arg[0] 

69 # https://docs.litellm.ai/docs/observability/custom_callback#callback-functions 

70 args is not None and len(args) > 1 and isinstance(args[1], dict) 

71 ): 

72 passed_kwargs: Final = args[1] 

73 parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs=passed_kwargs) 

74 if parent_otel_span is not None: 74 ↛ 78line 74 didn't jump to line 78 because the condition on line 74 was never true

75 # No metadata dump: identity rides on Baggage, and the full 

76 # request metadata (auth blob, response headers, tokens) must 

77 # not land on a span. 

78 asyncio.create_task( 

79 proxy_logging_obj.service_logging_obj.async_service_success_hook( 

80 service=ServiceTypes.BATCH_WRITE_TO_DB, 

81 call_type=func.__name__, 

82 parent_otel_span=parent_otel_span, 

83 duration=0.0, 

84 start_time=start_time, 

85 end_time=end_time, 

86 event_metadata=None, 

87 ) 

88 ) 

89 # end of logging to otel 

90 return result 

91 except Exception as e: 

92 end_time: datetime = datetime.now() 

93 await _handle_logging_db_exception( 

94 e=e, 

95 func=func, 

96 kwargs=kwargs, 

97 args=args, 

98 start_time=start_time, 

99 end_time=end_time, 

100 ) 

101 raise e 

102 

103 return wrapper 

104 

105 

106def _is_exception_related_to_db(e: Exception) -> bool: 

107 """ 

108 Returns True if the exception is related to the DB 

109 """ 

110 

111 import httpx 

112 from prisma.errors import PrismaError 

113 

114 return isinstance(e, (PrismaError, httpx.ConnectError, httpx.TimeoutException)) 

115 

116 

117async def _handle_logging_db_exception( 

118 e: Exception, 

119 func: Callable, 

120 kwargs: dict, 

121 args: tuple, 

122 start_time: datetime, 

123 end_time: datetime, 

124) -> None: 

125 from litellm.proxy.proxy_server import proxy_logging_obj 

126 

127 # don't log this as a DB Service Failure, if the DB did not raise an exception 

128 if _is_exception_related_to_db(e) is not True: 

129 return 

130 

131 await proxy_logging_obj.service_logging_obj.async_service_failure_hook( 

132 error=e, 

133 service=ServiceTypes.DB, 

134 call_type=func.__name__, 

135 parent_otel_span=kwargs.get("parent_otel_span"), 

136 duration=(end_time - start_time).total_seconds(), 

137 start_time=start_time, 

138 end_time=end_time, 

139 event_metadata=_safe_db_event_metadata(kwargs), 

140 )