Skip to content

fix(drizzle): use resilient SQL proxy to survive reconnection #901

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
jhaynie merged 1 commit into from
Feb 9, 2026
Merged

fix(drizzle): use resilient SQL proxy to survive reconnection #901

jhaynie merged 1 commit into from
Feb 9, 2026
+438 −3

Conversation

@jhaynie
Copy link
Member

@jhaynie jhaynie commented Feb 9, 2026 •
edited by coderabbitai bot
Loading

Summary

Fixes #899 — ERR_POSTGRES_CONNECTION_CLOSED still happening even with the custom @agentuity/drizzle package.

Root Cause

createPostgresDrizzle() passed client.raw (a static reference to the Bun.SQL instance) to Drizzle at creation time:

const db = drizzle({ client: client.raw, ... });

When the connection dropped and PostgresClient._reconnectLoop() succeeded, _reinitializeSql() created a new this._sql Bun.SQL instance — but Drizzle's internal BunSQLSession still held the old, dead reference. All subsequent Drizzle queries used the dead connection → ERR_POSTGRES_CONNECTION_CLOSED.

Fix

  • Resilient SQL Proxy (packages/drizzle/src/postgres.ts): Replace the static client.raw with a Proxy that always resolves client.raw at call time. After reconnection, Drizzle automatically uses the new connection.
  • Retry-wrapped unsafe(): The proxy intercepts unsafe() calls and routes them through client.executeWithRetry(), providing automatic retry with exponential backoff on transient connection errors. Supports .values() chaining that Drizzle uses internally.
  • Public executeWithRetry() (packages/postgres/src/client.ts): Exposes the existing private retry logic as a public method for use by integration layers like the Drizzle proxy.

Changes

File Change
packages/drizzle/src/postgres.ts Added createResilientSQLProxy(), replaced client.raw with proxy
packages/postgres/src/client.ts Added public executeWithRetry() method, bound in createCallableClient()
packages/drizzle/test/proxy.test.ts 10 new tests — proxy delegation, reconnection resilience, .values() chaining
packages/postgres/test/client.test.ts 8 new tests — executeWithRetry behavior, retry logic, error handling

Verification

  • ✅ bun run format — clean
  • ✅ bun run lint — 0 errors
  • ✅ bun run typecheck — 0 errors
  • ✅ bun run build — clean
  • ✅ bun run test — 212 pass, 0 fail (postgres + drizzle packages)

Summary by CodeRabbit

  • New Features

    • SQL operations now automatically retry on transient failures, improving database reliability during network interruptions.
    • Enhanced connection handling ensures fresh database connections are used after reconnections, preventing stale connection issues.

  • Tests

    • Added comprehensive test coverage for connection resilience, reconnection handling, and automatic retry behavior.

@jhaynie
fix(drizzle): use resilient SQL proxy to survive reconnection
2cc91ef 
Drizzle was initialized with a static reference to client.raw (the Bun.SQL
instance). When the connection dropped and PostgresClient reconnected, it
created a new Bun.SQL instance, but Drizzle still held the old dead reference,
causing ERR_POSTGRES_CONNECTION_CLOSED on all subsequent queries.

Fix: Replace the static client.raw reference with a dynamic Proxy that always
resolves client.raw at call time. The proxy also wraps unsafe() calls through
executeWithRetry for automatic retry on transient connection errors.

Closes #899
@coderabbitai
Copy link

coderabbitai bot commented Feb 9, 2026 •
edited
Loading

📝 Walkthrough

Walkthrough

This PR introduces a resilient SQL proxy in Drizzle that dynamically delegates to the current raw Postgres connection, enabling proper reconnection handling. It also exposes a public executeWithRetry method on PostgresClient, allowing external retry logic with automatic reconnection-aware execution.

Changes

Cohort / File(s) Summary
Postgres Client Retry API
packages/postgres/src/client.ts, packages/postgres/test/client.test.ts		 Added public executeWithRetry(operation, maxRetries?) method that wraps internal retry logic and is exposed on the callable client. Tests cover basic execution, promise/non-promise returns, retryable/non-retryable error handling, recovery, and closed-client behavior.
Drizzle Resilient SQL Proxy
packages/drizzle/src/postgres.ts, packages/drizzle/test/proxy.test.ts		 Introduced createResilientSQLProxy that delegates to the client's current raw connection, ensuring dynamic reconnection awareness. Wraps unsafe() with retry logic and returns a thenable supporting .values() for Bun compatibility. Tests validate delegation, retry wrapping, .values() chaining, reconnection handling, property propagation, and method binding.
🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check ✅ Passed The PR directly addresses the root cause identified in issue #899 by implementing a resilient SQL proxy that dynamically resolves the current Bun.SQL connection at call time, ensuring Drizzle uses the live connection after reconnection and wrapping calls with retry logic.
Out of Scope Changes check ✅ Passed All changes align with the stated objectives: adding the resilient proxy, exposing executeWithRetry, and comprehensive tests. No unrelated modifications or scope creep detected.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

