mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(proxy): keep operator-pinned max_idle_connection_lifetime URL values over the default
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI
parent
e4e09867f0
commit
e12a4243b6
2 changed files with 44 additions and 52 deletions
|
|
@ -7,8 +7,9 @@ import re
|
|||
import subprocess
|
||||
import sys
|
||||
import urllib.parse as urlparse
|
||||
from collections.abc import Iterable
|
||||
from collections.abc import Iterable, Mapping
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import click
|
||||
|
|
@ -64,7 +65,6 @@ def _build_db_connection_url_params(
|
|||
pool_timeout: float | None,
|
||||
connect_timeout: float | None = None,
|
||||
socket_timeout: float | None = None,
|
||||
idle_connection_lifetime: float | None = None,
|
||||
disable_prepared_statements: bool = False,
|
||||
extra_params: dict | None = None,
|
||||
) -> dict:
|
||||
|
|
@ -88,8 +88,6 @@ def _build_db_connection_url_params(
|
|||
params["connect_timeout"] = connect_timeout
|
||||
if socket_timeout is not None:
|
||||
params["socket_timeout"] = socket_timeout
|
||||
if idle_connection_lifetime is not None:
|
||||
params["max_idle_connection_lifetime"] = idle_connection_lifetime
|
||||
if disable_prepared_statements:
|
||||
params["pgbouncer"] = "true"
|
||||
if extra_params:
|
||||
|
|
@ -1085,7 +1083,9 @@ def run_server(
|
|||
db_connection_timeout: int | float | None = 60
|
||||
db_connect_timeout: int | float | None = None
|
||||
db_socket_timeout: int | float | None = None
|
||||
db_connection_idle_lifetime: int | float | None = ( # rebind-ok: overwritten from general_settings when a config is provided
|
||||
db_connection_idle_lifetime: (
|
||||
int | float | None
|
||||
) = ( # rebind-ok: overwritten from general_settings when a config is provided
|
||||
LiteLLMDatabaseConnectionPool.database_connection_idle_lifetime.value
|
||||
)
|
||||
db_disable_prepared_statements: bool = False
|
||||
|
|
@ -1190,7 +1190,7 @@ def run_server(
|
|||
db_connection_timeout = LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value
|
||||
db_connect_timeout = general_settings.get("database_connect_timeout")
|
||||
db_socket_timeout = general_settings.get("database_socket_timeout")
|
||||
db_connection_idle_lifetime = general_settings.get( # rebind-ok: default declared above for the no-config path
|
||||
db_connection_idle_lifetime = general_settings.get( # rebind-ok: default set for the no-config path
|
||||
"database_connection_idle_lifetime",
|
||||
LiteLLMDatabaseConnectionPool.database_connection_idle_lifetime.value,
|
||||
)
|
||||
|
|
@ -1261,10 +1261,18 @@ def run_server(
|
|||
pool_timeout=db_connection_timeout,
|
||||
connect_timeout=db_connect_timeout,
|
||||
socket_timeout=db_socket_timeout,
|
||||
idle_connection_lifetime=db_connection_idle_lifetime,
|
||||
disable_prepared_statements=db_disable_prepared_statements,
|
||||
extra_params=db_extra_connection_params,
|
||||
)
|
||||
# The idle lifetime is applied add-if-missing so a value the operator
|
||||
# pinned on the URL itself keeps winning over the built-in default.
|
||||
idle_lifetime_params: Final[Mapping[str, int | float]] = MappingProxyType(
|
||||
{
|
||||
"max_idle_connection_lifetime": lifetime
|
||||
for lifetime in (db_connection_idle_lifetime,)
|
||||
if lifetime is not None
|
||||
}
|
||||
)
|
||||
if os.getenv("DATABASE_URL", None) is not None:
|
||||
database_url = get_secret("DATABASE_URL", default_value=None)
|
||||
resolved_url: Final[str | None] = str(database_url) if database_url else None
|
||||
|
|
@ -1282,11 +1290,11 @@ def run_server(
|
|||
writer_url,
|
||||
connection_url_params,
|
||||
)
|
||||
os.environ["DATABASE_URL"] = modified_url
|
||||
os.environ["DATABASE_URL"] = add_missing_query_params(modified_url, idle_lifetime_params)
|
||||
if os.getenv("DIRECT_URL", None) is not None:
|
||||
database_url = os.getenv("DIRECT_URL")
|
||||
modified_url = append_query_params(database_url, connection_url_params)
|
||||
os.environ["DIRECT_URL"] = modified_url
|
||||
os.environ["DIRECT_URL"] = add_missing_query_params(modified_url, idle_lifetime_params)
|
||||
# The reader pool is a real pool against the same configured cap, so it
|
||||
# gets the allowlisted pool params. Schema-affecting ones, including any
|
||||
# the operator smuggled in through database_extra_connection_params, stay
|
||||
|
|
@ -1303,7 +1311,7 @@ def run_server(
|
|||
_with_query_value(read_replica_url, "options", reader_options)
|
||||
if reader_options
|
||||
else read_replica_url,
|
||||
reader_shareable_params(connection_url_params),
|
||||
reader_shareable_params(MappingProxyType({**idle_lifetime_params, **connection_url_params})),
|
||||
)
|
||||
subprocess.run(["prisma"], capture_output=True)
|
||||
is_prisma_runnable = True
|
||||
|
|
|
|||
|
|
@ -879,42 +879,25 @@ class TestProxyInitializationHelpers:
|
|||
assert appended_params["pgbouncer"] == "true"
|
||||
assert appended_params["statement_cache_size"] == 0
|
||||
|
||||
def test_build_db_connection_url_params_includes_idle_lifetime(self):
|
||||
from litellm.proxy.proxy_cli import _build_db_connection_url_params
|
||||
|
||||
params = _build_db_connection_url_params(
|
||||
connection_limit=10,
|
||||
pool_timeout=60,
|
||||
idle_connection_lifetime=60,
|
||||
)
|
||||
assert params["max_idle_connection_lifetime"] == 60
|
||||
|
||||
def test_build_db_connection_url_params_omits_none_idle_lifetime(self):
|
||||
from litellm.proxy.proxy_cli import _build_db_connection_url_params
|
||||
|
||||
params = _build_db_connection_url_params(
|
||||
connection_limit=10,
|
||||
pool_timeout=60,
|
||||
idle_connection_lifetime=None,
|
||||
)
|
||||
assert "max_idle_connection_lifetime" not in params
|
||||
|
||||
def test_build_db_connection_url_params_extra_overrides_idle_lifetime(self):
|
||||
from litellm.proxy.proxy_cli import _build_db_connection_url_params
|
||||
|
||||
params = _build_db_connection_url_params(
|
||||
connection_limit=10,
|
||||
pool_timeout=60,
|
||||
idle_connection_lifetime=60,
|
||||
extra_params={"max_idle_connection_lifetime": 300},
|
||||
)
|
||||
assert params["max_idle_connection_lifetime"] == 300
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"general_settings, expected_idle_lifetime",
|
||||
"general_settings, database_url, expected_idle_lifetime",
|
||||
[
|
||||
({}, 60),
|
||||
({"database_connection_idle_lifetime": 30}, 30),
|
||||
({}, "postgresql://test:test@localhost:5432/test", "60"),
|
||||
(
|
||||
{"database_connection_idle_lifetime": 30},
|
||||
"postgresql://test:test@localhost:5432/test",
|
||||
"30",
|
||||
),
|
||||
(
|
||||
{"database_connection_idle_lifetime": None},
|
||||
"postgresql://test:test@localhost:5432/test",
|
||||
None,
|
||||
),
|
||||
(
|
||||
{},
|
||||
"postgresql://test:test@localhost:5432/test?max_idle_connection_lifetime=300",
|
||||
"300",
|
||||
),
|
||||
],
|
||||
)
|
||||
@patch("subprocess.run")
|
||||
|
|
@ -930,6 +913,7 @@ class TestProxyInitializationHelpers:
|
|||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
general_settings,
|
||||
database_url,
|
||||
expected_idle_lifetime,
|
||||
):
|
||||
from click.testing import CliRunner
|
||||
|
|
@ -948,7 +932,7 @@ class TestProxyInitializationHelpers:
|
|||
mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock(
|
||||
return_value={
|
||||
"general_settings": {
|
||||
"database_url": "postgresql://test:test@localhost:5432/test",
|
||||
"database_url": database_url,
|
||||
**general_settings,
|
||||
}
|
||||
}
|
||||
|
|
@ -972,10 +956,6 @@ class TestProxyInitializationHelpers:
|
|||
patch( # test-quality-ok: keeps the boot test from binding a real port, same as sibling boot tests
|
||||
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
|
||||
) as mock_get_args,
|
||||
patch( # test-quality-ok: capture point for the assembled URL params, same as sibling boot tests
|
||||
"litellm.proxy.proxy_cli.append_query_params",
|
||||
side_effect=lambda url, params: str(url),
|
||||
) as mock_append_query_params,
|
||||
):
|
||||
mock_get_args.return_value = {
|
||||
"app": "litellm.proxy.proxy_server:app",
|
||||
|
|
@ -991,9 +971,13 @@ class TestProxyInitializationHelpers:
|
|||
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["max_idle_connection_lifetime"] == expected_idle_lifetime
|
||||
final_query = dict(
|
||||
urlparse.parse_qsl(urlparse.urlparse(os.environ["DATABASE_URL"]).query)
|
||||
)
|
||||
if expected_idle_lifetime is None:
|
||||
assert "max_idle_connection_lifetime" not in final_query
|
||||
else:
|
||||
assert final_query["max_idle_connection_lifetime"] == expected_idle_lifetime
|
||||
|
||||
def test_build_db_connection_url_params_disable_prepared_statements(self):
|
||||
from litellm.proxy.proxy_cli import _build_db_connection_url_params
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue