Skip to content

feat(rds): transparent stale-connection reconnect in RdsServiceGateway - #152

Merged
Kydoimos97 merged 1 commit into
releasefrom
fix/rds-gateway-reconnect
Aug 4, 2026
Merged

feat(rds): transparent stale-connection reconnect in RdsServiceGateway#152
Kydoimos97 merged 1 commit into
releasefrom
fix/rds-gateway-reconnect

Conversation

@Kydoimos97

Copy link
Copy Markdown
Contributor

New Features

  • RdsServiceGateway now recovers a stale database connection automatically. In single-connection mode the gateway opened one psycopg2 connection at init and reused it for the life of the process with no liveness check or reconnect. A long-running consumer (e.g. a pipeline whose connection sat idle across a long compute phase) would have that connection dropped server-side, and the next get_data/update_database failed with connection already closed. Callers now get transparent recovery: on a connection-level error the gateway reconnects and retries the operation once.

Fixes

  • Dead-connection rollback no longer masks the real error. The except blocks called conn.rollback(), which itself raises on a dead socket and hid the original failure. Those rollbacks are now guarded against OperationalError/InterfaceError.

Behavior / safety

  • Only psycopg2.OperationalError/InterfaceError (or messages like "connection already closed", "server closed the connection", "ssl connection has been closed", "connection not open", "terminating connection") trigger the reconnect+retry. Every other psycopg2.Error (IntegrityError, DataError, ProgrammingError, …) keeps its exact prior behavior — no retry.
  • The happy path and the raise_on_error semantics are unchanged for non-connection errors.
  • Retry is duplicate-safe: every update_database path performs all work then commits once at the very end, so a connection-level failure means nothing was committed and re-running the operation cannot double-write.
  • reconnect() rebuilds the single connection (or the pool) from the existing db_uri; it is a no-op in test mode.

Testing

  • tests/test_connect.py: 160 passed, ruff clean. New tests cover reconnect (single mode + test-mode no-op), connection-error classification, get_data reconnect-and-retry-once, and IntegrityError NOT triggering a reconnect for both get_data and update_database.

Context

  • Root cause of prod incident IR-1916 (ai-pipeline discarded a fully-trained model when its long-held RDS connection went stale before the final write). ai-pipeline shipped an app-layer guard already; this puts the fix in the shared library so every WrenchCL consumer inherits it. Ship as 6.3.0.

End-user impact

No user-facing surface. Operationally: services using RdsServiceGateway across long-lived processes stop failing on transient/idle connection drops.

RdsServiceGateway (single-connection mode) held one psycopg2 connection for
the life of the process with no liveness check, so a long-running consumer
whose connection was dropped server-side would fail its next query with
'connection already closed' — and the except-block rollback would itself throw
on the dead socket, masking the original error.

Adds reconnect() (rebuilds the single connection or the pool from db_uri) and
classifies connection-level errors; get_data and update_database now reconnect
and retry once on such an error only. Non-connection errors (IntegrityError,
DataError, etc.) keep their exact prior behavior, and rollback calls are
guarded so a dead-connection rollback can't mask the cause. update_database
commits once at the end of each path, so a whole-operation retry cannot
double-write.
@Kydoimos97
Kydoimos97 requested review from a team August 4, 2026 00:58
@Kydoimos97
Kydoimos97 merged commit 9bba3c7 into release Aug 4, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant