From 7142b0b61060292b5a353da65eb95185d281fb4d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 19:22:24 -0700 Subject: [PATCH 1/9] refactor PrismaDBExceptionHandler --- litellm/proxy/auth/auth_exception_handler.py | 38 +++----------------- litellm/proxy/db/exception_handler.py | 37 +++++++++++++++++++ 2 files changed, 41 insertions(+), 34 deletions(-) create mode 100644 litellm/proxy/db/exception_handler.py diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index c1a546b569f..05797381c6c 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -9,13 +9,9 @@ from fastapi import HTTPException, Request, status import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, - ProxyErrorTypes, - ProxyException, - UserAPIKeyAuth, -) +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import _get_request_ip_address +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes if TYPE_CHECKING: @@ -58,8 +54,8 @@ class UserAPIKeyAuthExceptionHandler: ) if ( - UserAPIKeyAuthExceptionHandler.should_allow_request_on_db_unavailable() - and UserAPIKeyAuthExceptionHandler.is_database_connection_error(e) + PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + and PrismaDBExceptionHandler.is_database_connection_error(e) ): # log this as a DB failure on prometheus proxy_logging_obj.service_logging_obj.service_failure_hook( @@ -125,29 +121,3 @@ class UserAPIKeyAuthExceptionHandler: param=getattr(e, "param", "None"), code=status.HTTP_401_UNAUTHORIZED, ) - - @staticmethod - def should_allow_request_on_db_unavailable() -> bool: - """ - Returns True if the request should be allowed to proceed despite the DB connection error - """ - from litellm.proxy.proxy_server import general_settings - - if general_settings.get("allow_requests_on_db_unavailable", False) is True: - return True - return False - - @staticmethod - def is_database_connection_error(e: Exception) -> bool: - """ - Returns True if the exception is from a database outage / connection error - """ - import prisma - - if isinstance(e, DB_CONNECTION_ERROR_TYPES): - return True - if isinstance(e, prisma.errors.PrismaError): - return True - if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: - return True - return False diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py new file mode 100644 index 00000000000..315811103d9 --- /dev/null +++ b/litellm/proxy/db/exception_handler.py @@ -0,0 +1,37 @@ +from litellm.proxy._types import ( + DB_CONNECTION_ERROR_TYPES, + ProxyErrorTypes, + ProxyException, +) + + +class PrismaDBExceptionHandler: + """ + Class to handle DB Exceptions or Connection Errors + """ + + @staticmethod + def should_allow_request_on_db_unavailable() -> bool: + """ + Returns True if the request should be allowed to proceed despite the DB connection error + """ + from litellm.proxy.proxy_server import general_settings + + if general_settings.get("allow_requests_on_db_unavailable", False) is True: + return True + return False + + @staticmethod + def is_database_connection_error(e: Exception) -> bool: + """ + Returns True if the exception is from a database outage / connection error + """ + import prisma + + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return True + if isinstance(e, prisma.errors.PrismaError): + return True + if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: + return True + return False From 497570b2a6e26905d7a7564dbb6bf92b6fc0b343 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 19:49:51 -0700 Subject: [PATCH 2/9] bug fix - allow pods to startup when DB is unavailable --- litellm/proxy/db/exception_handler.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 315811103d9..db73f9e9c93 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,8 +1,11 @@ +from typing import Union + from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, ProxyErrorTypes, ProxyException, ) +from litellm.secret_managers.main import str_to_bool class PrismaDBExceptionHandler: @@ -17,7 +20,12 @@ class PrismaDBExceptionHandler: """ from litellm.proxy.proxy_server import general_settings - if general_settings.get("allow_requests_on_db_unavailable", False) is True: + _allow_requests_on_db_unavailable: Union[bool, str] = general_settings.get( + "allow_requests_on_db_unavailable", False + ) + if isinstance(_allow_requests_on_db_unavailable, bool): + return _allow_requests_on_db_unavailable + if str_to_bool(_allow_requests_on_db_unavailable) is True: return True return False @@ -35,3 +43,19 @@ class PrismaDBExceptionHandler: if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: return True return False + + @staticmethod + def handle_db_exception(e: Exception): + """ + Primary handler for `allow_requests_on_db_unavailable` flag. Decides whether to raise a DB Exception or not based on the flag. + + - If exception is a DB Connection Error, and `allow_requests_on_db_unavailable` is True, + - Do not raise an exception, return None + - Else, raise the exception + """ + if ( + PrismaDBExceptionHandler.is_database_connection_error(e) + and PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + ): + return None + raise e From 88ef97b9d1d9b2ecad5b9258ff980930ff694033 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 19:50:57 -0700 Subject: [PATCH 3/9] allow proxy to startup on DB unavailable --- .../health_endpoints/_health_endpoints.py | 24 ++++--- litellm/proxy/proxy_server.py | 68 ++++++++++--------- 2 files changed, 50 insertions(+), 42 deletions(-) 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): From 15c04da73599316e4fa929787fd272fed49b5cda Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 20:00:10 -0700 Subject: [PATCH 4/9] refactor tests --- litellm/proxy/proxy_server.py | 1 + .../proxy/auth/test_auth_exception_handler.py | 72 ------------ .../proxy/db/test_exception_handler.py | 109 ++++++++++++++++++ 3 files changed, 110 insertions(+), 72 deletions(-) create mode 100644 tests/litellm/proxy/db/test_exception_handler.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e04e9f2132..59b4bc8fe34 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3393,6 +3393,7 @@ class ProxyStartupEvent: return prisma_client except Exception as e: PrismaDBExceptionHandler.handle_db_exception(e) + return None @classmethod def _init_dd_tracer(cls): diff --git a/tests/litellm/proxy/auth/test_auth_exception_handler.py b/tests/litellm/proxy/auth/test_auth_exception_handler.py index b44de86e05f..224bf24b570 100644 --- a/tests/litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/litellm/proxy/auth/test_auth_exception_handler.py @@ -29,78 +29,6 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler -# Test is_database_connection_error method -@pytest.mark.parametrize( - "prisma_error", - [ - PrismaError(), - DataError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), - UniqueViolationError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), - ForeignKeyViolationError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), - MissingRequiredValueError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), - RawQueryError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), - TableNotFoundError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), - RecordNotFoundError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), - HTTPClientClosedError(), - ClientNotConnectedError(), - ], -) -def test_is_database_connection_error_prisma_errors(prisma_error): - """ - Test that all Prisma errors are considered database connection errors - """ - handler = UserAPIKeyAuthExceptionHandler() - assert handler.is_database_connection_error(prisma_error) == True - - -def test_is_database_connection_generic_errors(): - """ - Test non-Prisma error cases for database connection checking - """ - handler = UserAPIKeyAuthExceptionHandler() - - # Test with ProxyException (DB connection) - db_proxy_exception = ProxyException( - message="DB Connection Error", - type=ProxyErrorTypes.no_db_connection, - param="test-param", - ) - assert handler.is_database_connection_error(db_proxy_exception) == True - - # Test with non-DB error - regular_exception = Exception("Regular error") - assert handler.is_database_connection_error(regular_exception) == False - - -# Test should_allow_request_on_db_unavailable method -@patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, -) -def test_should_allow_request_on_db_unavailable_true(): - handler = UserAPIKeyAuthExceptionHandler() - assert handler.should_allow_request_on_db_unavailable() == True - - -@patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": False}, -) -def test_should_allow_request_on_db_unavailable_false(): - handler = UserAPIKeyAuthExceptionHandler() - assert handler.should_allow_request_on_db_unavailable() == False - - @pytest.mark.asyncio @pytest.mark.parametrize( "prisma_error", diff --git a/tests/litellm/proxy/db/test_exception_handler.py b/tests/litellm/proxy/db/test_exception_handler.py new file mode 100644 index 00000000000..04f17ff19b4 --- /dev/null +++ b/tests/litellm/proxy/db/test_exception_handler.py @@ -0,0 +1,109 @@ +import asyncio +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest +from fastapi import HTTPException, Request, status +from prisma import errors as prisma_errors +from prisma.errors import ( + ClientNotConnectedError, + DataError, + ForeignKeyViolationError, + HTTPClientClosedError, + MissingRequiredValueError, + PrismaError, + RawQueryError, + RecordNotFoundError, + TableNotFoundError, + UniqueViolationError, +) + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler + + +# Test is_database_connection_error method +@pytest.mark.parametrize( + "prisma_error", + [ + PrismaError(), + DataError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + UniqueViolationError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + ForeignKeyViolationError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + MissingRequiredValueError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + RawQueryError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + TableNotFoundError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + RecordNotFoundError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + HTTPClientClosedError(), + ClientNotConnectedError(), + ], +) +def test_is_database_connection_error_prisma_errors(prisma_error): + """ + Test that all Prisma errors are considered database connection errors + """ + assert PrismaDBExceptionHandler.is_database_connection_error(prisma_error) == True + + +def test_is_database_connection_generic_errors(): + """ + Test non-Prisma error cases for database connection checking + """ + assert ( + PrismaDBExceptionHandler.is_database_connection_error( + Exception("Regular error") + ) + == False + ) + + # Test with ProxyException (DB connection) + db_proxy_exception = ProxyException( + message="DB Connection Error", + type=ProxyErrorTypes.no_db_connection, + param="test-param", + ) + assert ( + PrismaDBExceptionHandler.is_database_connection_error(db_proxy_exception) + == True + ) + + # Test with non-DB error + regular_exception = Exception("Regular error") + assert ( + PrismaDBExceptionHandler.is_database_connection_error(regular_exception) + == False + ) + + +# Test should_allow_request_on_db_unavailable method +@patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, +) +def test_should_allow_request_on_db_unavailable_true(): + assert PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() == True + + +@patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, +) +def test_should_allow_request_on_db_unavailable_false(): + assert PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() == False From 427580eff512fbdafec90157d443c4d1de63332d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 20:35:27 -0700 Subject: [PATCH 5/9] fix _setup_prisma_client --- litellm/proxy/db/exception_handler.py | 5 +++-- litellm/proxy/proxy_server.py | 17 ++++++++--------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index db73f9e9c93..79229f1e1bb 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -18,9 +18,10 @@ class PrismaDBExceptionHandler: """ Returns True if the request should be allowed to proceed despite the DB connection error """ - from litellm.proxy.proxy_server import general_settings + from litellm.proxy.proxy_server import proxy_config - _allow_requests_on_db_unavailable: Union[bool, str] = general_settings.get( + _general_settings = proxy_config.config + _allow_requests_on_db_unavailable: Union[bool, str] = _general_settings.get( "allow_requests_on_db_unavailable", False ) if isinstance(_allow_requests_on_db_unavailable, bool): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 59b4bc8fe34..67d88b9d322 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -457,6 +457,14 @@ 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( @@ -519,15 +527,6 @@ 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, From 34e58be36dcf16d7d139b0b51601f6be18df0a24 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 20:42:01 -0700 Subject: [PATCH 6/9] fix order of _setup_prisma_client --- litellm/proxy/db/exception_handler.py | 5 ++--- litellm/proxy/proxy_server.py | 23 ++++++++++++----------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 79229f1e1bb..db73f9e9c93 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -18,10 +18,9 @@ class PrismaDBExceptionHandler: """ Returns True if the request should be allowed to proceed despite the DB connection error """ - from litellm.proxy.proxy_server import proxy_config + from litellm.proxy.proxy_server import general_settings - _general_settings = proxy_config.config - _allow_requests_on_db_unavailable: Union[bool, str] = _general_settings.get( + _allow_requests_on_db_unavailable: Union[bool, str] = general_settings.get( "allow_requests_on_db_unavailable", False ) if isinstance(_allow_requests_on_db_unavailable, bool): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 67d88b9d322..c2044ac1d64 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -454,17 +454,6 @@ async def proxy_startup_event(app: FastAPI): import json init_verbose_loggers() - ### 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( @@ -521,6 +510,18 @@ async def proxy_startup_event(app: FastAPI): if isinstance(worker_config, dict): await initialize(**worker_config) + ### 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, + ) + ProxyStartupEvent._initialize_startup_logging( llm_router=llm_router, proxy_logging_obj=proxy_logging_obj, From 87f0201f8400bee3550e1c82befccfc637fc4436 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 20:44:25 -0700 Subject: [PATCH 7/9] test_handle_db_exception_with_connection_error --- .../proxy/db/test_exception_handler.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/litellm/proxy/db/test_exception_handler.py b/tests/litellm/proxy/db/test_exception_handler.py index 04f17ff19b4..e68c9b6a995 100644 --- a/tests/litellm/proxy/db/test_exception_handler.py +++ b/tests/litellm/proxy/db/test_exception_handler.py @@ -24,6 +24,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -107,3 +108,41 @@ def test_should_allow_request_on_db_unavailable_true(): ) def test_should_allow_request_on_db_unavailable_false(): assert PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() == False + + +@patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, +) +def test_handle_db_exception_with_connection_error(): + """ + Test that DB connection errors are handled gracefully when allow_requests_on_db_unavailable is True + """ + db_error = ClientNotConnectedError() + result = PrismaDBExceptionHandler.handle_db_exception(db_error) + assert result is None + + +@patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, +) +def test_handle_db_exception_raises_error(): + """ + Test that DB connection errors are raised when allow_requests_on_db_unavailable is False + """ + db_error = ClientNotConnectedError() + with pytest.raises(ClientNotConnectedError): + PrismaDBExceptionHandler.handle_db_exception(db_error) + + +def test_handle_db_exception_with_non_db_error(): + """ + Test that non-DB errors are always raised regardless of allow_requests_on_db_unavailable setting + """ + regular_error = litellm.BudgetExceededError( + current_cost=10, + max_budget=10, + ) + with pytest.raises(litellm.BudgetExceededError): + PrismaDBExceptionHandler.handle_db_exception(regular_error) From b6506f7bda5bb2935c97e6ca2b75e306f2c7111f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 20:56:39 -0700 Subject: [PATCH 8/9] test_db_health_readiness_check_with_prisma_error --- .../health_endpoints/test_health_endpoints.py | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 tests/litellm/proxy/health_endpoints/test_health_endpoints.py diff --git a/tests/litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/litellm/proxy/health_endpoints/test_health_endpoints.py new file mode 100644 index 00000000000..e2dd429357a --- /dev/null +++ b/tests/litellm/proxy/health_endpoints/test_health_endpoints.py @@ -0,0 +1,101 @@ +import asyncio +import json +import os +import sys +from datetime import datetime, timedelta +from unittest.mock import MagicMock, patch + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import pytest +from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, PrismaError + +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.health_endpoints._health_endpoints import ( + _db_health_readiness_check, + db_health_cache, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma_error", + [ + PrismaError(), + ClientNotConnectedError(), + HTTPClientClosedError(), + ], +) +async def test_db_health_readiness_check_with_prisma_error(prisma_error): + """ + Test that when prisma_client.health_check() raises a PrismaError and + allow_requests_on_db_unavailable is True, the function should not raise an error + and return the cached health status. + """ + # Mock the prisma client + mock_prisma_client = MagicMock() + mock_prisma_client.health_check.side_effect = prisma_error + + # Reset the health cache to a known state + global db_health_cache + db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(minutes=5), + } + + # Patch the imports and general_settings + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, + ): + + # Call the function + result = await _db_health_readiness_check() + + # Verify that the function called health_check + mock_prisma_client.health_check.assert_called_once() + + # Verify that the function returned the cache + assert result is not None + assert result["status"] == "unknown" # Should retain the status from the cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma_error", + [ + PrismaError(), + ClientNotConnectedError(), + HTTPClientClosedError(), + ], +) +async def test_db_health_readiness_check_with_error_and_flag_off(prisma_error): + """ + Test that when prisma_client.health_check() raises a DB error but + allow_requests_on_db_unavailable is False, the exception should be raised. + """ + # Mock the prisma client + mock_prisma_client = MagicMock() + mock_prisma_client.health_check.side_effect = prisma_error + + # Reset the health cache + global db_health_cache + db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(minutes=5), + } + + # Patch the imports and general_settings where the flag is False + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ): + + # The function should raise the exception + with pytest.raises(Exception) as excinfo: + await _db_health_readiness_check() + + # Verify that the raised exception is the same + assert excinfo.value == prisma_error From 05c38049feebd8cc113a361b3834331910dcb6d8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 21:04:36 -0700 Subject: [PATCH 9/9] docs prod.md --- docs/my-website/docs/proxy/prod.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/my-website/docs/proxy/prod.md b/docs/my-website/docs/proxy/prod.md index 314300f2a05..1c1cbeedb4e 100644 --- a/docs/my-website/docs/proxy/prod.md +++ b/docs/my-website/docs/proxy/prod.md @@ -114,6 +114,8 @@ When `allow_requests_on_db_unavailable` is set to `true`, LiteLLM will handle er |---------------|-------------------|----------------| | Prisma Errors | ✅ Request will be allowed | Covers issues like DB connection resets or rejections from the DB via Prisma, the ORM used by LiteLLM. | | Httpx Errors | ✅ Request will be allowed | Occurs when the database is unreachable, allowing the request to proceed despite the DB outage. | +| Pod Startup Behavior | ✅ Pods start regardless | LiteLLM Pods will start even if the database is down or unreachable, ensuring higher uptime guarantees for deployments. | +| Health/Readiness Check | ✅ Always returns 200 OK | The /health/readiness endpoint returns a 200 OK status to ensure that pods remain operational even when the database is unavailable. | LiteLLM Budget Errors or Model Errors | ❌ Request will be blocked | Triggered when the DB is reachable but the authentication token is invalid, lacks access, or exceeds budget limits. |