fix: authenticate PostgreSQL IAM connections on demand (#421)

* fix: authenticate PostgreSQL IAM connections on demand

Resolve IAM passwords when opening a connection so failed login attempts remain retryable.

- Keep pools available after idle disconnects and share concurrent authentication.
- Preserve AWS provider selection and reread refreshed credential files.
- Cover recovery with real pg and isolated AWS SDK checks.

* fix: correct PostgreSQL SSL configuration type

- Use the pool SSL type instead of the nonexistent pg.ConnectionOptions export.

---------

Co-authored-by: Daniel E. <daniel.e@teachingstrategies.com>

Author: deinspanjer

docs/config/toml.mdx

@@ -255,6 +255,34 @@ Sources define database connections. Each source represents a database that DBHu
   another AWS identity. When `aws_profile` is omitted, the AWS SDK default
   credential chain is used.
 
+  PostgreSQL generates a token when the driver opens a physical connection.
+  Existing pooled connections stay open; DBHub does not refresh PostgreSQL IAM
+  credentials on a timer. If a server closes an idle session, the next request
+  opens a replacement using a fresh token. Sources can have different server
+  idle timeouts without configuring a matching refresh interval in DBHub.
+
+  Credential renewal follows your AWS configuration. A `credential_process`
+  helper can renew credentials silently or perform its own browser login.
+  With AWS SSO, complete `aws sso login --profile <profile>` when the SDK
+  reports that login is required, then retry the request. A failed attempt
+  remains retryable; DBHub does not automatically replay SQL or repeatedly
+  retry authentication in the background. Concurrent PostgreSQL connections
+  for one source share an in-progress token request.
+
+  Use `lazy = true` to defer the initial connection and authentication until
+  the first request. Otherwise, the initial connection is checked at startup.
+  `connection_timeout` includes time spent waiting for the password callback;
+  a credential helper's browser login may outlast it. Completing login and
+  retrying is supported, but a socket timeout does not cancel the helper.
+
+  Replaced shared-credential-file values are read on the next authentication
+  attempt. Changes to environment variables in another shell cannot update a
+  running DBHub process; use a refreshable credential provider for temporary
+  credentials. Direct assume-role MFA prompts require a credential helper;
+  DBHub does not supply an interactive MFA callback.
+
+  MySQL and MariaDB continue to use periodic IAM pool refresh.
+
   ```toml
   [[sources]]
   id = "production"

src/connectors/__tests__/manager.test.ts

@@ -1,6 +1,6 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { ConnectorManager } from "../manager.js";
-import { ConnectorRegistry } from "../interface.js";
+import { ConnectorRegistry, type Connector, type ConnectorConfig } from "../interface.js";
 import { SSHTunnel } from "../../utils/ssh-tunnel.js";
 import type { SourceConfig } from "../../types/config.js";
 import { homedir } from "os";
@@ -278,7 +278,7 @@ describe("ConnectorManager IAM DSN rewrite", () => {
       dsn: "mysql://dbuser%40example.com:ignored@mydb.abc123.eu-west-1.rds.amazonaws.com:3306/mydb?connectTimeout=5000&sslmode=disable",
     };
 
-    const dsn = await (manager as any).buildConnectionDSN(source);
+    const dsn = await (manager as any).buildConnectionDSN(source, {});
 
     expect(mocks.generateRdsAuthToken).toHaveBeenCalledWith({
       hostname: "mydb.abc123.eu-west-1.rds.amazonaws.com",
@@ -481,3 +481,113 @@ describe("ConnectorManager IAM refresh recovery", () => {
     );
   });
 });
+
+describe("PostgreSQL IAM authentication on demand", () => {
+  const source: SourceConfig = {
+    id: "production", type: "postgres", host: "db.example.com", port: 5432,
+    database: "db", user: "db_user", aws_iam_auth: true,
+    aws_region: "us-east-1", aws_profile: "production", lazy: true,
+  };
+  let manager: ConnectorManager;
+  let config: ConnectorConfig;
+  let dsn: string;
+  const disconnect = vi.fn();
+  const connect = vi.fn();
+
+  beforeEach(() => {
+    vi.useFakeTimers();
+    vi.clearAllMocks();
+    mocks.generateRdsAuthToken.mockReset().mockResolvedValue("token");
+    connect.mockReset().mockImplementation(async (value: string, _init: string, options: ConnectorConfig) => {
+      dsn = value;
+      config = options;
+      await options.password?.();
+    });
+    manager = new ConnectorManager();
+    vi.spyOn(console, "error").mockImplementation(() => {});
+    vi.spyOn(ConnectorRegistry, "getConnectorForDSN").mockReturnValue({
+      clone: () => ({
+        id: "postgres", disconnect, connect,
+      }),
+    } as unknown as Connector);
+  });
+
+  afterEach(async () => {
+    await manager.disconnect();
+    vi.useRealTimers();
+    vi.restoreAllMocks();
+  });
+
+  it("authenticates new connections, never refreshes an idle source on a timer", async () => {
+    await manager.connectWithSources([source]);
+    expect(mocks.generateRdsAuthToken).not.toHaveBeenCalled();
+    await manager.ensureConnected(source.id);
+    expect(config.password).toBeTypeOf("function");
+    expect(dsn).not.toContain("token");
+    expect(dsn).toContain("sslmode=require");
+    expect(mocks.generateRdsAuthToken).toHaveBeenCalledTimes(1);
+
+    await vi.advanceTimersByTimeAsync(24 * 60 * 60 * 1000);
+    expect(disconnect).not.toHaveBeenCalled();
+    expect(mocks.generateRdsAuthToken).toHaveBeenCalledTimes(1);
+
+    // A pool opening several sockets shares only the in-flight token request.
+    const password = config.password!;
+    await Promise.all(Array.from({ length: 8 }, () => password()));
+    expect(mocks.generateRdsAuthToken).toHaveBeenCalledTimes(2);
+    mocks.generateRdsAuthToken.mockRejectedValueOnce(new Error("SSO session expired"));
+    await expect(password()).rejects.toThrow("SSO session expired");
+    await expect(password()).resolves.toBe("token");
+    expect(manager.getConnector(source.id)).toBeDefined();
+    expect(mocks.generateRdsAuthToken).toHaveBeenCalledTimes(4);
+  });
+
+  it("retries failed initial authentication only when another request arrives", async () => {
+    mocks.generateRdsAuthToken.mockRejectedValueOnce(new Error("SSO session expired"));
+    await manager.connectWithSources([source]);
+    await expect(manager.ensureConnected(source.id)).rejects.toThrow("SSO session expired");
+    await vi.advanceTimersByTimeAsync(24 * 60 * 60 * 1000);
+    expect(mocks.generateRdsAuthToken).toHaveBeenCalledTimes(1);
+    await Promise.all(Array.from({ length: 8 }, () => manager.ensureConnected(source.id)));
+    expect(mocks.generateRdsAuthToken).toHaveBeenCalledTimes(2);
+    expect(manager.getSourceIds()).toEqual([source.id]);
+    expect(manager.getConnector(source.id)).toBeDefined();
+  });
+
+  it("signs the original endpoint through SSH and cleans up a failed authentication", async () => {
+    mocks.looksLikeSSHAlias.mockReturnValue(false);
+    vi.spyOn(SSHTunnel.prototype, "establish").mockResolvedValue({
+      localPort: 15432, localHost: "127.0.0.1",
+    } as any);
+    const close = vi.spyOn(SSHTunnel.prototype, "close").mockResolvedValue();
+    mocks.generateRdsAuthToken.mockRejectedValueOnce(new Error("SSO session expired"));
+    await manager.connectWithSources([{ ...source, sslmode: "verify-full",
+      ssh_host: "bastion.example.com", ssh_user: "user", ssh_key: "/fake/key" }]);
+    await expect(manager.ensureConnected(source.id)).rejects.toThrow("SSO session expired");
+    expect(close).toHaveBeenCalledTimes(1);
+    await manager.ensureConnected(source.id);
+    expect(dsn).toContain("127.0.0.1:15432");
+    expect(dsn).toContain("sslmode=verify-full");
+    expect(mocks.generateRdsAuthToken).toHaveBeenLastCalledWith({
+      hostname: "db.example.com", port: 5432, username: "db_user",
+      region: "us-east-1", profile: "production",
+    });
+  });
+
+  it("shares a still-running credential helper after an initial socket timeout", async () => {
+    let finishLogin!: (token: string) => void;
+    mocks.generateRdsAuthToken.mockReturnValueOnce(new Promise<string>(resolve => { finishLogin = resolve; }));
+    connect.mockImplementationOnce(async (_dsn, _init, options: ConnectorConfig) => {
+      void options.password!().catch(() => {});
+      throw new Error("Connection timeout");
+    });
+    await manager.connectWithSources([source]);
+    await expect(manager.ensureConnected(source.id)).rejects.toThrow("Connection timeout");
+    const retry = manager.ensureConnected(source.id);
+    await vi.waitFor(() => expect(connect).toHaveBeenCalledTimes(2));
+    expect(mocks.generateRdsAuthToken).toHaveBeenCalledTimes(1);
+    finishLogin("token");
+    await retry;
+    expect(manager.getConnector(source.id)).toBeDefined();
+  });
+});

src/connectors/__tests__/postgres-auth.test.ts

@@ -0,0 +1,80 @@
+import { createServer, type Socket } from "node:net";
+import { once } from "node:events";
+import { afterEach, expect, it, vi } from "vitest";
+import { PostgresConnector } from "../postgres/index.js";
+
+// A loopback PostgreSQL protocol peer: real pg sockets and password handling,
+// no AWS access, real database, or SQL execution. Only auth/empty query responses.
+function message(type: string, body: Buffer) {
+  const length = Buffer.alloc(4);
+  length.writeInt32BE(body.length + 4);
+  return Buffer.concat([Buffer.from(type), length, body]);
+}
+const ready = message("Z", Buffer.from("I"));
+const sockets = new Set<Socket>();
+const peers = new Map<string, Socket>();
+const connectors: PostgresConnector[] = [];
+const server = createServer(socket => {
+  sockets.add(socket);
+  socket.on("close", () => sockets.delete(socket));
+  let startup = true;
+  let buffer: Buffer = Buffer.alloc(0);
+  socket.on("data", data => {
+    buffer = Buffer.concat([buffer, data]);
+    while (buffer.length >= (startup ? 4 : 5)) {
+      const length = buffer.readInt32BE(startup ? 0 : 1) + (startup ? 0 : 1);
+      if (buffer.length < length) return;
+      const packet = buffer.subarray(0, length);
+      buffer = buffer.subarray(length);
+      if (startup) {
+        const params = packet.subarray(8).toString().split("\0");
+        peers.set(params[params.indexOf("database") + 1], socket);
+        startup = false;
+        socket.write(message("R", Buffer.from([0, 0, 0, 3]))); // cleartext password request
+      } else if (packet[0] === 112) { // password
+        socket.write(Buffer.concat([message("R", Buffer.alloc(4)), ready]));
+      } else if (packet[0] === 81) { // simple query
+        socket.write(Buffer.concat([message("C", Buffer.from("SELECT 0\0")), ready]));
+      } else if (packet[0] === 88) { // terminate
+        socket.end();
+      }
+    }
+  });
+});
+
+afterEach(async () => {
+  await Promise.all(connectors.splice(0).map(connector => connector.disconnect()));
+  for (const socket of sockets) socket.destroy();
+  peers.clear();
+  if (server.listening) await new Promise<void>(resolve => server.close(() => resolve()));
+  vi.restoreAllMocks();
+});
+
+it("recovers on demand after idle session expiry and failed auth, independently per source", async () => {
+  vi.spyOn(console, "error").mockImplementation(() => {});
+  server.listen(0, "127.0.0.1");
+  await once(server, "listening");
+  const port = (server.address() as { port: number }).port;
+  const passwords = [vi.fn().mockResolvedValue("short-token"), vi.fn().mockResolvedValue("long-token")];
+  for (const [i, name] of ["short", "long"].entries()) {
+    const connector = new PostgresConnector();
+    connectors.push(connector);
+    await connector.connect(`postgres://user@127.0.0.1:${port}/${name}?sslmode=disable`, undefined, {
+      password: passwords[i], connectionTimeoutSeconds: i + 1,
+    });
+  }
+  expect(passwords.map(password => password.mock.calls.length)).toEqual([1, 1]);
+
+  // Server-side idle expiry is independent of IAM expiry. pg must remove this
+  // idle client without an unhandled pool error or disturbing the other source.
+  peers.get("short")!.end(message("E", Buffer.from("SFATAL\0C57P05\0Mterminating connection due to idle-session timeout\0\0")));
+  await vi.waitFor(() => expect(console.error).toHaveBeenCalledWith(
+    expect.any(String), expect.stringContaining("idle-session timeout")));
+  expect(passwords.map(password => password.mock.calls.length)).toEqual([1, 1]);
+
+  passwords[0].mockRejectedValueOnce(new Error("SSO session expired"));
+  await expect(connectors[0].getSchemas()).rejects.toThrow("SSO session expired");
+  await expect(connectors[0].getSchemas()).resolves.toEqual([]);
+  await expect(connectors[1].getSchemas()).resolves.toEqual([]);
+  expect(passwords.map(password => password.mock.calls.length)).toEqual([3, 1]);
+});

