scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(proxy): honor DATABASE_DISABLE_PREPARED_STATEMENTS in the litellm CLI (#42556)

Browse source 
* 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 <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-22 16:10:49 -05:00 • committed by GitHub
parent bbc702cd8b
commit 12dffbafd6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 179 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/proxy_cli.py
View file

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

7
tests/test_litellm/proxy/db/test_pgbouncer.py
View file

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

166
tests/test_litellm/proxy/test_proxy_cli.py
View file

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