diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c549f48126e..59cca71a5ce 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2425,6 +2425,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): database_connection_timeout: float | None = Field( 60, description="default timeout for a connection to the database" ) + database_connection_idle_lifetime: float | None = Field( + 60, + description=( + "Prisma `max_idle_connection_lifetime` URL param (seconds). Connections " + "idle longer than this are closed by the pool before a managed database " + "(RDS, Cloud SQL, Azure) silently drops them, preventing intermittent " + "`Error { kind: Closed }` failures. Set to null to fall back to " + "Prisma's built-in default (300s)." + ), + ) database_connect_timeout: float | None = Field( None, description=( diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 1a39016b3a3..d9ad36cd27d 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -82,6 +82,7 @@ CONNECTION_PARAM_KEYS: Final[frozenset[str]] = frozenset( "pool_timeout", "connect_timeout", "socket_timeout", + "max_idle_connection_lifetime", "pgbouncer", } ) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 8ac63ba25c9..d42bf9715e7 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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 @@ -56,6 +57,7 @@ telemetry: Final = None class LiteLLMDatabaseConnectionPool(Enum): database_connection_pool_limit = 10 database_connection_pool_timeout = 60 + database_connection_idle_lifetime = 60 def _build_db_connection_url_params( @@ -1081,6 +1083,11 @@ 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 + LiteLLMDatabaseConnectionPool.database_connection_idle_lifetime.value + ) db_disable_prepared_statements: bool = False db_extra_connection_params: dict | None = None db_statement_timeout: float | None = None @@ -1183,6 +1190,10 @@ 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 set for the no-config path + "database_connection_idle_lifetime", + LiteLLMDatabaseConnectionPool.database_connection_idle_lifetime.value, + ) _disable_prepared_statements: Final = general_settings.get("database_disable_prepared_statements", False) if isinstance(_disable_prepared_statements, str): from litellm.secret_managers.main import str_to_bool @@ -1253,6 +1264,15 @@ def run_server( 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 @@ -1270,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 @@ -1291,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 diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index 2552e52fb77..870504a008e 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -739,3 +739,12 @@ def test_unsupported_db_scheme_message_names_var_and_scheme(): assert "DIRECT_URL" in msg assert "sqlite" in msg assert "postgresql://" in msg + + +def test_reader_shareable_params_includes_idle_lifetime(): + from litellm.proxy.db.db_url_settings import reader_shareable_params + + shared = reader_shareable_params( + {"max_idle_connection_lifetime": 60, "schema": "other", "connection_limit": 10} + ) + assert shared == {"max_idle_connection_lifetime": 60, "connection_limit": 10} diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 6ea6f208bb5..0ad9049b228 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -879,6 +879,106 @@ class TestProxyInitializationHelpers: assert appended_params["pgbouncer"] == "true" assert appended_params["statement_cache_size"] == 0 + @pytest.mark.parametrize( + "general_settings, database_url, expected_idle_lifetime", + [ + ({}, "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") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: CLI boot test cannot run real prisma migrations, same as sibling boot tests + @patch( # test-quality-ok: CLI boot test cannot run real prisma migrations, same as sibling boot tests + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_db_connection_idle_lifetime_forwarded_to_url( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + general_settings, + database_url, + expected_idle_lifetime, + ): + 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": database_url, + **general_settings, + } + } + ) + + 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( # 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, + ): + 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}" + 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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e944062e15e..525cba154dd 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25411,6 +25411,12 @@ export interface components { * @description Prisma `connect_timeout` URL param (seconds). Bounds how long the engine waits to establish a new connection before failing. Defaults to Prisma's built-in value when unset. */ database_connect_timeout?: number | null; + /** + * Database Connection Idle Lifetime + * @description Prisma `max_idle_connection_lifetime` URL param (seconds). Connections idle longer than this are closed by the pool before a managed database (RDS, Cloud SQL, Azure) silently drops them, preventing intermittent `Error { kind: Closed }` failures. Set to null to fall back to Prisma's built-in default (300s). + * @default 60 + */ + database_connection_idle_lifetime: number | null; /** * Database Connection Pool Limit * @description default connection pool for prisma client connecting to postgres db