Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/db_url

1"""Assemble DATABASE_URL (+ optional DATABASE_URL_READ_REPLICA) from env.

3The CLI (`proxy_cli.py`) assembles ``DATABASE_URL`` from discrete

4``DATABASE_*`` env vars before Prisma initializes. The componentized

5entrypoints (gateway / backend / migrations) bypass the CLI by uvicorn'ing

6the app directly, so they call ``DatabaseURLSettings.from_env().apply_to_env()``

7to do the same thing before importing ``proxy_server``.

9The env var names this module reads are exactly the ones emitted by the

10``helm/litellm`` chart's ``litellm.serverEnv`` block

11(``helm/litellm/templates/_helpers.tpl``). Both auth styles and both

12endpoints are covered:

14 * Token auth (``IAM_TOKEN_DB_AUTH`` truthy for AWS RDS IAM, or

15 ``AZURE_POSTGRESQL_AUTH`` truthy for Azure Database for PostgreSQL with

16 Microsoft Entra ID): mint a short-lived token and embed it as the

17 password. The writer URL is always (re)written because the token is

18 freshly minted on every startup. The chart omits ``DATABASE_PASSWORD``

19 in this mode. Enabling both toggles is a startup error.

20 * Password auth: build a percent-encoded URL from ``DATABASE_PASSWORD``.

21 The chart emits the discrete ``DATABASE_*`` fields (never a

22 pre-assembled URL), so URL-reserved characters in the password survive

23 instead of corrupting the URL. A pre-existing ``DATABASE_URL`` — e.g.

24 one an operator pinned via ``extraEnv`` — is left untouched and wins.

26The read replica is opt-in via ``DATABASE_HOST_READ_REPLICA`` and never

27clobbers a pre-existing ``DATABASE_URL_READ_REPLICA``, so a token-auth writer

28can run alongside a password-auth reader (or a precomputed reader URL). Reader

29token auth is gated on the same global toggle as the writer: the chart only

30emits the reader token env vars when the writer also uses token auth.

31Reader-side fields fall back to the writer's user / name / schema / port /

32password when their ``*_READ_REPLICA`` counterpart is unset, and to the

33writer's connection params (pool size, timeouts, pgbouncer mode) for the

34ones the reader URL does not pin itself.

35"""

37import _ssl

38import hashlib

39import os

40import socket

41import ssl

42import struct

43import sys

44import tempfile

45import urllib.parse

46from collections.abc import Callable, Mapping, Sequence

47from functools import partial

48from pathlib import Path

49from types import MappingProxyType

50from typing import Annotated, Final, Protocol, TypeAlias, cast

52from pydantic import AliasChoices, BeforeValidator, Field

53from pydantic_settings import BaseSettings, SettingsConfigDict

55from litellm.proxy.db.pgbouncer import database_url_is_pooled

56from litellm.proxy.db.token_auth import (

57 AZURE_POSTGRESQL_AUTH_ENV_VAR,

58 DEFAULT_POSTGRES_PORT,

59 IAM_TOKEN_DB_AUTH_ENV_VAR,

60 DatabaseTokenAuth,

61 IAMEndpoint,

62 build_database_token_auth,

63 mint_database_token,

64 token_auth_flag_enabled,

65)

67IamTokenAuthFlag = Annotated[bool, BeforeValidator(partial(token_auth_flag_enabled, env_var=IAM_TOKEN_DB_AUTH_ENV_VAR))]

68AzureTokenAuthFlag = Annotated[

69 bool, BeforeValidator(partial(token_auth_flag_enabled, env_var=AZURE_POSTGRESQL_AUTH_ENV_VAR))

70]

72DISABLE_PREPARED_STATEMENTS_ENV_VAR: Final = "DATABASE_DISABLE_PREPARED_STATEMENTS"

73DisablePreparedStatementsFlag = Annotated[

74 bool, BeforeValidator(partial(token_auth_flag_enabled, env_var=DISABLE_PREPARED_STATEMENTS_ENV_VAR))

75]

76MAX_IDLE_CONNECTION_LIFETIME_ENV_VAR: Final = "DATABASE_MAX_IDLE_CONNECTION_LIFETIME"

