From 2ea0e53ec8bee35073d6b7d3419e80a42b602960 Mon Sep 17 00:00:00 2001 From: harish876 Date: Sat, 11 Apr 2026 00:37:10 +0000 Subject: [PATCH] Fix: Currently database_connection_timeout pulls value from the database_connection_pool_timeout param from the general settings which does not adhere to the settings in the doc. Add a fix to consider either database_connection_timeout or database_connection_pool_timeout --- litellm/proxy/proxy_cli.py | 7 +- tests/test_litellm/proxy/test_proxy_cli.py | 77 +++++++++++++++++++++- 2 files changed, 81 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c638e294268..046d8ac88f4 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -803,8 +803,11 @@ def run_server( # noqa: PLR0915 LiteLLMDatabaseConnectionPool.database_connection_pool_limit.value, ) db_connection_timeout = general_settings.get( - "database_connection_pool_timeout", - LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value, + "database_connection_timeout", + general_settings.get( + "database_connection_pool_timeout", + LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value, + ), ) if database_url and database_url.startswith("os.environ/"): original_dir = os.getcwd() diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 349fe76ed71..d891c2a9efe 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1,6 +1,6 @@ import os import sys -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import fastapi import pytest @@ -420,6 +420,81 @@ class TestProxyInitializationHelpers: call_args = mock_uvicorn_run.call_args assert call_args[1]["limit_max_requests"] == 123 + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_database_connection_timeout_key_is_used_for_pool_timeout( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( + return_value={ + "general_settings": { + "database_url": "postgresql://test:test@localhost:5432/test", + "database_connection_pool_limit": 5, + "database_connection_timeout": 30, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + + with patch.dict( + os.environ, clean_env, clear=True + ), patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, patch( + "litellm.proxy.proxy_cli.append_query_params", + side_effect=lambda url, params: ( + f"{url}?connection_limit={params['connection_limit']}&pool_timeout={params['pool_timeout']}" + ), + ) as mock_append_query_params: + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--config", "test-config.yaml", "--skip_server_startup"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_append_query_params.assert_called() + appended_params = mock_append_query_params.call_args.args[1] + assert appended_params["connection_limit"] == 5 + assert appended_params["pool_timeout"] == 30 + @patch.dict(os.environ, {}, clear=True) def test_construct_database_url_from_env_vars(self): """Test the construct_database_url_from_env_vars function with various scenarios"""