scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

fix: proactive RDS IAM token refresh to prevent 15-min connection failed (#18795)

Browse source 
* fix: proactive RDS IAM token refresh to prevent 15-min connection failures (#16220)

* fix: add noqa for PLR0915 in proxy_startup_event
This commit is contained in:
Harshit Jain 2026-01-08 23:53:36 +05:30 • committed by GitHub
parent 3ebec39b74
commit 516e4f8b96
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 602 additions and 79 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

309
litellm/proxy/db/prisma_client.py
View file

@ -17,50 +17,141 @@ from litellm.secret_managers.main import str_to_bool
class PrismaWrapper:
"""
Wrapper around Prisma client that handles RDS IAM token authentication.
When iam_token_db_auth is enabled, this wrapper:
1. Proactively refreshes IAM tokens before they expire (background task)
2. Falls back to synchronous refresh if a token is found expired
3. Uses proper locking to prevent race conditions during reconnection
RDS IAM tokens are valid for 15 minutes. This wrapper refreshes them
3 minutes before expiration to ensure uninterrupted database connectivity.
"""
# Buffer time in seconds before token expiration to trigger refresh
# Refresh 3 minutes (180 seconds) before the token expires
TOKEN_REFRESH_BUFFER_SECONDS = 180
# Fallback refresh interval if token parsing fails (10 minutes)
FALLBACK_REFRESH_INTERVAL_SECONDS = 600
def __init__(self, original_prisma: Any, iam_token_db_auth: bool):
self._original_prisma = original_prisma
self.iam_token_db_auth = iam_token_db_auth
# Background token refresh task management
self._token_refresh_task: Optional[asyncio.Task] = None
self._reconnection_lock = asyncio.Lock()
self._last_refresh_time: Optional[datetime] = None
def _extract_token_from_db_url(self, db_url: Optional[str]) -> Optional[str]:
"""
Extract the token (password) from the DATABASE_URL.
The token contains the AWS signature with X-Amz-Date and X-Amz-Expires parameters.
Important: We must parse the URL while it's still encoded to preserve structure,
then decode the password portion. Otherwise the '?' in the token breaks URL parsing.
"""
if db_url is None:
return None
try:
# Parse URL while still encoded to preserve structure
parsed = urllib.parse.urlparse(db_url)
if parsed.password:
# Now decode just the password/token
return urllib.parse.unquote(parsed.password)
return None
except Exception:
return None
def _parse_token_expiration(self, token: Optional[str]) -> Optional[datetime]:
"""
Parse the token to extract its expiration time.
Returns the datetime when the token expires, or None if parsing fails.
"""
if token is None:
return None
try:
# Token format: ...?X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=900&...
if "?" not in token:
return None
query_string = token.split("?", 1)[1]
params = urllib.parse.parse_qs(query_string)
expires_str = params.get("X-Amz-Expires", [None])[0]
date_str = params.get("X-Amz-Date", [None])[0]
if not expires_str or not date_str:
return None
token_created = datetime.strptime(date_str, "%Y%m%dT%H%M%SZ")
expires_in = int(expires_str)
return token_created + timedelta(seconds=expires_in)
except Exception as e:
verbose_proxy_logger.debug(f"Failed to parse token expiration: {e}")
return None
def _calculate_seconds_until_refresh(self) -> float:
"""
Calculate exactly how many seconds until we need to refresh the token.
Uses precise timing: sleeps until (token_expiration - buffer_seconds).
For a 15-minute (900s) token with 180s buffer, this returns ~720s (12 min).
Returns:
Number of seconds to sleep before the next refresh.
Returns 0 if token should be refreshed immediately.
Returns FALLBACK_REFRESH_INTERVAL_SECONDS if parsing fails.
"""
db_url = os.getenv("DATABASE_URL")
token = self._extract_token_from_db_url(db_url)
expiration_time = self._parse_token_expiration(token)
if expiration_time is None:
# If we can't parse the token, use fallback interval
verbose_proxy_logger.debug(
f"Could not parse token expiration, using fallback interval of "
f"{self.FALLBACK_REFRESH_INTERVAL_SECONDS}s"
)
return self.FALLBACK_REFRESH_INTERVAL_SECONDS
# Calculate when we should refresh (expiration - buffer)
refresh_at = expiration_time - timedelta(
seconds=self.TOKEN_REFRESH_BUFFER_SECONDS
)
# How long until refresh time?
now = datetime.utcnow()
seconds_until_refresh = (refresh_at - now).total_seconds()
# If already past refresh time, return 0 (refresh immediately)
return max(0, seconds_until_refresh)
def is_token_expired(self, token_url: Optional[str]) -> bool:
"""Check if the token in the given URL is expired."""
if token_url is None:
return True
# Decode the token URL to handle URL-encoded characters
decoded_url = urllib.parse.unquote(token_url)
# Parse the token URL
parsed_url = urllib.parse.urlparse(decoded_url)
token = self._extract_token_from_db_url(token_url)
expiration_time = self._parse_token_expiration(token)
# Parse the query parameters from the path component (if they exist there)
query_params = urllib.parse.parse_qs(parsed_url.query)
if expiration_time is None:
# If we can't parse the token, assume it's expired to trigger refresh
verbose_proxy_logger.debug(
"Could not parse token expiration, treating as expired"
)
return True
# Get expiration time from the query parameters
expires = query_params.get("X-Amz-Expires", [None])[0]
if expires is None:
raise ValueError("X-Amz-Expires parameter is missing or invalid.")
expires_int = int(expires)
# Get the token's creation time from the X-Amz-Date parameter
token_time_str = query_params.get("X-Amz-Date", [""])[0]
if not token_time_str:
raise ValueError("X-Amz-Date parameter is missing or invalid.")
# Ensure the token time string is parsed correctly
try:
token_time = datetime.strptime(token_time_str, "%Y%m%dT%H%M%SZ")
except ValueError as e:
raise ValueError(f"Invalid X-Amz-Date format: {e}")
# Calculate the expiration time
expiration_time = token_time + timedelta(seconds=expires_int)
# Current time in UTC
current_time = datetime.utcnow()
# Check if the token is expired
return current_time > expiration_time
return datetime.utcnow() > expiration_time
def get_rds_iam_token(self) -> Optional[str]:
"""Generate a new RDS IAM token and update DATABASE_URL."""
if self.iam_token_db_auth:
from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token
@ -74,7 +165,6 @@ class PrismaWrapper:
db_host=db_host, db_port=db_port, db_user=db_user
)
# print(f"token: {token}")
_db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}"
if db_schema:
_db_url += f"?schema={db_schema}"
@ -86,6 +176,7 @@ class PrismaWrapper:
async def recreate_prisma_client(
self, new_db_url: str, http_client: Optional[Any] = None
):
"""Disconnect and reconnect the Prisma client with a new database URL."""
from prisma import Prisma # type: ignore
try:
@ -100,21 +191,159 @@ class PrismaWrapper:
await self._original_prisma.connect()
async def start_token_refresh_task(self) -> None:
"""
Start the background token refresh task.
This task proactively refreshes RDS IAM tokens before they expire,
preventing connection failures. Should be called after the initial
Prisma client connection is established.
"""
if not self.iam_token_db_auth:
verbose_proxy_logger.debug(
"IAM token auth not enabled, skipping token refresh task"
)
return
if self._token_refresh_task is not None:
verbose_proxy_logger.debug("Token refresh task already running")
return
self._token_refresh_task = asyncio.create_task(self._token_refresh_loop())
verbose_proxy_logger.info(
"Started RDS IAM token proactive refresh background task"
)
async def stop_token_refresh_task(self) -> None:
"""
Stop the background token refresh task gracefully.
Should be called during application shutdown to clean up resources.
"""
if self._token_refresh_task is None:
return
self._token_refresh_task.cancel()
try:
await self._token_refresh_task
except asyncio.CancelledError:
pass
self._token_refresh_task = None
verbose_proxy_logger.info("Stopped RDS IAM token refresh background task")
async def _token_refresh_loop(self) -> None:
"""
Background loop that proactively refreshes RDS IAM tokens before expiration.
Uses precise timing: calculates the exact sleep duration until the token
needs to be refreshed (expiration - 3 minute buffer), then refreshes.
This is more efficient than polling, requiring only 1 wake-up per token cycle.
"""
verbose_proxy_logger.info(
f"RDS IAM token refresh loop started. "
f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration."
)
while True:
try:
# Calculate exactly how long to sleep until next refresh
sleep_seconds = self._calculate_seconds_until_refresh()
if sleep_seconds > 0:
verbose_proxy_logger.info(
f"RDS IAM token refresh scheduled in {sleep_seconds:.0f} seconds "
f"({sleep_seconds / 60:.1f} minutes)"
)
await asyncio.sleep(sleep_seconds)
# Refresh the token
verbose_proxy_logger.info("Proactively refreshing RDS IAM token...")
await self._safe_refresh_token()
except asyncio.CancelledError:
verbose_proxy_logger.info("RDS IAM token refresh loop cancelled")
break
except Exception as e:
verbose_proxy_logger.error(
f"Error in RDS IAM token refresh loop: {e}. "
f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..."
)
# On error, wait before retrying to avoid tight error loops
try:
await asyncio.sleep(self.FALLBACK_REFRESH_INTERVAL_SECONDS)
except asyncio.CancelledError:
break
async def _safe_refresh_token(self) -> None:
"""
Refresh the RDS IAM token with proper locking to prevent race conditions.
Uses an asyncio lock to ensure only one refresh operation happens at a time,
preventing multiple concurrent reconnection attempts.
"""
async with self._reconnection_lock:
new_db_url = self.get_rds_iam_token()
if new_db_url:
await self.recreate_prisma_client(new_db_url)
self._last_refresh_time = datetime.utcnow()
verbose_proxy_logger.info(
"RDS IAM token refreshed successfully. New token valid for ~15 minutes."
)
else:
verbose_proxy_logger.error(
"Failed to generate new RDS IAM token during proactive refresh"
)
def __getattr__(self, name: str):
"""
Proxy attribute access to the underlying Prisma client.
If IAM token auth is enabled and the token is expired, this method
provides a synchronous fallback to refresh the token. However, this
should rarely be needed since the background task proactively refreshes
tokens before they expire.
FIXED: Now properly waits for reconnection to complete before returning,
instead of the previous fire-and-forget pattern that caused the bug.
"""
original_attr = getattr(self._original_prisma, name)
if self.iam_token_db_auth:
db_url = os.getenv("DATABASE_URL")
if self.is_token_expired(db_url):
db_url = self.get_rds_iam_token()
loop = asyncio.get_event_loop()
if db_url:
# Check if token is expired (should be rare if background task is running)
if self.is_token_expired(db_url):
verbose_proxy_logger.warning(
"RDS IAM token expired in __getattr__ - proactive refresh may have failed. "
"Triggering synchronous fallback refresh..."
)
new_db_url = self.get_rds_iam_token()
if new_db_url:
loop = asyncio.get_event_loop()
if loop.is_running():
asyncio.run_coroutine_threadsafe(
self.recreate_prisma_client(db_url), loop
# FIXED: Actually wait for the reconnection to complete!
# The previous code used fire-and-forget which caused the bug.
future = asyncio.run_coroutine_threadsafe(
self.recreate_prisma_client(new_db_url), loop
)
try:
# Wait up to 30 seconds for reconnection
future.result(timeout=30)
verbose_proxy_logger.info(
"Synchronous token refresh completed successfully"
)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to refresh token synchronously: {e}"
)
raise
else:
asyncio.run(self.recreate_prisma_client(db_url))
asyncio.run(self.recreate_prisma_client(new_db_url))
# Get the NEW attribute after reconnection
original_attr = getattr(self._original_prisma, name)
else:
raise ValueError("Failed to get RDS IAM token")