77DATABASE_SSLMODE_ENV_VAR: Final = "DATABASE_SSLMODE"

78DATABASE_SSLROOTCERT_ENV_VAR: Final = "DATABASE_SSLROOTCERT"

80# schema.prisma pins `provider = "postgresql"`, so these are the only schemes

81# Prisma can actually connect with.

82SUPPORTED_DB_SCHEMES: Final[frozenset[str]] = frozenset({"postgresql", "postgres"})

83_MISSING_SCHEME: Final = "<missing scheme>"

86# An allowlist, deliberately not a denylist: only these pool and timeout params

87# follow the writer to the read replica, so nothing that decides which tables a

88# query resolves against (``schema``, or a ``search_path`` inside ``options``)

89# can ever repoint the reader. Without them the reader pool silently falls back

90# to Prisma's default size.

91CONNECTION_PARAM_KEYS: Final[frozenset[str]] = frozenset(

92 {

93 "connection_limit",

94 "pool_timeout",

95 "connect_timeout",

96 "socket_timeout",

97 "max_idle_connection_lifetime",

98 "pgbouncer",

99 }

100)

101

102# Quaint never tests pooled connections on checkout and keeps them idle for

103# 300s by default, past many infra idle timeouts, so dead sockets surface as

104# `Error { kind: Closed }`. 60s recycles them first; explicit values win.

105DEFAULT_MAX_IDLE_CONNECTION_LIFETIME: Final = 60

106IDLE_LIFETIME_DEFAULT_PARAMS: Final[Mapping[str, int]] = MappingProxyType(

107 {"max_idle_connection_lifetime": DEFAULT_MAX_IDLE_CONNECTION_LIFETIME}

108)

109

110

111def idle_lifetime_params(configured: float | None) -> Mapping[str, str | int | float]:

112 """The `max_idle_connection_lifetime` to add to URLs that do not pin one.

113

114 Applied via ``add_missing_query_params`` so a URL-pinned value always wins,

115 whether the operator configured `database_max_idle_connection_lifetime` or not.

116 """

117 if configured is None: 117 ↛ 119line 117 didn't jump to line 119 because the condition on line 117 was always true

118 return IDLE_LIFETIME_DEFAULT_PARAMS

119 return MappingProxyType({"max_idle_connection_lifetime": configured})

120

121

122def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) -> str:

123 """Return ``url`` with the ``params`` it does not already carry appended.

124

125 Params the operator pinned on the URL win, so a hand-tuned replica URL keeps

126 its values. Returns the URL untouched when there is nothing to add, leaving

127 its existing encoding alone.

128 """

129 parsed: Final = urllib.parse.urlsplit(url)

130 existing: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))

131 pinned: Final = frozenset(key for key, _ in existing)

132 additions: Final = tuple((key, str(value)) for key, value in params.items() if key not in pinned)

133 if not additions: 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true

134 return url

135 query: Final = urllib.parse.urlencode(existing + additions)

136 return urllib.parse.urlunsplit(parsed._replace(query=query))

137

138

139LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"})

140PRISMA_TLS_PARAM_KEYS: Final[frozenset[str]] = frozenset({"sslmode", "sslcert", "sslaccept"})

141PEM_CERT_HEADER: Final = b"-----BEGIN CERTIFICATE-----"

142PG_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103)

143TLS_PROBE_TIMEOUT_SECONDS: Final = 10.0

144

145RootCertResolver: TypeAlias = Callable[[str, str, int], str] # mutable-ok: Callable parameter syntax

146

147

148class _VerifiedChainSource(Protocol):

149 def get_verified_chain(self) -> Sequence[_ssl.Certificate] | None: ... 149 ↛ exitline 149 didn't return from function 'get_verified_chain' because

150

151

152def _verified_chain_der(tls: ssl.SSLSocket) -> tuple[bytes, ...]:

153 if sys.version_info >= (3, 13):

154 return tuple(tls.get_verified_chain())

155 legacy: Final = cast( # cast-ok: the stub omits _sslobj, the C object has get_verified_chain since 3.10

156 "_VerifiedChainSource | None",

157 tls._sslobj, # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # public API only from 3.13

158 )

