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