mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
5c213127e8
Azure Database for PostgreSQL Flexible Server takes a Microsoft Entra ID access token as the connection password, and those tokens last about an hour, so a proxy pointed at one dies shortly after boot unless something keeps minting fresh ones Set AZURE_POSTGRESQL_AUTH=True (or pass --azure_postgresql_auth) alongside DATABASE_HOST, DATABASE_USER, and DATABASE_NAME, and the proxy mints a token at startup, assembles the connection URL around it, and refreshes it in the background for as long as the process runs. That is the same shape IAM_TOKEN_DB_AUTH already had for AWS RDS, so the two now share one code path: a tagged union picks the minting strategy once, and the wrapper, the read replica, and the refresh loop all read the choice off it instead of each guessing from the environment. Setting both toggles is a startup error, in the chart as well as in Python The helm chart gets database.writer.useAzureEntraAuth and the matching reader knob next to the existing useIAMAuth Fixes #29661 Co-authored-by: David Balatoni <balcsida@gmail.com>
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
import os
|
|
from collections.abc import Generator
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
|
|
DB_ENV_KEYS = (
|
|
"IAM_TOKEN_DB_AUTH",
|
|
"AZURE_POSTGRESQL_AUTH",
|
|
"DATABASE_URL",
|
|
"DIRECT_URL",
|
|
"DATABASE_URL_READ_REPLICA",
|
|
"DATABASE_HOST",
|
|
"DATABASE_PORT",
|
|
"DATABASE_USER",
|
|
"DATABASE_USERNAME",
|
|
"DATABASE_NAME",
|
|
"DATABASE_SCHEMA",
|
|
"DATABASE_PASSWORD",
|
|
"DATABASE_HOST_READ_REPLICA",
|
|
"DATABASE_PORT_READ_REPLICA",
|
|
"DATABASE_USER_READ_REPLICA",
|
|
"DATABASE_USERNAME_READ_REPLICA",
|
|
"DATABASE_NAME_READ_REPLICA",
|
|
"DATABASE_SCHEMA_READ_REPLICA",
|
|
"DATABASE_PASSWORD_READ_REPLICA",
|
|
)
|
|
|
|
_db_env_snapshot_key = pytest.StashKey[dict[str, Optional[str]]]()
|
|
|
|
|
|
def _db_env_snapshot() -> dict[str, Optional[str]]:
|
|
return {key: os.environ.get(key) for key in DB_ENV_KEYS}
|
|
|
|
|
|
@pytest.hookimpl(wrapper=True)
|
|
def pytest_runtest_setup(item: pytest.Item) -> Generator[None, None, None]:
|
|
item.stash[_db_env_snapshot_key] = _db_env_snapshot()
|
|
return (yield)
|
|
|
|
|
|
@pytest.hookimpl(wrapper=True)
|
|
def pytest_runtest_teardown(item: pytest.Item, nextitem: Optional[pytest.Item]) -> Generator[None, None, None]:
|
|
result = yield
|
|
before = item.stash[_db_env_snapshot_key]
|
|
leaked = {key: value for key, value in _db_env_snapshot().items() if value != before[key]}
|
|
for key, original in before.items():
|
|
if original is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = original
|
|
assert not leaked, (
|
|
f"{item.nodeid} leaked DB env vars past monkeypatch teardown: {leaked}. "
|
|
"Product code under test writes DATABASE_URL(_READ_REPLICA) into os.environ as a side effect; "
|
|
"monkeypatch only restores keys it has a record for, so a value written to a previously unset "
|
|
"key survives the test and poisons every later test in this pytest-xdist worker process "
|
|
"(DB-backed e2e tests arm themselves on DATABASE_URL and then fail to connect). "
|
|
"Use the unset_database_url fixture (or monkeypatch.setenv) so restoration is registered."
|
|
)
|
|
return result
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_entra_token_provider_cache() -> Generator[None, None, None]:
|
|
"""The Entra provider factory is cached process-wide so one Azure credential serves
|
|
the whole proxy; that cache would otherwise carry one test's stub into the next."""
|
|
from litellm.proxy.db.token_auth import build_azure_entra_token_provider
|
|
|
|
build_azure_entra_token_provider.cache_clear()
|
|
yield
|
|
build_azure_entra_token_provider.cache_clear()
|
|
|
|
|
|
@pytest.fixture
|
|
def unset_database_url(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("DATABASE_URL", "about-to-be-unset")
|
|
monkeypatch.delenv("DATABASE_URL")
|