159 chain: Final = () if legacy is None else legacy.get_verified_chain() or ()

160 return tuple(cert.public_bytes(_ssl.ENCODING_DER) for cert in chain)

161

162

163def _server_trust_anchor(cafile: str, host: str, port: int) -> bytes | None:

164 try:

165 context: Final = ssl.create_default_context(cafile=cafile)

166 with socket.create_connection((host, port), timeout=TLS_PROBE_TIMEOUT_SECONDS) as raw:

167 raw.sendall(PG_SSL_REQUEST)

168 if raw.recv(1) != b"S":

169 return None

170 with context.wrap_socket(raw, server_hostname=host) as tls:

171 chain: Final = _verified_chain_der(tls)

172 except (OSError, ValueError):

173 return None

174 return chain[-1] if chain else None

175

176

177def pin_bundle_root(cert_path: str, host: str, port: int) -> str:

178 """Reduce a multi-root CA bundle to the one root that verifies ``host``.

179

180 Prisma's ``sslcert`` loads a single PEM certificate (native-tls

181 ``Certificate::from_pem``), so pointing it at a bundle such as the AWS RDS

182 global bundle trusts only the first of its 108 regional roots and the

183 handshake fails with "unable to get local issuer certificate" for every

184 other region. A single-certificate file is returned as is. For a bundle,

185 one verifying handshake (chain and hostname, whole bundle as trust store)

186 identifies the trust anchor the server actually chains to, which is

187 written to a single-certificate file for Prisma. If the probe fails the

188 bundle path is returned unchanged, so Prisma fails closed exactly as

189 before rather than trusting anything the bundle would not.

190 """

191 try:

192 if Path(cert_path).read_bytes().count(PEM_CERT_HEADER) < 2:

193 return cert_path

194 except OSError:

195 return cert_path

196 root: Final = _server_trust_anchor(cert_path, host, port)

197 if root is None:

198 return cert_path

199 pinned: Final = Path(tempfile.gettempdir()) / f"litellm-sslcert-{hashlib.sha256(root).hexdigest()[:16]}.pem"

200 return str(pinned) if _replace_file(pinned, ssl.DER_cert_to_PEM_cert(root)) else cert_path

201

202

203def _replace_file(target: Path, content: str) -> bool:

204 """Write ``content`` to a private temp file and rename it over ``target``, so

205 readers never see a partial file and a symlink planted at ``target`` is

206 replaced rather than followed."""

207 try:

208 fd, staged = tempfile.mkstemp(dir=target.parent, prefix=f"{target.name}.")

209 except OSError:

210 return False

211 try:

212 with os.fdopen(fd, "w") as handle:

213 handle.write(content)

214 os.replace(staged, target)

215 except OSError:

216 Path(staged).unlink(missing_ok=True)

217 return False

218 return True

219

220

221def translate_libpq_ssl_params(url: str, resolve_root_cert: RootCertResolver = pin_bundle_root) -> str:

222 """Rewrite libpq's certificate-verification params into Prisma's dialect.

223

224 Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert``

225 (a single CA certificate) and ``sslaccept=strict``. It silently discards

226 ``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to

227 ``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no

228 certificate check at all. ``verify-ca`` and ``verify-full`` both become

229 ``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes

230 ``sslcert`` (run through ``resolve_root_cert``, which pins a multi-root

231 bundle down to the server's root), and either one turns on

232 ``sslaccept=strict`` (chain and hostname), matching libpq where a root

233 cert makes ``require`` verify. Prisma params the operator pinned

234 themselves win; anything else is left untouched.

235 """

236 parsed: Final = urllib.parse.urlsplit(url)

237 pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))

238 keys: Final = frozenset(key for key, _ in pairs)

239 wants_verify: Final = any(key == "sslmode" and value in LIBPQ_VERIFY_SSLMODES for key, value in pairs)

240 if not wants_verify and "sslrootcert" not in keys: 240 ↛ 242line 240 didn't jump to line 242 because the condition on line 240 was always true

241 return url

242 translated: Final = tuple(

243 ("sslmode", "require") if key == "sslmode" and value in LIBPQ_VERIFY_SSLMODES else (key, value)

244 for key, value in pairs

245 if key != "sslrootcert"

246 )

