fix(postgres): enforce query timeout on server (#415)

Author: rabberbock

src/connectors/__tests__/dsn-parser.test.ts

@@ -99,6 +99,18 @@ describe('DSN Parser - PostgreSQL SSL Modes', () => {
   });
 });
 
+describe('DSN Parser - PostgreSQL query timeout', () => {
+  it('configures a server-side statement timeout before the client fallback', async () => {
+    const parser = new PostgresConnector().dsnParser;
+    const config = await parser.parse('postgres://user:pass@localhost:5432/db', {
+      queryTimeoutSeconds: 30,
+    });
+
+    expect(config.statement_timeout).toBe(30_000);
+    expect(config.query_timeout).toBe(35_000);
+  });
+});
+
 describe('DSN Parser - AWS IAM Authentication', () => {
   describe('MySQL', () => {
     const connector = new MySQLConnector();

src/connectors/__tests__/postgres.integration.test.ts

@@ -214,6 +214,65 @@ describe('PostgreSQL Connector Integration Tests', () => {
   postgresTest.createErrorHandlingTests();
   postgresTest.createSSLTests();
   describe('PostgreSQL-specific Features', () => {
+    it('should cancel a timed-out query on the PostgreSQL server', async () => {
+      const timedConnector = new PostgresConnector();
+      const observer = new PostgresConnector();
+      const probe = 'dbhub_query_timeout_probe';
+
+      const runningProbeCount = async (): Promise<number> => {
+        const result = await observer.executeSQL(
+          `SELECT count(*)::int AS count
+           FROM pg_stat_activity
+           WHERE state = 'active'
+             AND query LIKE '%${probe}%'
+             AND query NOT LIKE '%pg_stat_activity%'`,
+          {}
+        );
+        return result.resultSets[0].rows[0].count;
+      };
+
+      const waitFor = async (
+        predicate: () => Promise<boolean>,
+        timeoutMs: number
+      ): Promise<boolean> => {
+        const deadline = Date.now() + timeoutMs;
+        while (Date.now() < deadline) {
+          if (await predicate()) return true;
+          await new Promise((resolve) => setTimeout(resolve, 50));
+        }
+        return false;
+      };
+
+      try {
+        await timedConnector.connect(postgresTest.connectionString, undefined, {
+          queryTimeoutSeconds: 1,
+        });
+        await observer.connect(postgresTest.connectionString);
+
+        const query = timedConnector.executeSQL(
+          `SELECT pg_sleep(10), '${probe}'`,
+          { readonly: true }
+        );
+        const settled = query.then(
+          () => null,
+          (error) => error as NodeJS.ErrnoException
+        );
+
+        expect(await waitFor(async () => (await runningProbeCount()) === 1, 5_000)).toBe(true);
+
+        const error = await settled;
+        expect(error).toBeInstanceOf(Error);
+        expect(error?.code).toBe('57014');
+        expect(await waitFor(async () => (await runningProbeCount()) === 0, 2_000)).toBe(true);
+
+        const after = await timedConnector.executeSQL('SELECT 1 AS ok', { readonly: true });
+        expect(after.resultSets[0].rows[0].ok).toBe(1);
+      } finally {
+        await timedConnector.disconnect();
+        await observer.disconnect();
+      }
+    }, 20_000);
+
     it('should execute multiple statements with transaction support', async () => {
       const result = await postgresTest.connector.executeSQL(`
         INSERT INTO users (name, email, age) VALUES ('Multi User 1', 'multi1@example.com', 30);
@@ -756,4 +815,4 @@ describe('PostgreSQL Connector Integration Tests', () => {
       }
     });
   });
-});
+});

src/connectors/postgres/index.ts

@@ -26,6 +26,8 @@ import { splitSQLStatements } from "../../utils/sql-parser.js";
 import { FailedToReadCertificate } from "./failed-to-read-certificate.js";
 import { closeQuietly } from "../../utils/resource-cleanup.js";
 
+const POSTGRES_CLIENT_QUERY_TIMEOUT_GRACE_MS = 5_000;
+
 /**
  * PostgreSQL DSN Parser
  * Handles DSN strings like: postgres://user:password@localhost:5432/dbname?sslmode=disable
@@ -113,10 +115,14 @@ class PostgresDSNParser implements DSNParser {
         poolConfig.connectionTimeoutMillis = connectionTimeoutSeconds * 1000;
       }
 
-      // Apply query timeout if specified (client-side timeout)
+      // Apply the configured limit on the server so a timed-out statement does
+      // not keep running after DBHub stops waiting for it. Retain the client-side
+      // timeout as a fallback, with enough grace for PostgreSQL's cancellation
+      // response to arrive first.
       if (queryTimeoutSeconds !== undefined) {
-        // pg library expects query_timeout in milliseconds
-        poolConfig.query_timeout = queryTimeoutSeconds * 1000;
+        const queryTimeoutMs = queryTimeoutSeconds * 1000;
+        poolConfig.statement_timeout = queryTimeoutMs;
+        poolConfig.query_timeout = queryTimeoutMs + POSTGRES_CLIENT_QUERY_TIMEOUT_GRACE_MS;
       }
 
       return poolConfig;