diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 036b3174b6..5b11c25b2b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2817,6 +2817,13 @@ class PrismaClient: ) raise e + @backoff.on_exception( + backoff.expo, + Exception, + max_tries=3, + max_time=10, + on_backoff=on_backoff, + ) async def health_check(self): """ Health check endpoint for the prisma client @@ -2848,7 +2855,18 @@ class PrismaClient: raise e async def _get_spend_logs_row_count(self) -> int: - try: + """ + Get the row count from LiteLLM_SpendLogs table using PostgreSQL system statistics. + """ + + @backoff.on_exception( + backoff.expo, + Exception, + max_tries=3, + max_time=10, + on_backoff=on_backoff, + ) + async def _fetch_row_count() -> int: sql_query = """ SELECT reltuples::BIGINT FROM pg_class @@ -2856,12 +2874,22 @@ class PrismaClient: """ result = await self.db.query_raw(query=sql_query) return result[0]["reltuples"] + + try: + return await _fetch_row_count() except Exception as e: verbose_proxy_logger.error( f"Error getting LiteLLM_SpendLogs row count: {e}" ) return 0 + @backoff.on_exception( + backoff.expo, + Exception, + max_tries=3, + max_time=10, + on_backoff=on_backoff, + ) async def _set_spend_logs_row_count_in_proxy_state(self) -> None: """ Set the `LiteLLM_SpendLogs`row count in proxy state. diff --git a/tests/proxy_unit_tests/test_prisma_client_backoff_retry.py b/tests/proxy_unit_tests/test_prisma_client_backoff_retry.py new file mode 100644 index 0000000000..ca0519f0d9 --- /dev/null +++ b/tests/proxy_unit_tests/test_prisma_client_backoff_retry.py @@ -0,0 +1,374 @@ +""" +Test backoff retry mechanisms for PrismaClient methods during _setup_prisma_client state. + +This test validates that intermittent database connection issues are handled correctly +with exponential backoff retries for critical startup operations. +""" + +import asyncio +import pytest +import time +from unittest.mock import AsyncMock, MagicMock, patch, call +from unittest.mock import Mock +import sys +import os + +# Add project root to path +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.proxy.utils import PrismaClient, ProxyLogging +from prisma.errors import PrismaError, ClientNotConnectedError +import httpx +import backoff + + +@pytest.fixture +def mock_prisma_client(): + """Create a mock PrismaClient with necessary attributes""" + client = MagicMock(spec=PrismaClient) + + # Mock database connection + client.db = AsyncMock() + client.db.query_raw = AsyncMock() + + # Mock proxy logging object + client.proxy_logging_obj = AsyncMock() + client.proxy_logging_obj.failure_handler = AsyncMock() + + return client + + +@pytest.fixture +def mock_proxy_logging(): + """Create a mock ProxyLogging object""" + proxy_logging = AsyncMock(spec=ProxyLogging) + proxy_logging.failure_handler = AsyncMock() + return proxy_logging + + +@pytest.fixture +def connection_errors(): + """Common database connection errors that should trigger retries""" + return [ + httpx.ConnectError("Connection failed"), + httpx.TimeoutException("Request timeout"), + ClientNotConnectedError(), + PrismaError("Database connection lost"), + ConnectionError("Network connection failed"), + OSError("Connection refused"), + ] + + +class TestPrismaClientBackoffRetry: + """Test suite for PrismaClient backoff retry mechanisms""" + + @pytest.mark.asyncio + async def test_health_check_success_no_retry( + self, mock_prisma_client, mock_proxy_logging + ): + """Test health_check succeeds immediately without retries""" + # Mock successful query response + mock_prisma_client.db.query_raw.return_value = [{"result": 1}] + + # Create real PrismaClient instance with mocked db + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + # Call health_check + result = await client.health_check() + + # Verify success + assert result == [{"result": 1}] + mock_prisma_client.db.query_raw.assert_called_once_with("SELECT 1") + + @pytest.mark.asyncio + async def test_health_check_retry_then_success( + self, mock_prisma_client, mock_proxy_logging, connection_errors + ): + """Test health_check retries on connection errors and eventually succeeds""" + # Mock first two calls to fail, third to succeed + mock_prisma_client.db.query_raw.side_effect = [ + connection_errors[0], # First call fails + connection_errors[1], # Second call fails + [{"result": 1}], # Third call succeeds + ] + + # Create real PrismaClient instance + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + # Measure execution time to verify backoff delay + start_time = time.time() + result = await client.health_check() + end_time = time.time() + + # Verify eventual success + assert result == [{"result": 1}] + + # Verify retry attempts (3 calls total) + assert mock_prisma_client.db.query_raw.call_count == 3 + + # Verify backoff delay occurred (should be at least a few milliseconds) + assert end_time - start_time > 0.01 + + @pytest.mark.asyncio + async def test_health_check_max_retries_exceeded( + self, mock_prisma_client, mock_proxy_logging, connection_errors + ): + """Test health_check fails after max retries (3) are exceeded""" + # Mock all calls to fail + mock_prisma_client.db.query_raw.side_effect = connection_errors[0] + + # Create real PrismaClient instance + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + # Expect final exception after retries + with pytest.raises(httpx.ConnectError): + await client.health_check() + + # Verify max retries attempted (3 attempts) + assert mock_prisma_client.db.query_raw.call_count == 3 + + @pytest.mark.asyncio + async def test_get_spend_logs_row_count_success_no_retry( + self, mock_prisma_client, mock_proxy_logging + ): + """Test _get_spend_logs_row_count succeeds immediately""" + # Mock successful query response + mock_prisma_client.db.query_raw.return_value = [{"reltuples": 1000}] + + # Create real PrismaClient instance + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + result = await client._get_spend_logs_row_count() + + assert result == 1000 + mock_prisma_client.db.query_raw.assert_called_once() + + @pytest.mark.asyncio + async def test_get_spend_logs_row_count_retry_then_success( + self, mock_prisma_client, mock_proxy_logging, connection_errors + ): + """Test _get_spend_logs_row_count retries and eventually succeeds""" + # Mock first call fails, second succeeds + mock_prisma_client.db.query_raw.side_effect = [ + connection_errors[2], # First call fails (ClientNotConnectedError) + [{"reltuples": 500}], # Second call succeeds + ] + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + result = await client._get_spend_logs_row_count() + + assert mock_prisma_client.db.query_raw.call_count == 2 + + @pytest.mark.asyncio + async def test_get_spend_logs_row_count_handles_errors_gracefully( + self, mock_prisma_client, mock_proxy_logging + ): + """Test _get_spend_logs_row_count returns 0 on persistent errors""" + # Mock all calls to fail + mock_prisma_client.db.query_raw.side_effect = PrismaError("Persistent DB error") + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + result = await client._get_spend_logs_row_count() + + assert result == 0 + assert mock_prisma_client.db.query_raw.call_count == 3 # Max retries attempted + + @pytest.mark.asyncio + async def test_set_spend_logs_row_count_in_proxy_state_success( + self, mock_prisma_client, mock_proxy_logging + ): + """Test _set_spend_logs_row_count_in_proxy_state succeeds""" + # Mock successful query response + mock_prisma_client.db.query_raw.return_value = [{"reltuples": 2000}] + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + # Mock proxy_state + with patch("litellm.proxy.proxy_server.proxy_state") as mock_proxy_state: + mock_proxy_state.set_proxy_state_variable = Mock() + + await client._set_spend_logs_row_count_in_proxy_state() + + # Verify proxy state was updated + mock_proxy_state.set_proxy_state_variable.assert_called_once_with( + variable_name="spend_logs_row_count", value=2000 + ) + + @pytest.mark.asyncio + async def test_set_spend_logs_row_count_retry_behavior( + self, mock_prisma_client, mock_proxy_logging, connection_errors + ): + """Test _set_spend_logs_row_count_in_proxy_state retries on database errors""" + mock_prisma_client.db.query_raw.side_effect = [ + connection_errors[3], # First call fails (PrismaError) + [{"reltuples": 1500}], # Second call succeeds + ] + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + with patch("litellm.proxy.proxy_server.proxy_state") as mock_proxy_state: + mock_proxy_state.set_proxy_state_variable = Mock() + + await client._set_spend_logs_row_count_in_proxy_state() + + assert mock_prisma_client.db.query_raw.call_count == 2 + mock_proxy_state.set_proxy_state_variable.assert_called_once_with( + variable_name="spend_logs_row_count", value=1500 + ) + + @pytest.mark.asyncio + async def test_backoff_configuration_parameters(self, mock_proxy_logging): + """Test that backoff decorators are configured with correct parameters""" + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + + # Check that methods have backoff decorators + assert hasattr(client.health_check, "__wrapped__") + assert hasattr(client._set_spend_logs_row_count_in_proxy_state, "__wrapped__") + + # Verify backoff configuration exists (methods should have retry behavior) + # This is implicit verification - the decorators are applied in the source code + + @pytest.mark.asyncio + async def test_multiple_connection_error_types( + self, mock_prisma_client, mock_proxy_logging + ): + """Test that different types of connection errors all trigger retries""" + error_types = [ + httpx.ConnectError("Connection error"), + httpx.TimeoutException("Timeout error"), + ClientNotConnectedError(), + PrismaError("Database error"), + ConnectionError("Network error"), + OSError("OS-level connection error"), + ] + + for error_type in error_types: + # Reset mock for each test + mock_prisma_client.db.query_raw.reset_mock() + mock_prisma_client.db.query_raw.side_effect = [ + error_type, # First call fails + [{"result": 1}], # Second call succeeds + ] + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + # Should succeed after retry + result = await client.health_check() + assert result == [{"result": 1}] + assert mock_prisma_client.db.query_raw.call_count == 2 + + @pytest.mark.asyncio + async def test_setup_prisma_client_integration( + self, mock_prisma_client, mock_proxy_logging + ): + """Test simulated _setup_prisma_client flow with intermittent failures""" + # This simulates the actual flow that happens in proxy_server.py _setup_prisma_client + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + # Simulate intermittent failures followed by success + call_count = 0 + + def mock_query_side_effect(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count <= 2: # First two calls fail + raise httpx.ConnectError("Intermittent connection issue") + else: + if "SELECT 1" in str(args): + return [{"result": 1}] + else: + return [{"reltuples": 1000}] + + mock_prisma_client.db.query_raw.side_effect = mock_query_side_effect + + with patch("litellm.proxy.proxy_server.proxy_state") as mock_proxy_state: + mock_proxy_state.set_proxy_state_variable = Mock() + + # Execute the two critical calls from _setup_prisma_client + health_result = await client.health_check() + await client._set_spend_logs_row_count_in_proxy_state() + + # Verify both operations eventually succeeded despite initial failures + assert health_result == [{"result": 1}] + mock_proxy_state.set_proxy_state_variable.assert_called_once() + + # Verify retries occurred (should have more than 2 total query calls) + assert mock_prisma_client.db.query_raw.call_count >= 4 + + @pytest.mark.asyncio + async def test_backoff_timing_constraints( + self, mock_prisma_client, mock_proxy_logging, connection_errors + ): + """Test that backoff respects max_time constraint (10 seconds)""" + # Mock all calls to fail to test max_time + mock_prisma_client.db.query_raw.side_effect = connection_errors[0] + + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db = mock_prisma_client.db + client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj + + start_time = time.time() + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + end_time = time.time() + duration = end_time - start_time + + # Should not exceed max_time of 10 seconds significantly + # Adding small buffer for test execution overhead + assert duration < 12.0 + + # Should have attempted max_tries (3) retries + assert mock_prisma_client.db.query_raw.call_count == 3 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])