247 root_cert: Final = tuple(

248 ("sslcert", resolve_root_cert(value, parsed.hostname or "", parsed.port or int(DEFAULT_POSTGRES_PORT)))

249 for key, value in pairs

250 if key == "sslrootcert" and "sslcert" not in keys

251 )

252 strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),)

253 query: Final = urllib.parse.urlencode(translated + root_cert + strict)

254 return urllib.parse.urlunsplit(parsed._replace(query=query))

255

256

257def reader_shareable_params(params: Mapping[str, str | int | float]) -> Mapping[str, str | int | float]:

258 """Return the subset of ``params`` the read replica is allowed to inherit."""

259 return MappingProxyType({key: value for key, value in params.items() if key in CONNECTION_PARAM_KEYS})

260

261

262def connection_params_from_url(url: str) -> Mapping[str, str | int | float]:

263 """Return the connection params on ``url`` that the read replica shares."""

264 return reader_shareable_params(

265 MappingProxyType({key: value for key, value in urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query)})

266 )

267

268

269def token_refresh_params_from_url(url: str) -> Mapping[str, str | int | float]:

270 """Return the params a re-minted token URL carries over from the URL it replaces.

271

272 The pool and timeout params plus Prisma's TLS params (already translated from

273 libpq spelling), so a refreshed URL keeps verifying the server the way the

274 first one did.

275 """

276 kept: Final = CONNECTION_PARAM_KEYS | PRISMA_TLS_PARAM_KEYS

277 return MappingProxyType(

278 {key: value for key, value in urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query) if key in kept}

279 )

280

281

282def unsupported_db_scheme(database_url: str) -> str | None:

283 """Return the connection URL scheme when it is not PostgreSQL, else None.

284

285 A `sqlite://` / `mysql://` URL can never connect against the

286 postgresql-only datasource, but the resulting Prisma failure is opaque and

287 version-dependent (a confusing migration error, or a startup that never

288 binds). Callers use this to reject the URL up front with an actionable

289 error instead.

290

291 A schemeless value (e.g. a malformed DSN like ``user:pass@host/db``) yields

292 the ``_MISSING_SCHEME`` placeholder rather than the raw URL, so callers that

293 log the return value never echo embedded credentials.

294 """

295 scheme: Final = urllib.parse.urlsplit(database_url).scheme.lower()

296 if scheme in SUPPORTED_DB_SCHEMES: 296 ↛ 298line 296 didn't jump to line 298 because the condition on line 296 was always true

297 return None

298 return scheme or _MISSING_SCHEME

299

300

301def unsupported_db_scheme_message(env_var: str, scheme: str) -> str:

302 """Operator-facing message naming the offending env var and scheme."""

303 return (

304 f"{env_var} uses unsupported scheme '{scheme}'. LiteLLM's database "

305 "features (virtual keys, store_model_in_db, spend tracking) require "

306 "PostgreSQL; use a 'postgresql://' connection string. SQLite and other "

307 "engines are not supported. "

308 "See https://docs.litellm.ai/docs/proxy/virtual_keys"

309 )

310

311

312class DatabaseURLSettings(BaseSettings):

313 """Discrete ``DATABASE_*`` env vars, loaded once at process start.

314

315 Field names are internal; ``validation_alias`` pins each one to the exact

316 env var the helm chart emits. ``DATABASE_USER`` doubles as

317 ``DATABASE_USERNAME`` for parity with ``construct_database_url_from_env_vars``.

318 """

319

320 model_config = SettingsConfigDict(case_sensitive=False, extra="ignore")

321

322 iam_token_db_auth: IamTokenAuthFlag = Field(default=False, validation_alias=IAM_TOKEN_DB_AUTH_ENV_VAR)

323 azure_postgresql_auth: AzureTokenAuthFlag = Field(default=False, validation_alias=AZURE_POSTGRESQL_AUTH_ENV_VAR)

324 disable_prepared_statements: DisablePreparedStatementsFlag = Field(

325 default=False, validation_alias=DISABLE_PREPARED_STATEMENTS_ENV_VAR

326 )

327 max_idle_connection_lifetime: int | None = Field(

328 default=None, validation_alias=MAX_IDLE_CONNECTION_LIFETIME_ENV_VAR

329 )

