fix(proxy_server.py): do a health check on db before returning if proxy ready (if db connected)

This commit is contained in:
Krrish Dholakia 2024-02-06 12:57:05 -08:00
parent e542aa10cc
commit e5fec98e1f
2 changed files with 24 additions and 11 deletions

View file

@ -472,8 +472,6 @@ class PrismaClient:
reset_at: Optional[datetime] = None,
):
try:
print_verbose("PrismaClient: get_data")
response: Any = None
if token is not None or (table_name is not None and table_name == "key"):
# check if plain text or hash
@ -885,6 +883,21 @@ class PrismaClient:
)
raise e
async def health_check(self):
"""
Health check endpoint for the prisma client
"""
sql_query = """
SELECT 1
FROM "LiteLLM_VerificationToken"
LIMIT 1
"""
# Execute the raw query
# The asterisk before `user_id_list` unpacks the list into separate arguments
response = await self.db.query_raw(sql_query)
return response
class DBClient:
"""