src/connectors/interface.ts

@@ -138,6 +138,8 @@ export interface ExecuteOptions {
  * Different databases may use different subset of these options
  */
 export interface ConnectorConfig {
+  /** Password resolved for each new physical connection (PostgreSQL). */
+  password?: () => Promise<string>;
   /** Connection timeout in seconds (PostgreSQL, MySQL, MariaDB, SQL Server) */
   connectionTimeoutSeconds?: number;
   /** Query timeout in seconds (PostgreSQL, MySQL, MariaDB, SQL Server) */

src/connectors/manager.ts

@@ -30,6 +30,9 @@ export class ConnectorManager {
   // Lazy connection support
   private lazySources: Map<string, SourceConfig> = new Map(); // Sources pending lazy connection
   private pendingConnections: Map<string, Promise<void>> = new Map(); // Prevent race conditions
+  // A socket timeout does not cancel a credential helper. Share its in-flight
+  // attempt across retries of the same config, but never across changed profiles.
+  private pendingIamTokens = new WeakMap<SourceConfig, Promise<string>>();
 
   constructor() {
     if (!managerInstance) {
@@ -142,8 +145,9 @@ export class ConnectorManager {
    */
   private async connectSource(source: SourceConfig): Promise<void> {
     const sourceId = source.id;
+    const config: ConnectorConfig = {};
     // Build DSN from source config
-    const dsn = await this.buildConnectionDSN(source);
+    const dsn = await this.buildConnectionDSN(source, config);
     console.error(`  - ${sourceId}: ${redactDSN(dsn)}`);
 
     // Setup SSH tunnel if needed
@@ -256,7 +260,6 @@ export class ConnectorManager {
     (connector as any).sourceId = sourceId;
 
     // Build config for database-specific options
-    const config: ConnectorConfig = {};
     if (source.connection_timeout !== undefined) {
       config.connectionTimeoutSeconds = source.connection_timeout;
     }
@@ -318,7 +321,7 @@ export class ConnectorManager {
     // Store source config (for API exposure)
     this.sourceConfigs.set(sourceId, source);
 
-    // Keep AWS IAM auth sources fresh by rotating pool credentials before token expiry.
+    // MySQL/MariaDB still rotate pools; PostgreSQL authenticates on demand.
     this.scheduleIamRefresh(source);
   }
 
@@ -497,7 +500,7 @@ export class ConnectorManager {
       clearTimeout(existingTimer);
       this.iamRefreshTimers.delete(sourceId);
     }
-    if (!source.aws_iam_auth) {
+    if (!source.aws_iam_auth || source.type === "postgres") {
       return;
     }
 
@@ -566,7 +569,7 @@ export class ConnectorManager {
    * Build a connection DSN, optionally replacing password with
    * an AWS RDS IAM auth token when aws_iam_auth is enabled.
    */
-  private async buildConnectionDSN(source: SourceConfig): Promise<string> {
+  private async buildConnectionDSN(source: SourceConfig, config: ConnectorConfig = {}): Promise<string> {
     const dsn = buildDSNFromSource(source);
 
     if (!source.aws_iam_auth) {
@@ -597,13 +600,22 @@ export class ConnectorManager {
       );
     }
 
-    const token = await generateRdsAuthToken({
-      hostname,
-      port,
-      username,
-      region: source.aws_region,
-      profile: source.aws_profile,
-    });
+    // Share concurrent authentication attempts, never retain a failed promise or
+    // cache a token's lifetime ourselves. The AWS provider owns credential refresh.
+    const password = () => {
+      let pending = this.pendingIamTokens.get(source);
+      if (!pending) {
+        pending = generateRdsAuthToken({
+          hostname, port, username, region: source.aws_region!, profile: source.aws_profile,
+        }).finally(() => { this.pendingIamTokens.delete(source); });
+        this.pendingIamTokens.set(source, pending);
+      }
+      return pending;
+    };
+    if (source.type === "postgres") {
+      config.password = password;
+    }
+    const token = source.type === "postgres" ? "" : await password();
 
     const queryParams = new Map(parsed.searchParams);
     const currentSslMode = queryParams.get("sslmode");

src/connectors/postgres/index.ts

@@ -105,7 +105,7 @@ class PostgresDSNParser implements DSNParser {
         port: url.port ? parseInt(url.port) : 5432,
         database: url.pathname ? url.pathname.substring(1) : '', // Remove leading '/' if exists
         user: url.username,
-        password: url.password,
+        password: config?.password ?? url.password,
       };
 
       let sslmode: string | undefined;
@@ -168,7 +168,7 @@ class PostgresDSNParser implements DSNParser {
       } else if (sslmode === "require") {
         poolConfig.ssl = { rejectUnauthorized: false };
       } else if (sslmode === "verify-ca" || sslmode === "verify-full") {
-        const sslConfig: pg.ConnectionOptions["ssl"] & object = { rejectUnauthorized: true };
+        const sslConfig: pg.PoolConfig["ssl"] & object = { rejectUnauthorized: true };
         // verify-ca checks the certificate chain but does not verify the server hostname,
         // matching libpq behavior. verify-full (the default with rejectUnauthorized: true)
         // verifies both the certificate chain and the hostname.