330

331 database_sslmode: str | None = Field(default=None, validation_alias=DATABASE_SSLMODE_ENV_VAR)

332 database_sslrootcert: str | None = Field(default=None, validation_alias=DATABASE_SSLROOTCERT_ENV_VAR)

333

334 # Writer

335 database_url: str | None = Field(default=None, validation_alias="DATABASE_URL")

336 direct_url: str | None = Field(default=None, validation_alias="DIRECT_URL")

337 database_host: str | None = Field(default=None, validation_alias="DATABASE_HOST")

338 database_port: str = Field(default=DEFAULT_POSTGRES_PORT, validation_alias="DATABASE_PORT")

339 database_user: str | None = Field(

340 default=None,

341 validation_alias=AliasChoices("DATABASE_USER", "DATABASE_USERNAME"),

342 )

343 database_name: str | None = Field(default=None, validation_alias="DATABASE_NAME")

344 database_schema: str | None = Field(default=None, validation_alias="DATABASE_SCHEMA")

345 database_password: str | None = Field(default=None, validation_alias="DATABASE_PASSWORD")

346

347 # Read replica

348 database_url_read_replica: str | None = Field(default=None, validation_alias="DATABASE_URL_READ_REPLICA")

349 database_host_read_replica: str | None = Field(default=None, validation_alias="DATABASE_HOST_READ_REPLICA")

350 database_port_read_replica: str | None = Field(default=None, validation_alias="DATABASE_PORT_READ_REPLICA")

351 database_user_read_replica: str | None = Field(

352 default=None,

353 validation_alias=AliasChoices("DATABASE_USER_READ_REPLICA", "DATABASE_USERNAME_READ_REPLICA"),

354 )

355 database_name_read_replica: str | None = Field(default=None, validation_alias="DATABASE_NAME_READ_REPLICA")

356 database_schema_read_replica: str | None = Field(default=None, validation_alias="DATABASE_SCHEMA_READ_REPLICA")

357 database_password_read_replica: str | None = Field(default=None, validation_alias="DATABASE_PASSWORD_READ_REPLICA")

358

359 @classmethod

360 def from_env(cls) -> "DatabaseURLSettings":

361 """Load the settings from ``os.environ`` (read at call time)."""

362 return cls()

363

364 def token_auth(self) -> DatabaseTokenAuth | None:

365 """The token strategy the toggles ask for, or ``None`` for password auth.

366

367 Raises ``RuntimeError`` when both toggles are on, since the password can only

368 come from one source.

369 """

370 return build_database_token_auth(

371 iam_token_db_auth=self.iam_token_db_auth,

372 azure_postgresql_auth=self.azure_postgresql_auth,

373 )

374

375 def tls_params(self) -> Mapping[str, str]:

376 """``sslmode`` / ``sslrootcert`` query params for every URL assembled from the discrete vars.

377

378 A root cert on its own means ``verify-full``: under libpq's default

379 ``prefer`` the CA would never be consulted, and PgBouncer would dial

380 Postgres unverified with the bundle loaded.

381 """

382 sslmode: Final = self.database_sslmode or ("verify-full" if self.database_sslrootcert else None)

383 return MappingProxyType(

384 {

385 key: value

386 for key, value in (

387 ("sslmode", sslmode),

388 ("sslrootcert", self.database_sslrootcert),

389 )

390 if value

391 }

392 )

393

394 def build_writer_url(self) -> str | None:

395 """Return the writer URL to set, or ``None`` to leave it as-is.

396

397 Raises ``RuntimeError`` (naming the offending vars) when token auth is

398 enabled but a required field is missing — the proxy cannot recover

399 from this and a clear startup error beats a Prisma connect failure.

400 A ``DATABASE_URL`` the supervisor pointed at the in-container PgBouncer

401 is kept even under token auth: the pooler renews the token upstream.

402 """

403 assembled: Final = self._assemble_writer_url()

404 if assembled is None:

405 return None

406 return add_missing_query_params(assembled, self.tls_params())

407

408 def _assemble_writer_url(self) -> str | None:

409 auth: Final = self.token_auth()

