fix(proxy): honor DATABASE_DISABLE_PREPARED_STATEMENTS in componentized entrypoints (#38363)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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.

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

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