From c11c654b8efaa60bded2c6a57267f57587d60bbe Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:39:04 -0700 Subject: [PATCH] fix(proxy): honor DATABASE_DISABLE_PREPARED_STATEMENTS in componentized entrypoints (#38363) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_url_settings.py | 17 ++++ litellm/proxy/db/token_auth.py | 2 +- .../proxy/db/test_db_url_settings.py | 78 +++++++++++++++++++ 3 files changed, 96 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 0918b9039da..1a39016b3a3 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -60,6 +60,11 @@ AzureTokenAuthFlag = Annotated[ bool, BeforeValidator(partial(token_auth_flag_enabled, env_var=AZURE_POSTGRESQL_AUTH_ENV_VAR)) ] +DISABLE_PREPARED_STATEMENTS_ENV_VAR: Final = "DATABASE_DISABLE_PREPARED_STATEMENTS" +DisablePreparedStatementsFlag = Annotated[ + bool, BeforeValidator(partial(token_auth_flag_enabled, env_var=DISABLE_PREPARED_STATEMENTS_ENV_VAR)) +] + # schema.prisma pins `provider = "postgresql"`, so these are the only schemes # Prisma can actually connect with. SUPPORTED_DB_SCHEMES: Final[frozenset[str]] = frozenset({"postgresql", "postgres"}) @@ -153,6 +158,9 @@ class DatabaseURLSettings(BaseSettings): iam_token_db_auth: IamTokenAuthFlag = Field(default=False, validation_alias=IAM_TOKEN_DB_AUTH_ENV_VAR) azure_postgresql_auth: AzureTokenAuthFlag = Field(default=False, validation_alias=AZURE_POSTGRESQL_AUTH_ENV_VAR) + disable_prepared_statements: DisablePreparedStatementsFlag = Field( + default=False, validation_alias=DISABLE_PREPARED_STATEMENTS_ENV_VAR + ) # Writer database_url: str | None = Field(default=None, validation_alias="DATABASE_URL") @@ -375,6 +383,15 @@ class DatabaseURLSettings(BaseSettings): self._raise_for_unsupported_scheme() wrote_writer: Final = self.apply_writer_url_to_env() + # DATABASE_DISABLE_PREPARED_STATEMENTS maps to Prisma's `pgbouncer=true` + # URL param, same as the CLI's `database_disable_prepared_statements` + # config key. An explicit `pgbouncer` value already on the URL wins. + if self.disable_prepared_statements: + for env_var in ("DATABASE_URL", "DIRECT_URL"): + url = os.environ.get(env_var) + if url: + os.environ[env_var] = add_missing_query_params(url, MappingProxyType({"pgbouncer": "true"})) + # The reader inherits the writer's connection params (pool size, timeouts, # pgbouncer mode). Without this the reader pool ignores the configured cap # and falls back to Prisma's `num_physical_cpus * 2 + 1` default. diff --git a/litellm/proxy/db/token_auth.py b/litellm/proxy/db/token_auth.py index e1f84d1c04c..32c83c4f404 100644 --- a/litellm/proxy/db/token_auth.py +++ b/litellm/proxy/db/token_auth.py @@ -62,7 +62,7 @@ def token_auth_flag_enabled(value: str | bool | None, *, env_var: str) -> bool: return False raise ValueError( f"{env_var}={value!r} is not a recognized boolean. Set it to one of " - f"{', '.join(sorted(TRUTHY_TOKEN_AUTH_VALUES))} to turn token auth on, or to one of " + f"{', '.join(sorted(TRUTHY_TOKEN_AUTH_VALUES))} to turn it on, or to one of " f"{', '.join(sorted(v for v in FALSY_TOKEN_AUTH_VALUES if v))} to turn it off." ) 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 0ceec49de12..2552e52fb77 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -34,6 +34,7 @@ def _apply() -> bool: _MANAGED_DB_ENV_VARS = ( "IAM_TOKEN_DB_AUTH", "AZURE_POSTGRESQL_AUTH", + "DATABASE_DISABLE_PREPARED_STATEMENTS", "DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA", @@ -656,6 +657,83 @@ def test_reader_url_left_alone_when_writer_has_no_params(monkeypatch): ) +# --------------------------------------------------------------------------- +# DATABASE_DISABLE_PREPARED_STATEMENTS +# --------------------------------------------------------------------------- + + +def test_disable_prepared_statements_appends_pgbouncer_to_assembled_writer(monkeypatch): + monkeypatch.setenv("DATABASE_DISABLE_PREPARED_STATEMENTS", "true") + monkeypatch.setenv("DATABASE_HOST", "writer.example.com") + monkeypatch.setenv("DATABASE_USER", "litellm") + monkeypatch.setenv("DATABASE_NAME", "litellm_db") + monkeypatch.setenv("DATABASE_PASSWORD", "s3cr3t") + + assert _apply() is True + assert os.environ["DATABASE_URL"] == ( + "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?pgbouncer=true" + ) + assert "DIRECT_URL" not in os.environ + + +def test_disable_prepared_statements_appends_pgbouncer_to_pinned_writer(monkeypatch): + """The componentized entrypoints (gateway / backend / migrations) receive a + pinned DATABASE_URL and call apply_to_env; without the pgbouncer param Prisma + keeps named prepared statements and 42P05 collisions surface behind a + transaction-pooling pgbouncer.""" + monkeypatch.setenv("DATABASE_DISABLE_PREPARED_STATEMENTS", "true") + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db") + + assert _apply() is False + assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=true" + + +def test_disable_prepared_statements_respects_a_pinned_pgbouncer_value(monkeypatch): + monkeypatch.setenv("DATABASE_DISABLE_PREPARED_STATEMENTS", "true") + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=false") + + _apply() + + assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=false" + + +def test_disable_prepared_statements_applies_to_direct_url(monkeypatch): + monkeypatch.setenv("DATABASE_DISABLE_PREPARED_STATEMENTS", "true") + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db") + monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@direct.example.com:5432/litellm_db") + + _apply() + + assert os.environ["DIRECT_URL"] == "postgresql://u:p@direct.example.com:5432/litellm_db?pgbouncer=true" + + +def test_reader_inherits_pgbouncer_from_disable_prepared_statements(monkeypatch): + monkeypatch.setenv("DATABASE_DISABLE_PREPARED_STATEMENTS", "true") + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db") + + _apply() + + query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query) + assert query["pgbouncer"] == ["true"] + + +def test_disable_prepared_statements_off_leaves_urls_alone(monkeypatch): + monkeypatch.setenv("DATABASE_DISABLE_PREPARED_STATEMENTS", "false") + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db") + + _apply() + + assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db" + + +def test_disable_prepared_statements_rejects_an_unreadable_value(monkeypatch): + monkeypatch.setenv("DATABASE_DISABLE_PREPARED_STATEMENTS", "enabled") + + with pytest.raises(ValidationError, match="DATABASE_DISABLE_PREPARED_STATEMENTS"): + DatabaseURLSettings.from_env() + + def test_unsupported_db_scheme_message_names_var_and_scheme(): msg = unsupported_db_scheme_message("DIRECT_URL", "sqlite") assert "DIRECT_URL" in msg