410 if auth is not None and database_url_is_pooled():

411 return None

412 if auth is not None:

413 missing: Final = tuple(

414 env

415 for env, val in (

416 ("DATABASE_HOST", self.database_host),

417 ("DATABASE_USER", self.database_user),

418 ("DATABASE_NAME", self.database_name),

419 )

420 if not val

421 )

422 if missing:

423 raise RuntimeError(

424 f"{auth.env_var} is enabled but required DB env var(s) "

425 f"are unset: {', '.join(missing)}. Set them so the writer "

426 f"DATABASE_URL can be assembled with a minted {auth.label}."

427 )

428 endpoint: Final = IAMEndpoint(

429 host=cast(str, self.database_host),

430 port=self.database_port,

431 user=cast(str, self.database_user),

432 name=cast(str, self.database_name),

433 schema=self.database_schema,

434 )

435 return endpoint.build_url(mint_database_token(auth, endpoint))

436

437 # Password auth: an operator-pinned DATABASE_URL always wins.

438 if self.database_url:

439 return None

440 if self.database_host and self.database_user and self.database_name:

441 return self._password_url(

442 user=self.database_user,

443 password=self.database_password,

444 host=self.database_host,

445 port=self.database_port,

446 name=self.database_name,

447 schema=self.database_schema,

448 )

449 return None

450

451 def build_reader_url(self) -> str | None:

452 """Return the read-replica URL to set, or ``None`` to leave it as-is.

453

454 Opt-in via ``DATABASE_HOST_READ_REPLICA``; never clobbers a

455 pre-existing ``DATABASE_URL_READ_REPLICA``. Reader fields fall back

456 to the writer's values.

457 """

458 assembled: Final = self._assemble_reader_url()

459 if assembled is None:

460 return None

461 return add_missing_query_params(assembled, self.tls_params())

462

463 def _assemble_reader_url(self) -> str | None:

464 if not self.database_host_read_replica:

465 return None # reader is opt-in

466 if self.database_url_read_replica:

467 return None # never clobber an operator-supplied reader URL

468

469 host: Final = self.database_host_read_replica

470 port: Final = self.database_port_read_replica or self.database_port

471 user: Final = self.database_user_read_replica or self.database_user

472 name: Final = self.database_name_read_replica or self.database_name

473 schema: Final = self.database_schema_read_replica or self.database_schema

474 password: Final = self.database_password_read_replica or self.database_password

475

476 auth: Final = self.token_auth()

477 if auth is not None:

478 missing: Final = tuple(

479 env

480 for env, val in (

481 ("DATABASE_USER[_READ_REPLICA]", user),

482 ("DATABASE_NAME[_READ_REPLICA]", name),

483 )

484 if not val

485 )

486 if missing:

487 raise RuntimeError(

488 f"{auth.env_var} is enabled and DATABASE_HOST_READ_REPLICA "

489 "is set, but the reader could not resolve: "

490 f"{', '.join(missing)} (no *_READ_REPLICA value and no "

491 "writer fallback). Set the reader fields or the writer "

492 "defaults."

493 )

494 endpoint: Final = IAMEndpoint(

495 host=host,

496 port=port,

497 user=cast(str, user),

498 name=cast(str, name),

499 schema=schema,

500 )

501 return endpoint.build_url(mint_database_token(auth, endpoint))

502

503 if user and name:

504 return self._password_url(

505 user=user,

506 password=password,

507 host=host,

508 port=port,

509 name=name,

510 schema=schema,

511 )

512 return None

513

514 @staticmethod

515 def _password_url(

516 *,

517 user: str,

518 password: str | None,

519 host: str,

520 port: str,

521 name: str,

522 schema: str | None,

523 ) -> str:

524 """Percent-encode credentials into a ``postgresql://`` URL.

525

526 Parity with ``construct_database_url_from_env_vars`` in

527 ``proxy/utils.py``; ``password`` may be empty for a passwordless URL.

528 """

529 quote: Final = urllib.parse.quote_plus

530 user_p: Final = quote(user)

531 name_p: Final = quote(name)

532 if password:

533 url = f"postgresql://{user_p}:{quote(password)}@{host}:{port}/{name_p}"

