diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 34e7d34bbf3..9de845397aa 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -20,6 +20,7 @@ from litellm.proxy._types import ( WebhookEvent, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.health_check import ( _clean_endpoint_data, _update_litellm_params_for_health_check, @@ -381,20 +382,23 @@ async def _db_health_readiness_check(): global db_health_cache # Note - Intentionally don't try/except this so it raises an exception when it fails + try: + # if timedelta is less than 2 minutes return DB Status + time_diff = datetime.now() - db_health_cache["last_updated"] + if db_health_cache["status"] != "unknown" and time_diff < timedelta(minutes=2): + return db_health_cache - # if timedelta is less than 2 minutes return DB Status - time_diff = datetime.now() - db_health_cache["last_updated"] - if db_health_cache["status"] != "unknown" and time_diff < timedelta(minutes=2): + if prisma_client is None: + db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} + return db_health_cache + + await prisma_client.health_check() + db_health_cache = {"status": "connected", "last_updated": datetime.now()} return db_health_cache - - if prisma_client is None: - db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} + except Exception as e: + PrismaDBExceptionHandler.handle_db_exception(e) return db_health_cache - await prisma_client.health_check() - db_health_cache = {"status": "connected", "last_updated": datetime.now()} - return db_health_cache - @router.get( "/settings", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index eefc0bdd06e..0e04e9f2132 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -176,6 +176,7 @@ from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES from litellm.proxy.credential_endpoints.endpoints import router as credential_router +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.guardrails.guardrail_endpoints import router as guardrails_router @@ -456,15 +457,6 @@ async def proxy_startup_event(app: FastAPI): ### LOAD MASTER KEY ### # check if master key set in environment - load from there master_key = get_secret("LITELLM_MASTER_KEY", None) # type: ignore - # check if DATABASE_URL in environment - load from there - if prisma_client is None: - _db_url: Optional[str] = get_secret("DATABASE_URL", None) # type: ignore - prisma_client = await ProxyStartupEvent._setup_prisma_client( - database_url=_db_url, - proxy_logging_obj=proxy_logging_obj, - user_api_key_cache=user_api_key_cache, - ) - ## CHECK PREMIUM USER verbose_proxy_logger.debug( "litellm.proxy.proxy_server.py::startup() - CHECKING PREMIUM USER - {}".format( @@ -527,6 +519,15 @@ async def proxy_startup_event(app: FastAPI): redis_usage_cache=redis_usage_cache, ) + # check if DATABASE_URL in environment - load from there + if prisma_client is None: + _db_url: Optional[str] = get_secret("DATABASE_URL", None) # type: ignore + prisma_client = await ProxyStartupEvent._setup_prisma_client( + database_url=_db_url, + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + ) + ## JWT AUTH ## ProxyStartupEvent._initialize_jwt_auth( general_settings=general_settings, @@ -3362,33 +3363,36 @@ class ProxyStartupEvent: - Sets up prisma client - Adds necessary views to proxy """ - prisma_client: Optional[PrismaClient] = None - if database_url is not None: - try: - prisma_client = PrismaClient( - database_url=database_url, proxy_logging_obj=proxy_logging_obj - ) - except Exception as e: - raise e + try: + prisma_client: Optional[PrismaClient] = None + if database_url is not None: + try: + prisma_client = PrismaClient( + database_url=database_url, proxy_logging_obj=proxy_logging_obj + ) + except Exception as e: + raise e - await prisma_client.connect() + await prisma_client.connect() - ## Add necessary views to proxy ## - asyncio.create_task( - prisma_client.check_view_exists() - ) # check if all necessary views exist. Don't block execution + ## Add necessary views to proxy ## + asyncio.create_task( + prisma_client.check_view_exists() + ) # check if all necessary views exist. Don't block execution - asyncio.create_task( - prisma_client._set_spend_logs_row_count_in_proxy_state() - ) # set the spend logs row count in proxy state. Don't block execution + asyncio.create_task( + prisma_client._set_spend_logs_row_count_in_proxy_state() + ) # set the spend logs row count in proxy state. Don't block execution - # run a health check to ensure the DB is ready - if ( - get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) - is not True - ): - await prisma_client.health_check() - return prisma_client + # run a health check to ensure the DB is ready + if ( + get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) + is not True + ): + await prisma_client.health_check() + return prisma_client + except Exception as e: + PrismaDBExceptionHandler.handle_db_exception(e) @classmethod def _init_dd_tracer(cls):