97
litellm/proxy/proxy_server.py
View file

@ -533,9 +533,9 @@ except ImportError:
server_root_path = os.getenv("SERVER_ROOT_PATH", "")
_license_check = LicenseCheck()
premium_user: bool = _license_check.is_premium()
premium_user_data: Optional["EnterpriseLicenseData"] = (
_license_check.airgapped_license_data
)
premium_user_data: Optional[
"EnterpriseLicenseData"
] = _license_check.airgapped_license_data
global_max_parallel_request_retries_env: Optional[str] = os.getenv(
"LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
)
@ -658,7 +658,7 @@ async def _initialize_shared_aiohttp_session():
@asynccontextmanager
async def proxy_startup_event(app: FastAPI):
async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session
import json
@ -788,6 +788,17 @@ async def proxy_startup_event(app: FastAPI):
except Exception as e:
verbose_proxy_logger.error(f"Error closing shared aiohttp session: {e}")
# Shutdown event - stop RDS IAM token refresh background task
if (
prisma_client is not None
and hasattr(prisma_client, "db")
and hasattr(prisma_client.db, "stop_token_refresh_task")
):
try:
await prisma_client.db.stop_token_refresh_task()
except Exception as e:
verbose_proxy_logger.error(f"Error stopping token refresh task: {e}")
await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues]
@ -1083,9 +1094,7 @@ try:
# In non-root Docker, we restructure in /var/lib/litellm/ui.
try:
_restructure_ui_html_files(ui_path)
verbose_proxy_logger.info(
f"Restructured UI directory: {ui_path}"
)
verbose_proxy_logger.info(f"Restructured UI directory: {ui_path}")
except PermissionError as e:
verbose_proxy_logger.exception(
f"Permission error while restructuring UI directory {ui_path}: {e}"
@ -1171,9 +1180,9 @@ master_key: Optional[str] = None
config_agents: Optional[List[AgentConfig]] = None
otel_logging = False
prisma_client: Optional[PrismaClient] = None
shared_aiohttp_session: Optional["ClientSession"] = (
None # Global shared session for connection reuse
)
shared_aiohttp_session: Optional[
"ClientSession"
] = None # Global shared session for connection reuse
user_api_key_cache = DualCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
@ -1181,9 +1190,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
dual_cache=user_api_key_cache
)
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
redis_usage_cache: Optional[RedisCache] = (
None # redis cache used for tracking spend, tpm/rpm limits
)
redis_usage_cache: Optional[
RedisCache
] = None # redis cache used for tracking spend, tpm/rpm limits
polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False
polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache
user_custom_auth = None
@ -1522,9 +1531,9 @@ async def update_cache( # noqa: PLR0915
_id = "team_id:{}".format(team_id)
try:
# Fetch the existing cost for the given user
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
await user_api_key_cache.async_get_cache(key=_id)
)
existing_spend_obj: Optional[
LiteLLM_TeamTable
] = await user_api_key_cache.async_get_cache(key=_id)
if existing_spend_obj is None:
# do nothing if team not in api key cache
return
@ -1876,7 +1885,6 @@ class ProxyConfig:
"environment_variables" in config_to_save
and config_to_save["environment_variables"]
):
# decrypt the environment_variables - in case a caller function has already encrypted the environment_variables
decrypted_env_vars = self._decrypt_and_set_db_env_variables(
environment_variables=config_to_save["environment_variables"],
@ -2794,21 +2802,21 @@ class ProxyConfig:
verbose_proxy_logger.debug(f"_alerting_callbacks: {general_settings}")
if _alerting_callbacks is None:
return
# Ensure proxy_logging_obj.alerting is set for all alerting types
_alerting_value = general_settings.get("alerting", None)
verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}")
verbose_proxy_logger.debug(
f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}"
)
proxy_logging_obj.update_values(
alerting=_alerting_value,
alerting_threshold=general_settings.get("alerting_threshold", 600),
alert_types=general_settings.get("alert_types", None),
alert_to_webhook_url=general_settings.get(
"alert_to_webhook_url", None
),
alert_to_webhook_url=general_settings.get("alert_to_webhook_url", None),
alerting_args=general_settings.get("alerting_args", None),
redis_cache=redis_usage_cache,
)
for _alert in _alerting_callbacks:
if _alert == "slack":
# [OLD] v0 implementation - already handled by update_values above
@ -3279,7 +3287,7 @@ class ProxyConfig:
proxy_logging_obj: ProxyLogging
"""
_general_settings = config_data.get("general_settings", {})
if _general_settings is not None and "alerting" in _general_settings:
if (
general_settings is not None
@ -3294,7 +3302,8 @@ class ProxyConfig:
_merged_alerting = list(_yaml_alerting.union(_db_alerting))
# Preserve order: YAML values first, then DB values
_merged_alerting = list(general_settings["alerting"]) + [
item for item in _general_settings["alerting"]
item
for item in _general_settings["alerting"]
if item not in general_settings["alerting"]
]
verbose_proxy_logger.debug(
@ -3605,7 +3614,6 @@ class ProxyConfig:
await self._init_vector_stores_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="vector_store_indexes"):
await self._init_vector_store_indexes_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="mcp"):
@ -3804,10 +3812,10 @@ class ProxyConfig:
)
try:
guardrails_in_db: List[Guardrail] = (
await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
guardrails_in_db: List[
Guardrail
] = await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
verbose_proxy_logger.debug(
"guardrails from the DB %s", str(guardrails_in_db)
@ -4134,9 +4142,9 @@ async def initialize( # noqa: PLR0915
user_api_base = api_base
dynamic_config[user_model]["api_base"] = api_base
if api_version:
os.environ["AZURE_API_VERSION"] = (
api_version # set this for azure - litellm can read this from the env
)
os.environ[
"AZURE_API_VERSION"
] = api_version # set this for azure - litellm can read this from the env
if max_tokens: # model-specific param
dynamic_config[user_model]["max_tokens"] = max_tokens
if temperature: # model-specific param
@ -4654,10 +4662,14 @@ class ProxyStartupEvent:
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Responses cost check job scheduled successfully")
verbose_proxy_logger.info(
"Responses cost check job scheduled successfully"
)
except Exception as e:
verbose_proxy_logger.debug(f"Failed to setup responses cost checking: {e}")
verbose_proxy_logger.debug(
f"Failed to setup responses cost checking: {e}"
)
verbose_proxy_logger.debug(
"Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..."
)
@ -4826,6 +4838,14 @@ class ProxyStartupEvent:
await prisma_client.connect()
## Start RDS IAM token refresh background task if enabled ##
# This proactively refreshes IAM tokens before they expire,
# preventing the 15-minute connection failure bug (#16220)
if hasattr(prisma_client, "db") and hasattr(
prisma_client.db, "start_token_refresh_task"
):
await prisma_client.db.start_token_refresh_task()
## Add necessary views to proxy ##
asyncio.create_task(
prisma_client.check_view_exists()
@ -5937,7 +5957,6 @@ async def realtime_websocket_endpoint(
),
user_api_key_dict=Depends(user_api_key_auth_websocket),
):
await websocket.accept()
# Only use explicit parameters, not all query params
@ -9525,9 +9544,9 @@ async def get_config_list(
hasattr(sub_field_info, "description")
and sub_field_info.description is not None
):
nested_fields[idx].field_description = (
sub_field_info.description
)
nested_fields[
idx
].field_description = sub_field_info.description
idx += 1
_stored_in_db = None

275
tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py Normal file
View file

@ -0,0 +1,275 @@
"""
Tests for the RDS IAM token proactive refresh implementation.
Tests for GitHub Issue #16220: RDS IAM authentication connection failures after 15 minutes.
The fix implements:
1. Proactive background token refresh (refreshes 3 min before expiration)
2. Precise sleep timing (1 wake-up per token cycle instead of polling)
3. Proper locking during reconnection
4. Fixed __getattr__ fallback that now waits for reconnection
Run these tests:
poetry run pytest tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py -v -s
"""
import asyncio
import os
import urllib.parse
from datetime import datetime, timedelta
from unittest.mock import MagicMock, patch
import pytest
class TestPrismaWrapperTokenRefresh:
"""Tests for the PrismaWrapper RDS IAM token refresh implementation."""
@pytest.fixture
def setup_env(self):
"""Setup environment variables for testing."""
os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com"
os.environ["DATABASE_PORT"] = "5432"
os.environ["DATABASE_USER"] = "test_user"
os.environ["DATABASE_NAME"] = "test_db"
os.environ["IAM_TOKEN_DB_AUTH"] = "True"
yield
# Cleanup
for key in [
"DATABASE_HOST",
"DATABASE_PORT",
"DATABASE_USER",
"DATABASE_NAME",
"DATABASE_URL",
"IAM_TOKEN_DB_AUTH",
"DATABASE_SCHEMA",
]:
os.environ.pop(key, None)
def _generate_mock_token(self, expires_in_seconds: int = 900) -> str:
"""Generate a mock IAM token with expiration info."""
now = datetime.utcnow()
date_str = now.strftime("%Y%m%dT%H%M%SZ")
# Build the token like AWS does
token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires={expires_in_seconds}&X-Amz-Signature=abc123"
return urllib.parse.quote(token, safe="")
def _set_database_url_with_token(self, expires_in_seconds: int = 900):
"""Set DATABASE_URL with a mock token."""
token = self._generate_mock_token(expires_in_seconds)
os.environ[
"DATABASE_URL"
] = f"postgresql://test_user:{token}@test-host:5432/test_db"
@pytest.mark.asyncio
async def test_is_token_expired_fresh(self, setup_env):
"""Test that fresh token is not detected as expired."""
from litellm.proxy.db.prisma_client import PrismaWrapper
mock_prisma = MagicMock()
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
self._set_database_url_with_token(expires_in_seconds=900)
db_url = os.getenv("DATABASE_URL")
assert wrapper.is_token_expired(db_url) is False
@pytest.mark.asyncio
async def test_is_token_expired_old(self, setup_env):
"""Test that old token is detected as expired."""
from litellm.proxy.db.prisma_client import PrismaWrapper
mock_prisma = MagicMock()
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
# Create an expired token
old_date = datetime.utcnow() - timedelta(seconds=901)
date_str = old_date.strftime("%Y%m%dT%H%M%SZ")
token = (
f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=900&X-Amz-Signature=abc"
)
encoded_token = urllib.parse.quote(token, safe="")
db_url = f"postgresql://test_user:{encoded_token}@test-host:5432/test_db"
assert wrapper.is_token_expired(db_url) is True
@pytest.mark.asyncio
async def test_start_stop_token_refresh_task(self, setup_env):
"""Test that token refresh task starts and stops correctly."""
from litellm.proxy.db.prisma_client import PrismaWrapper
mock_prisma = MagicMock()
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
# Set a valid token
self._set_database_url_with_token(expires_in_seconds=900)
# Start the task
await wrapper.start_token_refresh_task()
assert wrapper._token_refresh_task is not None
assert not wrapper._token_refresh_task.done()
# Stop the task
await wrapper.stop_token_refresh_task()
assert wrapper._token_refresh_task is None
@pytest.mark.asyncio
async def test_start_task_not_enabled(self, setup_env):
"""Test that task doesn't start when IAM auth is not enabled."""
from litellm.proxy.db.prisma_client import PrismaWrapper
mock_prisma = MagicMock()
# IAM auth disabled
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)
await wrapper.start_token_refresh_task()
assert wrapper._token_refresh_task is None
@pytest.mark.asyncio
async def test_is_token_expired_null(self, setup_env):
"""Test that None token is treated as expired."""
from litellm.proxy.db.prisma_client import PrismaWrapper
mock_prisma = MagicMock()
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
assert wrapper.is_token_expired(None) is True
class TestTokenExpirationParsing:
"""Tests for token expiration parsing utilities."""
def test_parse_token_expiration_valid(self):
"""Test parsing expiration from a valid token."""
from litellm.proxy.db.prisma_client import PrismaWrapper
mock_prisma = MagicMock()
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
# Create a token with known expiration
token = "mock-token?X-Amz-Date=20240101T120000Z&X-Amz-Expires=900&X-Amz-Signature=abc"
expiration = wrapper._parse_token_expiration(token)
assert expiration is not None
expected = datetime(2024, 1, 1, 12, 0, 0) + timedelta(seconds=900)
assert expiration == expected
def test_parse_token_expiration_invalid(self):
"""Test that invalid token returns None."""
from litellm.proxy.db.prisma_client import PrismaWrapper
mock_prisma = MagicMock()
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
# Invalid tokens
assert wrapper._parse_token_expiration(None) is None
assert wrapper._parse_token_expiration("no-query-params") is None
assert wrapper._parse_token_expiration("?missing=params") is None
class TestBackgroundRefreshLoop:
"""Tests for the background refresh loop timing."""
@pytest.fixture
def setup_env(self):
"""Setup environment variables for testing."""
os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com"
os.environ["DATABASE_PORT"] = "5432"
os.environ["DATABASE_USER"] = "test_user"
os.environ["DATABASE_NAME"] = "test_db"
yield
# Cleanup
for key in [
"DATABASE_HOST",
"DATABASE_PORT",
"DATABASE_USER",
"DATABASE_NAME",
"DATABASE_URL",
]:
os.environ.pop(key, None)
@pytest.mark.asyncio
async def test_calculate_seconds_fallback_when_no_url(self, setup_env):
"""Test that fallback is used when DATABASE_URL is not set."""
from litellm.proxy.db.prisma_client import PrismaWrapper
mock_prisma = MagicMock()
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
# Don't set DATABASE_URL
seconds = wrapper._calculate_seconds_until_refresh()
# Should return fallback interval
assert seconds == wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS
# ============================================================================
# DEMONSTRATION SCRIPT
# ============================================================================
async def demonstrate_fix():
"""
Demonstrates the fix for the RDS IAM token expiration bug.
Shows how the proactive refresh prevents the 15-minute connection failure.
"""
# Import the actual implementation
try:
from litellm.proxy.db.prisma_client import PrismaWrapper
except ImportError:
return
# Setup mock environment
os.environ["DATABASE_HOST"] = "mock-rds.region.rds.amazonaws.com"
os.environ["DATABASE_PORT"] = "5432"
os.environ["DATABASE_USER"] = "iam_user"
os.environ["DATABASE_NAME"] = "litellm"
# Create initial token (expires in 10 seconds for demo)
now = datetime.utcnow()
date_str = now.strftime("%Y%m%dT%H%M%SZ")
token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=10&X-Amz-Signature=abc123"
encoded_token = urllib.parse.quote(token, safe="")
os.environ[
"DATABASE_URL"
] = f"postgresql://iam_user:{encoded_token}@mock-rds:5432/litellm"
# Create mock prisma client
mock_prisma = MagicMock()
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
# Override buffer for faster demo
wrapper.TOKEN_REFRESH_BUFFER_SECONDS = 3
wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS = 5
_ = wrapper._calculate_seconds_until_refresh() # Verify calculation works
db_url = os.getenv("DATABASE_URL")
is_expired = wrapper.is_token_expired(db_url)
assert is_expired is False, "Fresh token should not be expired!"
# Mock the _token_refresh_loop to prevent it from actually running
async def mock_loop():
try:
await asyncio.sleep(1000)
except asyncio.CancelledError:
pass
with patch.object(wrapper, "_token_refresh_loop", side_effect=mock_loop):
await wrapper.start_token_refresh_task()
await wrapper.stop_token_refresh_task()
# Cleanup
for key in [
"DATABASE_HOST",
"DATABASE_PORT",
"DATABASE_USER",
"DATABASE_NAME",
"DATABASE_URL",
]:
os.environ.pop(key, None)
if __name__ == "__main__":
asyncio.run(demonstrate_fix())