534 else:

535 url = f"postgresql://{user_p}@{host}:{port}/{name_p}"

536 if schema:

537 url += f"?schema={schema}"

538 return url

539

540 def _raise_for_unsupported_scheme(self) -> None:

541 """Reject an operator-pinned non-PostgreSQL writer / direct / reader URL.

542

543 The componentized entrypoints (gateway / backend / migrations) call

544 ``apply_to_env`` and then hand the URL straight to Prisma, bypassing

545 the CLI's own guard. A pinned URL flows through untouched, so validate

546 the same three vars the CLI guard checks (DATABASE_URL, DIRECT_URL, and

547 the read replica) rather than letting Prisma stall on an unusable scheme.

548 """

549 for env_var, url in (

550 ("DATABASE_URL", self.database_url),

551 ("DIRECT_URL", self.direct_url),

552 ("DATABASE_URL_READ_REPLICA", self.database_url_read_replica),

553 ):

554 if not url:

555 continue

556 bad_scheme = unsupported_db_scheme(url)

557 if bad_scheme is not None:

558 raise RuntimeError(unsupported_db_scheme_message(env_var, bad_scheme))

559

560 def apply_writer_url_to_env(self) -> bool:

561 """Write just the assembled writer URL into ``os.environ``.

562

563 Split out because the CLI shares this minting path but resolves the read

564 replica separately, so it must not pick up reader behavior on the way. The

565 CLI runs its own scheme guard over the pinned URLs, so unlike

566 ``apply_to_env`` this does not repeat it.

567 """

568 writer_url: Final = self.build_writer_url()

569 if writer_url is None:

570 return False

571 os.environ["DATABASE_URL"] = writer_url

572 # Normalize the toggles so downstream readers (PrismaWrapper's token

573 # refresh) reliably see token auth on, regardless of spelling.

574 if self.iam_token_db_auth:

575 os.environ[IAM_TOKEN_DB_AUTH_ENV_VAR] = "True"

576 if self.azure_postgresql_auth:

577 os.environ[AZURE_POSTGRESQL_AUTH_ENV_VAR] = "True"

578 return True

579

580 def apply_to_env(self) -> bool:

581 """Write the assembled URL(s) into ``os.environ``.

582

583 Returns True iff this call set ``DATABASE_URL`` (token mint, or

584 password auth that assembled a fresh URL). False means there was

585 nothing to do — an operator-pinned URL, or no discrete fields.

586 """

587 self._raise_for_unsupported_scheme()

588 wrote_writer: Final = self.apply_writer_url_to_env()

589

590 for env_var in ("DATABASE_URL", "DIRECT_URL"):

591 url = os.environ.get(env_var)

592 if url:

593 os.environ[env_var] = translate_libpq_ssl_params(url)

594

595 # DATABASE_DISABLE_PREPARED_STATEMENTS maps to Prisma's `pgbouncer=true`

596 # URL param, same as the CLI's `database_disable_prepared_statements`

597 # config key. An explicit `pgbouncer` value already on the URL wins.

598 if self.disable_prepared_statements:

599 for env_var in ("DATABASE_URL", "DIRECT_URL"):

600 url = os.environ.get(env_var)

601 if url:

602 os.environ[env_var] = add_missing_query_params(url, MappingProxyType({"pgbouncer": "true"}))

603

604 lifetime_params: Final = idle_lifetime_params(self.max_idle_connection_lifetime)

605 for env_var in ("DATABASE_URL", "DIRECT_URL"):

606 url = os.environ.get(env_var)

607 if url:

608 os.environ[env_var] = add_missing_query_params(url, lifetime_params)

609

610 # The reader inherits the writer's connection params (pool size, timeouts,

611 # pgbouncer mode). Without this the reader pool ignores the configured cap

612 # and falls back to Prisma's `num_physical_cpus * 2 + 1` default.

613 reader_url: Final = self.build_reader_url() or self.database_url_read_replica

614 if reader_url is not None:

615 os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params(

616 translate_libpq_ssl_params(reader_url),

617 connection_params_from_url(os.environ.get("DATABASE_URL", "")),

618 )

619

620 return wrote_writer

Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/db_url_settings.py: 34%

245 statements