From 12dffbafd6dce04a0ea81c93b35ff83f871027d6 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 16:10:49 -0500 Subject: [PATCH] fix(proxy): honor DATABASE_DISABLE_PREPARED_STATEMENTS in the litellm CLI (#42556) * fix(proxy): honor DATABASE_DISABLE_PREPARED_STATEMENTS in the litellm CLI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): pooled pgbouncer url keeps a single pgbouncer=true when the upstream already carries it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): reject a malformed DATABASE_DISABLE_PREPARED_STATEMENTS even when the config already enables it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_cli.py | 7 +- tests/test_litellm/proxy/db/test_pgbouncer.py | 7 + tests/test_litellm/proxy/test_proxy_cli.py | 166 ++++++++++++++++++ 3 files changed, 179 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 78885461724..40140974198 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1283,6 +1283,7 @@ def run_server( if os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None: from litellm.proxy.db.db_url_settings import ( + DISABLE_PREPARED_STATEMENTS_ENV_VAR, add_missing_query_params, idle_lifetime_params, reader_shareable_params, @@ -1305,12 +1306,16 @@ def run_server( sys.exit(1) from litellm.secret_managers.main import get_secret + env_disable_prepared_statements: Final = token_auth_flag_enabled( + os.getenv(DISABLE_PREPARED_STATEMENTS_ENV_VAR), env_var=DISABLE_PREPARED_STATEMENTS_ENV_VAR + ) + disable_prepared_statements: Final = db_disable_prepared_statements or env_disable_prepared_statements connection_url_params: Final = _build_db_connection_url_params( connection_limit=db_connection_pool_limit, pool_timeout=db_connection_timeout, connect_timeout=db_connect_timeout, socket_timeout=db_socket_timeout, - disable_prepared_statements=db_disable_prepared_statements, + disable_prepared_statements=disable_prepared_statements, extra_params=db_extra_connection_params, ) lifetime_params: Final = idle_lifetime_params(general_settings.get("database_max_idle_connection_lifetime")) diff --git a/tests/test_litellm/proxy/db/test_pgbouncer.py b/tests/test_litellm/proxy/db/test_pgbouncer.py index bf7df3077ea..c69c8d015a3 100644 --- a/tests/test_litellm/proxy/db/test_pgbouncer.py +++ b/tests/test_litellm/proxy/db/test_pgbouncer.py @@ -122,6 +122,13 @@ class TestPlanPgBouncer: "pgbouncer": "true", } + def test_an_upstream_that_already_disables_prepared_statements_gets_a_single_pgbouncer_flag(self): + pooled: Final = _plan("postgresql://app:pw@db/litellm?connection_limit=5&pgbouncer=true").pooled_url + assert urllib.parse.parse_qsl(urllib.parse.urlsplit(pooled).query) == [ + ("connection_limit", "5"), + ("pgbouncer", "true"), + ] + @pytest.mark.parametrize("hop_param", ["channel_binding=require", "gssencmode=require"]) def test_transport_params_for_the_postgres_hop_stay_off_the_plain_tcp_loopback_url(self, hop_param: str): pooled: Final = _plan(f"postgresql://app:pw@db/litellm?connection_limit=5&{hop_param}").pooled_url diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index a38470d1fdf..d2d71d6df05 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1174,6 +1174,172 @@ class TestProxyInitializationHelpers: else: assert "pgbouncer" not in appended_params + @pytest.mark.parametrize( + "env_value, config_value, expect_pgbouncer", + [ + ("true", None, True), + ("1", None, True), + ("false", None, False), + (None, None, False), + ("true", False, True), + ("false", True, True), + ], + ) + @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_disable_prepared_statements_env_var_forwarded_to_url( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + env_value, + config_value, + expect_pgbouncer, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + general_settings = {"database_url": "postgresql://test:test@localhost:5432/test"} + if config_value is not None: + general_settings["database_disable_prepared_statements"] = config_value + 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": general_settings} + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL", "DATABASE_DISABLE_PREPARED_STATEMENTS") + } + if env_value is not None: + clean_env["DATABASE_DISABLE_PREPARED_STATEMENTS"] = env_value + + 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: str(url), + ) 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}" + appended_params = mock_append_query_params.call_args.args[1] + if expect_pgbouncer: + assert appended_params["pgbouncer"] == "true", appended_params + else: + assert "pgbouncer" not in appended_params, appended_params + + @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_malformed_disable_prepared_statements_env_var_is_rejected_even_when_config_enables_it( + 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_disable_prepared_statements": True, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL", "DATABASE_DISABLE_PREPARED_STATEMENTS") + } + clean_env["DATABASE_DISABLE_PREPARED_STATEMENTS"] = "enabled" + + 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: str(url), + ) 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 isinstance(result.exception, ValueError), f"exit_code={result.exit_code}, output={result.output}" + assert "DATABASE_DISABLE_PREPARED_STATEMENTS" in str(result.exception), result.exception + mock_append_query_params.assert_not_called() + @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")