feat(postgres-js): add OCC retry with exponential backoff (#495)

## Summary

Added automatic retry with exponential backoff + jitter for OCC
conflicts on begin() transactions, with constructor-level and per-call
configuration. Fixed region parsing to gracefully fall back to
`AWS_REGION` / `AWS_DEFAULT_REGION` env vars.

## Changes

- `src/occ-retry.ts` — retry module (resolveRetryConfig,
executeWithRetry, isOCCError, validateRetryConfig, calculateBackoff)
- `src/client.ts` — wrapBegin() patches sql.begin() for both TCP and
WebSocket clients; region parsing fix
- `src/index.ts` — export isOCCError, OCCRetryConfig, Logger types
- `test/occ-retry.test.ts` — unit tests for isOCCError,
resolveRetryConfig, validation
- `test/integration/dsql.integration.test.ts` — 3 OCC integration tests
- `README.md` — OCC Retry documentation section + config table entries
- `example/` — updated preferred and no-pool examples with retry usage

## Test plan

- [x] Unit tests pass (46 tests — 2 suites)
- [x] TypeScript compiles clean
- [x] Integration tests pass against DSQL cluster (33 tests —
constructor opt-in, per-call opt-in, non-OCC passthrough)

*Issue #, if available:*
N/A

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Co-authored-by: Prabakaran Annadurai <prabaw@amazon.com>

Author: praba2210

node/postgres-js/README.md

@@ -42,6 +42,7 @@ The Aurora DSQL Connector for Postgres.js is designed to understand these requir
 - **Full TypeScript Support** - Provides full type safety
 - **AWS Credentials Support** - Supports various AWS credential providers (default, profile-based, custom)
 - **Connection Pooling Compatibility** - Works seamlessly with Postgres.js' built-in connection pooling
+- **OCC Retry** - Automatic retry with exponential backoff on optimistic concurrency conflicts
 
 ## Quick start guide
 
@@ -127,6 +128,50 @@ const sql = AuroraDSQLPostgres({
 });
 ```
 
+### OCC Retry
+
+Aurora DSQL uses optimistic concurrency control (OCC). Transactions that conflict are rejected at
+commit time with error codes `OC000`, `OC001`, or `40001`. The connector can automatically retry
+`begin()` transactions on these conflicts with exponential backoff.
+
+```typescript
+import { auroraDSQLPostgres } from '@aws/aurora-dsql-postgresjs-connector';
+
+// Opt-in at connection level — all begin() calls will retry on OCC conflicts
+const sql = auroraDSQLPostgres({
+  host: 'your-cluster.dsql.us-east-1.on.aws',
+  username: 'admin',
+  retry: true,
+});
+
+// No per-call config needed
+await sql.begin(async (tx) => {
+  await tx`INSERT INTO accounts(id, balance) VALUES(${1}, ${100})`;
+});
+```
+
+The retry config accepts a partial object — unset fields use defaults:
+
+| Field          | Default | Description                        |
+| -------------- | ------- | ---------------------------------- |
+| `maxRetries`   | 3       | Number of retry attempts after the initial try (0 = no retry) |
+| `baseDelayMs`  | 1       | Initial backoff delay in ms        |
+| `maxDelayMs`   | 100     | Maximum base backoff delay in ms (jitter is added on top) |
+| `jitterFactor` | 0.25    | Jitter as a fraction of delay      |
+
+Per-call opt-in/out:
+
+```typescript
+// Per-call opt-in (no constructor config needed)
+await sql.begin(async (tx) => { ... }, { retry: true });
+await sql.begin(async (tx) => { ... }, { retry: { maxRetries: 10 } });
+
+// Per-call opt-out (overrides constructor config)
+await sql.begin(async (tx) => { ... }, { retry: false });
+```
+
+**Note:** The callback may be invoked multiple times on OCC conflicts. Keep it idempotent - avoid non-database side effects (HTTP calls, message publishes, external counters) inside the transaction body.
+
 ### Configuration Options
 
 | Option                      | Type                             | Required | Description                                              |
@@ -137,6 +182,8 @@ const sql = AuroraDSQLPostgres({
 | `region`                    | `string?`                        | No       | AWS region (auto-detected from hostname if not provided) |
 | `customCredentialsProvider` | `AwsCredentialIdentityProvider?` | No       | Custom AWS credentials provider                          |
 | `tokenDurationSecs`         | `number?`                        | No       | Token expiration time in seconds                         |
+| `retry`                     | `Partial<OCCRetryConfig> \| boolean` | No   | OCC retry configuration (see above)                      |
+| `logger`                    | `Logger`                         | No       | Logger with `error` method and optional `debug` (e.g. `console`)  |
 
 All standard [Postgres.js options](https://github.com/porsager/postgres?tab=readme-ov-file#connection-details) are also supported.
 

node/postgres-js/example/README.md

@@ -59,6 +59,7 @@ The example demonstrates the following operations:
 
 - Opening a connection to an Aurora DSQL cluster
 - Creating a table
+- Transactional writes with automatic OCC retry
 - Inserting and querying data
 
 The example is designed to work with both admin and non-admin users:

node/postgres-js/example/src/alternatives/no_connection_pool/example_with_no_connection_pool.js

@@ -14,6 +14,7 @@ function getConnection(clusterEndpoint, user) {
   return auroraDSQLPostgres({
     host: clusterEndpoint,
     user: user,
+    retry: true,
     // Other DSQL options:
     // region: 'us-east-1',
     // profile: awsProfile,
@@ -46,9 +47,11 @@ async function example() {
                      telephone VARCHAR(20)
                  )`;
 
-    // Insert some data
-    await client`INSERT INTO ${client(schema)}.owner(name, city, telephone)
-                 VALUES ('John Doe', 'Anytown', '555-555-0150')`;
+    // Transactional write with OCC retry
+    await client.begin(async (tx) => {
+      await tx`INSERT INTO ${tx(schema)}.owner(name, city, telephone)
+               VALUES ('John Doe', 'Anytown', '555-555-0150')`;
+    });
 
     // Check that data is inserted by reading it back
     const result = await client`SELECT id, city

node/postgres-js/example/src/example_preferred.js

@@ -15,6 +15,7 @@ function createPooledConnection(clusterEndpoint, user) {
     max: 10, // Connection pool size
     idle_timeout: 30, // Idle connection timeout in seconds
     connect_timeout: 10, // Connection timeout in seconds
+    retry: true,
   });
 }
 
@@ -43,6 +44,30 @@ async function example() {
     await Promise.all(workers);
 
     console.log("Connection pool with concurrent connections exercised successfully");
+
+    // Create table
+    await sql`CREATE TABLE IF NOT EXISTS owner (
+      id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+      name VARCHAR(30) NOT NULL,
+      city VARCHAR(80) NOT NULL,
+      telephone VARCHAR(20)
+    )`;
+
+    // Transactional write
+    await sql.begin(async (tx) => {
+      await tx`INSERT INTO owner(name, city, telephone) VALUES(${"John Doe"}, ${"Anytown"}, ${"555-555-1900"})`;
+    });
+
+    // Verify the write
+    const result = await sql`SELECT name, city FROM owner WHERE name = ${"John Doe"}`;
+    assert.strictEqual(result[0].name, "John Doe");
+    assert.strictEqual(result[0].city, "Anytown");
+    console.log(`Inserted: name=${result[0].name}, city=${result[0].city}`);
+
+    // Clean up
+    await sql`DELETE FROM owner WHERE name = ${"John Doe"}`;
+
+    console.log("Transactional write completed successfully");
   } catch (error) {
     console.error(error);
     throw error;

node/postgres-js/src/client.ts

@@ -9,6 +9,7 @@ import {
 } from "@aws-sdk/types";
 import { DsqlSigner, DsqlSignerConfig } from "@aws-sdk/dsql-signer";
 import { createPostgresWs } from "./postgres-web-socket";
+import { Logger, OCCRetryConfig, resolveRetryConfig, executeWithRetry } from "./occ-retry";
 
 // Version is injected at build time via tsdown
 declare const __VERSION__: string;
@@ -94,6 +95,45 @@ function setupDsqlSigner(opts: any): { signerConfig: DsqlSignerConfig; } {
   return { signerConfig };
 }
 
+function wrapBegin(sql: postgres.Sql<any>, opts: { retry?: Partial<OCCRetryConfig> | boolean; logger?: Logger }): void {
+  if (!sql.begin) return;
+  const { retry: constructorRetry, logger } = opts;
+  if (constructorRetry && constructorRetry !== true) {
+    resolveRetryConfig(constructorRetry, undefined);
+  }
+  const originalBegin = sql.begin.bind(sql);
+  // postgres-js does not call this.begin internally, so direct mutation is safe.
+  (sql as any).begin = (first: any, second?: any, third?: any) => {
+    let beginOpts: string | undefined;
+    let fn: (sql: any) => any;
+    let callOpts: { retry?: Partial<OCCRetryConfig> | boolean } | undefined;
+
+    if (typeof first === "string") {
+      beginOpts = first;
+      fn = second;
+      callOpts = third;
+    } else {
+      fn = first;
+      callOpts = second;
+    }
+
+    if (typeof fn !== "function") {
+      throw new TypeError("sql.begin: callback function is required");
+    }
+
+    const config = resolveRetryConfig(constructorRetry, callOpts?.retry);
+    if (!config) {
+      return beginOpts ? originalBegin(beginOpts, fn) : originalBegin(fn);
+    }
+
+    return executeWithRetry(
+      () => beginOpts ? originalBegin(beginOpts, fn) : originalBegin(fn),
+      config,
+      logger,
+    );
+  };
+}
+
 export function auroraDSQLPostgres<
   T extends Record<string, postgres.PostgresType> = {}
 >(
@@ -156,9 +196,12 @@ export function auroraDSQLPostgres<
     ...opts,
     pass: () => getToken(signer, opts.username),
   };
-  return typeof urlOrOptions === "string"
+  const sql = typeof urlOrOptions === "string"
     ? postgres(urlOrOptions, postgresOpts)
     : postgres(postgresOpts);
+
+  wrapBegin(sql, { retry: opts.retry, logger: opts.logger });
+  return sql;
 }
 
 export function auroraDSQLWsPostgres<
@@ -238,9 +281,13 @@ export function auroraDSQLWsPostgres<
 
   opts.port = 443;
 
-  return typeof urlOrOptions === "string"
-    ? postgres(urlOrOptions, opts)
-    : postgres(opts);
+  const postgresOpts: postgres.Options<T> = { ...opts };
+  const sql = typeof urlOrOptions === "string"
+    ? postgres(urlOrOptions, postgresOpts)
+    : postgres(postgresOpts);
+
+  wrapBegin(sql, { retry: opts.retry, logger: opts.logger });
+  return sql;
 }
 
 
@@ -337,6 +384,10 @@ export interface AuroraDSQLConfig<T extends Record<string, PostgresType<T>>> ext
 
   customCredentialsProvider?: AwsCredentialIdentity | AwsCredentialIdentityProvider;
 
+  retry?: Partial<OCCRetryConfig> | boolean;
+
+  logger?: Logger;
+
 }
 export interface AuroraDSQLWsConfig<T extends Record<string, PostgresType<T>>>
   extends Omit<postgres.Options<T>, "socket" | "ssl" | "port"> {
@@ -348,4 +399,6 @@ export interface AuroraDSQLWsConfig<T extends Record<string, PostgresType<T>>>
   connectionCheck?: boolean;
   connectionId?: string;
   onReservedConnectionClose?: (connectionId?: string) => void;
+  retry?: Partial<OCCRetryConfig> | boolean;
+  logger?: Logger;
 }

node/postgres-js/src/index.ts

@@ -3,4 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 export { auroraDSQLPostgres, auroraDSQLWsPostgres } from './client.js';
-export type { AuroraDSQLConfig, AuroraDSQLWsConfig } from './client.js';
+export { isOCCError } from './occ-retry.js';
+export type { AuroraDSQLConfig, AuroraDSQLWsConfig } from './client.js';
+export type { OCCRetryConfig, Logger } from './occ-retry.js';

node/postgres-js/src/occ-retry.ts

@@ -0,0 +1,128 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+interface Logger {
+  debug?(msg: string, ...args: unknown[]): void;
+  error(msg: string, ...args: unknown[]): void;
+}
+
+interface OCCRetryConfig {
+  maxRetries: number;
+  baseDelayMs: number;
+  maxDelayMs: number;
+  jitterFactor: number;
+}
+
+const DEFAULT_CONFIG: OCCRetryConfig = {
+  maxRetries: 3,
+  baseDelayMs: 1,
+  maxDelayMs: 100,
+  jitterFactor: 0.25,
+};
+
+function isOCCError(error: unknown): boolean {
+  if (!error || typeof error !== "object") {
+    return false;
+  }
+
+  const dbError = error as { code?: string };
+  if (!dbError.code) {
+    return false;
+  }
+
+  return dbError.code === "OC000" || dbError.code === "OC001" || dbError.code === "40001";
+}
+
+function validateRetryConfig(config: OCCRetryConfig): void {
+  for (const [k, v] of Object.entries(config)) {
+    if (typeof v !== "number" || !Number.isFinite(v)) {
+      throw new Error(`${k} must be a finite number`);
+    }
+  }
+  if (!Number.isInteger(config.maxRetries)) {
+    throw new Error("maxRetries must be an integer");
+  }
+  if (config.maxRetries < 0) {
+    throw new Error("maxRetries must be >= 0");
+  }
+  if (config.maxRetries > 100) {
+    throw new Error("maxRetries must not exceed 100");
+  }
+  if (config.baseDelayMs <= 0) {
+    throw new Error("baseDelayMs must be greater than 0");
+  }
+  if (config.maxDelayMs < config.baseDelayMs) {
+    throw new Error("maxDelayMs must be >= baseDelayMs");
+  }
+  if (config.jitterFactor < 0 || config.jitterFactor > 1) {
+    throw new Error("jitterFactor must be between 0.0 and 1.0");
+  }
+}
+
+// Returns null when retry is not opted-in (no config from either source).
+function resolveRetryConfig(
+  constructorRetry: Partial<OCCRetryConfig> | boolean | undefined,
+  perCallRetry: Partial<OCCRetryConfig> | boolean | undefined,
+): OCCRetryConfig | null {
+  if (perCallRetry === false) return null;
+  if (perCallRetry === undefined && (constructorRetry === undefined || constructorRetry === false)) return null;
+
+  const base = (constructorRetry === true || constructorRetry === undefined || constructorRetry === false)
+    ? {}
+    : constructorRetry;
+  const override = (perCallRetry === true || perCallRetry === undefined)
+    ? {}
+    : perCallRetry;
+
+  const resolved: OCCRetryConfig = { ...DEFAULT_CONFIG, ...base, ...override };
+  validateRetryConfig(resolved);
+  return resolved;
+}
+
+function calculateBackoff(config: OCCRetryConfig, attempt: number): number {
+  const exponent = Math.min(attempt - 1, 31);
+  const delay = Math.min(config.baseDelayMs * Math.pow(2.0, exponent), config.maxDelayMs);
+  const jitter = delay * Math.random() * config.jitterFactor;
+  return delay + jitter;
+}
+
+async function executeWithRetry<T>(
+  executeCallback: () => Promise<T>,
+  config: OCCRetryConfig,
+  logger?: Logger,
+): Promise<T> {
+  if (config.maxRetries <= 0) {
+    return executeCallback();
+  }
+
+  let lastError: Error | undefined;
+
+  for (let attempt = 1; attempt <= config.maxRetries + 1; attempt++) {
+    try {
+      return await executeCallback();
+    } catch (error) {
+      if (!isOCCError(error)) {
+        throw error;
+      }
+
+      lastError = error as Error;
+
+      if (attempt <= config.maxRetries) {
+        const delay = calculateBackoff(config, attempt);
+        logger?.debug?.(
+          `OCC conflict on attempt ${attempt}/${config.maxRetries + 1}, retrying in ${delay.toFixed(1)}ms`,
+        );
+
+        await new Promise<void>((resolve) => setTimeout(resolve, delay));
+      }
+    }
+  }
+
+  logger?.error(`OCC retry exhausted after ${config.maxRetries + 1} attempts`);
+  throw lastError!;
+}
+
+export type { Logger, OCCRetryConfig };
+export { resolveRetryConfig, executeWithRetry, isOCCError };