No actionable comments were generated in the recent review. 🎉

🧹 Recent nitpick comments
packages/postgres/test/client.test.ts (1)

142-151: Consider verifying the specific error type.

The test confirms that operations reject after close, but it doesn't verify that the error is specifically a ConnectionClosedError. This would make the test more robust against regressions.

💡 Optional: Verify specific error type 
 	it('throws after close', async () => {
 		const client = createTestClient();
 		await client.close();
 
-		await expect(
-			client.executeWithRetry(async () => {
-				return 'should not reach';
-			})
-		).rejects.toThrow();
+		await expect(
+			client.executeWithRetry(async () => {
+				return 'should not reach';
+			})
+		).rejects.toThrow('Client has been closed');
 	});
packages/drizzle/test/proxy.test.ts (1)

184-209: Consider removing unused variable for clarity.

The variable _beginFn on line 189 is assigned but never used. The test correctly demonstrates the behavior by accessing proxy.begin again after swapping raw, making this variable unnecessary.

🧹 Remove unused variable 
 	it('binds methods to the current raw (not stale)', async () => {
 		const { client } = createMockClient();
 		const proxy = createResilientSQLProxy(client);
 
-		// Get begin from proxy
-		const _beginFn = proxy.begin;
-
 		// Create a new raw with a different begin mock
 		const newBegin = mock(() => Promise.resolve('new-tx'));
📜 Recent review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 97bece6 and 2cc91ef.

📒 Files selected for processing (4)
  • packages/drizzle/src/postgres.ts
  • packages/drizzle/test/proxy.test.ts
  • packages/postgres/src/client.ts
  • packages/postgres/test/client.test.ts
🧰 Additional context used
📓 Path-based instructions (8)
packages/drizzle/**/*.{ts,tsx}

📄 CodeRabbit inference engine (packages/drizzle/AGENTS.md)

packages/drizzle/**/*.{ts,tsx}: Use createPostgresDrizzle() factory function to create Drizzle ORM instances
Define schemas using Drizzle's pgTable and column types rather than other schema definition methods
Use full TypeScript support with schema inference from Drizzle ORM for type-safe queries
Use drizzleAdapter from better-auth when integrating with @agentuity/auth

Files:

  • packages/drizzle/test/proxy.test.ts
  • packages/drizzle/src/postgres.ts
packages/drizzle/**/*.{test,spec}.{ts,tsx}

📄 CodeRabbit inference engine (packages/drizzle/AGENTS.md)

Ensure tests run against a PostgreSQL instance and use @agentuity/test-utils for mocking

Files:

  • packages/drizzle/test/proxy.test.ts
**/*.{ts,tsx,js,jsx,json,md}

📄 CodeRabbit inference engine (AGENTS.md)

Use Prettier for code formatting with tabs (width 3), single quotes, and semicolons

Files:

  • packages/drizzle/test/proxy.test.ts
  • packages/postgres/test/client.test.ts
  • packages/postgres/src/client.ts
  • packages/drizzle/src/postgres.ts
**/*.{ts,tsx}

📄 CodeRabbit inference engine (AGENTS.md)

Use TypeScript in strict mode with ESNext target and bundler moduleResolution

Files:

  • packages/drizzle/test/proxy.test.ts
  • packages/postgres/test/client.test.ts
  • packages/postgres/src/client.ts
  • packages/drizzle/src/postgres.ts
packages/**/*.{ts,tsx}

📄 CodeRabbit inference engine (AGENTS.md)

Use StructuredError from @agentuity/core for error handling

Files:

  • packages/drizzle/test/proxy.test.ts
  • packages/postgres/test/client.test.ts
  • packages/postgres/src/client.ts
  • packages/drizzle/src/postgres.ts
packages/*/test/**/*.{ts,tsx}

📄 CodeRabbit inference engine (AGENTS.md)

packages/*/test/**/*.{ts,tsx}: Place test files in test/ folder, never in src/ or __tests__/
Import from ../src/ in test files
Use @agentuity/test-utils for mocks in tests

Files:

  • packages/drizzle/test/proxy.test.ts
  • packages/postgres/test/client.test.ts
packages/postgres/**/*.{ts,tsx}

📄 CodeRabbit inference engine (packages/postgres/AGENTS.md)

packages/postgres/**/*.{ts,tsx}: Use tagged template literals as the primary query interface for PostgreSQL queries in @agentuity/postgres
Implement exponential backoff with jitter for reconnection attempts
Use StructuredError from @agentuity/core for error handling in @agentuity/postgres
Ensure full TypeScript type safety throughout the postgres package implementation
Use lazy connection establishment by default, with optional preconnection via config.preconnect flag
Automatically reconnect when connection is closed unexpectedly, network errors occur, or PostgreSQL server restarts
Detect SIGTERM/SIGINT signals to prevent reconnection during graceful shutdown
Register all postgres clients in a global registry for coordinated shutdown
Support transactions through begin(), commit(), and rollback() methods
Support connection pooling in @agentuity/postgres
Automatically register a shutdown hook with @agentuity/runtime when available for coordinated postgres client shutdown

Files:

  • packages/postgres/test/client.test.ts
  • packages/postgres/src/client.ts
packages/postgres/**/*.{test,spec}.{ts,tsx}

📄 CodeRabbit inference engine (packages/postgres/AGENTS.md)

Tests require a running PostgreSQL instance and should use @agentuity/test-utils for mocking

Files:

  • packages/postgres/test/client.test.ts
🧠 Learnings (1)
📚 Learning: 2025-12-21T00:31:41.858Z 
Learnt from: jhaynie
Repo: agentuity/sdk PR: 274
File: packages/cli/src/cmd/build/vite/server-bundler.ts:12-41
Timestamp: 2025-12-21T00:31:41.858Z
Learning: In Bun runtime, BuildMessage and ResolveMessage are global types and are not exported from the bun module. Do not import { BuildMessage } from 'bun' or similar; these types are available globally and should be used without import. This applies to all TypeScript files that target the Bun runtime within the repository.

Applied to files:

  • packages/drizzle/test/proxy.test.ts
  • packages/postgres/test/client.test.ts
  • packages/postgres/src/client.ts
  • packages/drizzle/src/postgres.ts
🧬 Code graph analysis (1)
packages/drizzle/test/proxy.test.ts (2)
packages/postgres/src/client.ts (2)
  • CallablePostgresClient (731-733)
  • raw (265-267)
packages/drizzle/src/postgres.ts (1)
  • createResilientSQLProxy (53-92)
🔇 Additional comments (9)
packages/postgres/src/client.ts (2)

638-651: LGTM! Clean public API exposure.

The new public executeWithRetry method is a well-documented thin wrapper that correctly delegates to the private _executeWithRetry. The generic type parameter is properly propagated, and the JSDoc clearly explains its purpose for integration layers like the Drizzle resilient proxy.

789-789: LGTM!

Correctly binds the new executeWithRetry method to the client instance, consistent with the pattern used for other methods in the callable client.

packages/postgres/test/client.test.ts (2)

1-21: LGTM! Well-structured test setup.

The createTestClient() helper is a pragmatic approach to test the retry logic without requiring a real database connection. Setting _connected = true allows testing the executeWithRetry behavior in isolation.

23-66: LGTM! Comprehensive retry behavior tests.

These tests effectively cover:

  • Basic async operation execution
  • Non-promise (synchronous) return values
  • Retry behavior on retryable errors with proper error code simulation

The error code simulation using (err as unknown as Record<string, string>).code correctly mimics the error structure that isRetryableError checks.

packages/drizzle/src/postgres.ts (2)

43-92: LGTM! Elegant resilient proxy implementation.

The proxy correctly addresses the core issue by:

  1. Always resolving client.raw at access time (line 59) - ensuring fresh connection post-reconnect
  2. Re-resolving client.raw inside the retry callback (line 71) - critical for mid-retry reconnection scenarios
  3. Wrapping unsafe() with executeWithRetry for automatic retry on transient errors
  4. Providing .values() chaining that matches Bun's SQLQuery interface
  5. Binding other methods to the current raw for correct this context

The dual-call behavior when using .values() (separate executeWithRetry invocations) is the correct design choice to ensure each operation independently benefits from retry logic.

155-166: LGTM! Clean integration.

The replacement of the static client.raw reference with resilientSQL proxy is the key fix for the issue. The comments clearly explain the rationale for using the resilient proxy.

packages/drizzle/test/proxy.test.ts (3)

17-48: LGTM! Well-designed mock client.

The createMockClient() helper provides an elegant way to simulate a CallablePostgresClient with controllable behavior. The _setRaw helper is particularly useful for testing reconnection scenarios by allowing the underlying raw connection to be swapped.

50-105: LGTM! Core proxy behavior tests.

These tests effectively validate:

  • Delegation to current raw.unsafe() with correct parameters
  • Integration with executeWithRetry for retry logic
  • .values() chaining creating a separate executor (correctly expects 2 executeWithRetry calls)
  • Reconnection handling by verifying the new raw is used after _setRaw

126-161: LGTM! Critical reconnection-during-retry test.

This test validates the most important behavior: that client.raw is re-resolved inside the retry callback. By overriding executeWithRetry to swap the raw before calling op(), it proves that the proxy correctly picks up the post-reconnection instance even mid-retry.

✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.

Comment @coderabbitai help to get the list of available commands and usage tips.

@github-actions
Copy link

github-actions bot commented Feb 9, 2026

📦 Canary Packages Published

version: 1.0.4-2cc91ef

Packages
Package Version URL
@agentuity/server 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-server-1.0.4-2cc91ef.tgz
@agentuity/drizzle 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-drizzle-1.0.4-2cc91ef.tgz
@agentuity/frontend 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-frontend-1.0.4-2cc91ef.tgz
@agentuity/workbench 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-workbench-1.0.4-2cc91ef.tgz
@agentuity/cli 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-cli-1.0.4-2cc91ef.tgz
@agentuity/opencode 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-opencode-1.0.4-2cc91ef.tgz
@agentuity/claude-code 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-claude-code-1.0.4-2cc91ef.tgz
@agentuity/evals 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-evals-1.0.4-2cc91ef.tgz
@agentuity/auth 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-auth-1.0.4-2cc91ef.tgz
@agentuity/runtime 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-runtime-1.0.4-2cc91ef.tgz
@agentuity/react 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-react-1.0.4-2cc91ef.tgz
@agentuity/core 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-core-1.0.4-2cc91ef.tgz
@agentuity/postgres 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-postgres-1.0.4-2cc91ef.tgz
@agentuity/schema 1.0.4-2cc91ef https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-schema-1.0.4-2cc91ef.tgz
Install

Add to your package.json:

{
  "dependencies": {
    "@agentuity/server": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-server-1.0.4-2cc91ef.tgz",
    "@agentuity/drizzle": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-drizzle-1.0.4-2cc91ef.tgz",
    "@agentuity/frontend": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-frontend-1.0.4-2cc91ef.tgz",
    "@agentuity/workbench": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-workbench-1.0.4-2cc91ef.tgz",
    "@agentuity/cli": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-cli-1.0.4-2cc91ef.tgz",
    "@agentuity/opencode": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-opencode-1.0.4-2cc91ef.tgz",
    "@agentuity/claude-code": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-claude-code-1.0.4-2cc91ef.tgz",
    "@agentuity/evals": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-evals-1.0.4-2cc91ef.tgz",
    "@agentuity/auth": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-auth-1.0.4-2cc91ef.tgz",
    "@agentuity/runtime": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-runtime-1.0.4-2cc91ef.tgz",
    "@agentuity/react": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-react-1.0.4-2cc91ef.tgz",
    "@agentuity/core": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-core-1.0.4-2cc91ef.tgz",
    "@agentuity/postgres": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-postgres-1.0.4-2cc91ef.tgz",
    "@agentuity/schema": "https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-schema-1.0.4-2cc91ef.tgz"
  }
}

Or install directly:

bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-server-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-drizzle-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-frontend-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-workbench-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-cli-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-opencode-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-claude-code-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-evals-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-auth-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-runtime-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-react-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-core-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-postgres-1.0.4-2cc91ef.tgz
bun add https://agentuity-sdk-objects.t3.storageapi.dev/npm/1.0.4-2cc91ef/agentuity-schema-1.0.4-2cc91ef.tgz

@jhaynie jhaynie merged commit fca6ec6 into main Feb 9, 2026
15 checks passed
@jhaynie jhaynie deleted the task/postgres-db-closed-899 branch February 9, 2026 23:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

POSTGRES_CONNECTION_CLOSED still happening even with the custom @agentuity/drizzle package.

1 participant

@jhaynie