fix(proxy): avoid DB outage during planned RDS IAM rotation#34749
fix(proxy): avoid DB outage during planned RDS IAM rotation#34749mubashir1osmani wants to merge 2 commits into
Conversation
Greptile SummaryThis PR changes planned RDS IAM token rotation to warm-connect and swap Prisma clients before retiring the prior engine.
Confidence Score: 2/5This PR should not merge until planned rotation preserves active database operations and handles cancellation without leaving inconsistent client state or an old engine running. The replacement connects before the swap, but the old engine is terminated after an arbitrary delay regardless of active work, and cancellation can interrupt the multi-stage transition without rollback or retirement cleanup. Files Needing Attention: litellm/proxy/db/prisma_client.py
|
| Filename | Overview |
|---|---|
| litellm/proxy/db/prisma_client.py | Introduces connect-before-swap IAM rotation, but fixed-duration draining can terminate active operations and task cancellation can leave token/client state divergent or leak the old engine. |
| tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py | Adds focused ordering and rollback coverage, but does not exercise long-running operations or cancellation during replacement. |
Reviews (1): Last reviewed commit: "fix(proxy): warm rotate Prisma client fo..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if old_engine_pid > 0: | ||
| await asyncio.sleep(self.PLANNED_REPLACEMENT_DRAIN_SECONDS) | ||
| await self._kill_engine_process(old_engine_pid) |
There was a problem hiding this comment.
When a database operation remains in flight on the old Prisma client for longer than 0.5 seconds during rotation, this path unconditionally terminates its engine, causing a transport failure for the active query and a request-visible database error for callers without reconnect retry handling.
Knowledge Base Used: Database Schema and Proxy DB Access Layer
| except Exception: | ||
| if previous_db_url is None: | ||
| os.environ.pop(self._db_url_env_var, None) | ||
| else: | ||
| os.environ[self._db_url_env_var] = previous_db_url | ||
| raise |
There was a problem hiding this comment.
When shutdown cancels the refresh task during replacement, CancelledError bypasses this except Exception rollback. Cancellation before the swap leaves the environment URL updated while the old client remains active, and cancellation after the swap during the drain or kill leaves the old engine running after its watcher has moved to the replacement.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
1 participant
TLDR
Problem this solves:
How it solves it:
Relevant issues
Follow-up to #29176
Prior incident: https://docs.litellm.ai/blog/prisma-reconnect-blocking-incident
Linear ticket
Pre-Submission checklist
Screenshots / Proof of Fix
Live IAM-enabled RDS validation is pending. This PR is a draft until continuous database operations pass across at least three forced token rotations
Focused regression result at
564cb7f147: 70 passed across planned restart, IAM expiry, routing, and cold recovery suitesType
Bug Fix
Test
Changes
Planned IAM refresh now connects a replacement Prisma client while the current client continues serving. It swaps only after connect succeeds, allows a bounded drain, then retires the old query engine using the existing non-blocking SIGTERM and SIGKILL path
If replacement connect fails, the current client and previous database URL remain active. Broken-engine recovery keeps its existing kill-first ordering
Regression tests pin connect-before-swap, swap-before-kill, failure rollback, and repeated old-engine retirement
RCA
Notion RCA: https://app.notion.com/p/3aa43b8acdab8175bb5cf40a87cda81c
The latest completed Rust run reported 346 passed, 75 failed, and 28 skipped. Database-backed operations failed together around planned Prisma SIGTERM events with
All connection attempts failedGrafana showed the failure every 720 seconds, matching the 15-minute RDS IAM token lifetime minus the 3-minute refresh buffer. Liveness and readiness stayed healthy, and there was no corresponding memory or CPU pressure. A newly rolled gateway reproduced the issue at its first rotation, ruling out long gateway uptime as the cause
The proactive IAM refresh path reused the cold Prisma recovery primitive. That primitive intentionally terminates the current engine before constructing its replacement to avoid prisma-client-py's event-loop-blocking
disconnect()behavior. The ordering is correct for an already broken engine but creates an availability gap when the engine is healthy and the restart is plannedThe prior incident fix preserved gateway liveness by replacing blocking shutdown with direct process signals. This follow-up addresses database continuity by separating planned warm rotation from unplanned cold recovery
The test gap was continuous database traffic across a token boundary. Existing verification covered liveness, eventual reconnect, and refresh coordination but did not assert that the old client remained available until its replacement connected
Grafana evidence: https://berriai.grafana.net/d/mu5lrjt?from=1785113400000&to=1785121800000
Final Attestation