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

This commit is contained in:
harish876 2026-04-11 00:37:10 +00:00
parent 5f49f29f4e
commit 2ea0e53ec8
2 changed files with 81 additions and 3 deletions

View file

@ -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()

View file

@ -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"""