feat(query): add queryTimeoutSecs param for pool timeout override

src/lib/db.ts

@@ -124,9 +124,10 @@ export const init: (config: PoolConfig) => {
           async () => {
             // Use statement_timeout AND idle_session_timeout to ensure the connection will be killed even if idle after
             // timeout time.
-            const statementTimeoutQueryPrefix = statementQueryTimeout
-              ? `SET statement_timeout='${statementQueryTimeout}s'; SET idle_session_timeout='${statementQueryTimeout}s';`
-              : ''
+            const statementTimeoutQueryPrefix =
+              statementQueryTimeout !== undefined
+                ? `SET statement_timeout='${statementQueryTimeout}s'; SET idle_session_timeout='${statementQueryTimeout}s';`
+                : ''
             // node-postgres need a statement_timeout to kill the connection when timeout is reached
             // otherwise the query will keep running on the database even if query timeout was reached
             // This need to be added at query and not connection level because poolers (pgbouncer) doesn't

src/server/routes/query.ts

@@ -17,11 +17,11 @@ export default async (fastify: FastifyInstance) => {
   fastify.post<{
     Headers: { pg: string; 'x-pg-application-name'?: string }
     Body: { query: string; parameters?: unknown[] }
-    Querystring: { statementTimeoutSecs?: number }
+    Querystring: { statementTimeoutSecs?: number; queryTimeoutSecs?: number }
   }>('/', async (request, reply) => {
     const statementTimeoutSecs = request.query.statementTimeoutSecs
     errorOnEmptyQuery(request)
-    const config = createConnectionConfig(request)
+    const config = createConnectionConfig(request, request.query.queryTimeoutSecs)
     const pgMeta = new PostgresMeta(config)
     const { data, error } = await pgMeta.query(request.body.query, {
       trackQueryInSentry: true,

src/server/utils.ts

@@ -23,9 +23,19 @@ export const extractRequestForLogging = (request: FastifyRequest) => {
   }
 }
 
-export function createConnectionConfig(request: FastifyRequest): PoolConfig {
+export function createConnectionConfig(
+  request: FastifyRequest,
+  queryTimeoutSecs?: number | string
+): PoolConfig {
   const connectionString = request.headers.pg as string
-  const config = { ...DEFAULT_POOL_CONFIG, connectionString }
+  const timeout = queryTimeoutSecs !== undefined ? Number(queryTimeoutSecs) : undefined
+  const config = {
+    ...DEFAULT_POOL_CONFIG,
+    connectionString,
+    ...(timeout !== undefined && {
+      query_timeout: timeout === 0 ? undefined : timeout * 1000,
+    }),
+  }
 
   // Override application_name if custom one provided in header
   if (request.headers['x-pg-application-name']) {

test/server/query-timeout.ts

@@ -2,15 +2,16 @@ import { expect, test, describe } from 'vitest'
 import { app } from './utils'
 import { pgMeta } from '../lib/utils'
 
-const TIMEOUT = (Number(process.env.PG_QUERY_TIMEOUT_SECS) ?? 10) + 2
-const STATEMENT_TIMEOUT = (Number(process.env.PG_QUERY_TIMEOUT_SECS) ?? 10) + 1
+const PG_QUERY_TIMEOUT = Number(process.env.PG_QUERY_TIMEOUT_SECS) ?? 10
+const TIMEOUT = PG_QUERY_TIMEOUT + 2
+const STATEMENT_TIMEOUT = PG_QUERY_TIMEOUT + 1
+const CUSTOM_QUERY_TIMEOUT = 2
 
 describe('test query timeout', () => {
   test(
-    `query timeout after ${TIMEOUT}s and connection cleanup`,
+    `pool timeout after ${TIMEOUT}s with statementTimeoutSecs and connection cleanup`,
     async () => {
       const query = `SELECT pg_sleep(${TIMEOUT + 10});`
-      // Execute a query that will sleep for 10 seconds
       const res = await app.inject({
         method: 'POST',
         path: '/query',
@@ -20,30 +21,25 @@ describe('test query timeout', () => {
         },
       })
 
-      // Check that we get the proper timeout error response
-      expect(res.statusCode).toBe(408) // Request Timeout
+      expect(res.statusCode).toBe(408)
       expect(res.json()).toMatchObject({
         error: expect.stringContaining('Query read timeout'),
       })
-      // wait one second for the statement timeout to take effect
       await new Promise((resolve) => setTimeout(resolve, 1000))
 
-      // Verify that the connection has been cleaned up by checking active connections
       const connectionsRes = await pgMeta.query(`
       SELECT * FROM pg_stat_activity where application_name = 'postgres-meta 0.0.0-automated' and query ILIKE '%${query}%';
     `)
 
-      // Should have no active connections except for our current query
       expect(connectionsRes.data).toHaveLength(0)
     },
     TIMEOUT * 1000
   )
 
   test(
-    'query without timeout parameter should not have timeout',
+    'absent queryTimeoutSecs uses default pool timeout',
     async () => {
       const query = `SELECT pg_sleep(${TIMEOUT + 10});`
-      // Execute a query that will sleep for 10 seconds without specifying timeout
       const res = await app.inject({
         method: 'POST',
         path: '/query',
@@ -52,22 +48,79 @@ describe('test query timeout', () => {
         },
       })
 
-      // Check that we get the proper timeout error response
-      expect(res.statusCode).toBe(408) // Request Timeout
+      expect(res.statusCode).toBe(408)
       expect(res.json()).toMatchObject({
         error: expect.stringContaining('Query read timeout'),
       })
-      // wait one second
       await new Promise((resolve) => setTimeout(resolve, 1000))
 
-      // Verify that the connection has not been cleaned up sinice there is no statementTimetout
+      // No statementTimeout was set, so the PG-side query is still running
       const connectionsRes = await pgMeta.query(`
       SELECT * FROM pg_stat_activity where application_name = 'postgres-meta 0.0.0-automated' and query ILIKE '%${query}%';
     `)
 
-      // Should have no active connections except for our current query
       expect(connectionsRes.data).toHaveLength(1)
     },
     TIMEOUT * 1000
   )
+
+  test(
+    'queryTimeoutSecs=0 disables pool-level timeout',
+    async () => {
+      const sleepSecs = PG_QUERY_TIMEOUT + 1
+      const res = await app.inject({
+        method: 'POST',
+        path: '/query',
+        query: 'queryTimeoutSecs=0',
+        payload: {
+          query: `SELECT pg_sleep(${sleepSecs});`,
+        },
+      })
+
+      expect(res.statusCode).toBe(200)
+      expect(res.json()).toEqual([{ pg_sleep: '' }])
+    },
+    (PG_QUERY_TIMEOUT + 5) * 1000
+  )
+
+  test(
+    'custom queryTimeoutSecs overrides default pool timeout',
+    async () => {
+      const res = await app.inject({
+        method: 'POST',
+        path: '/query',
+        query: `queryTimeoutSecs=${CUSTOM_QUERY_TIMEOUT}`,
+        payload: {
+          query: `SELECT pg_sleep(${PG_QUERY_TIMEOUT + 1});`,
+        },
+      })
+
+      expect(res.statusCode).toBe(408)
+      expect(res.json()).toMatchObject({
+        error: expect.stringContaining('Query read timeout'),
+      })
+    },
+    (CUSTOM_QUERY_TIMEOUT + 5) * 1000
+  )
+
+  test(
+    'queryTimeoutSecs=0 with statementTimeoutSecs still enforces statement timeout',
+    async () => {
+      const res = await app.inject({
+        method: 'POST',
+        path: '/query',
+        query: `queryTimeoutSecs=0&statementTimeoutSecs=${STATEMENT_TIMEOUT}`,
+        payload: {
+          query: `SELECT pg_sleep(${STATEMENT_TIMEOUT + 5});`,
+        },
+      })
+
+      // Statement timeout fires (not pool timeout), producing a DatabaseError
+      expect(res.statusCode).toBe(400)
+      expect(res.json()).toMatchObject({
+        error: expect.stringContaining('canceling statement due to statement timeout'),
+      })
+    },
+    (STATEMENT_TIMEOUT + 5) * 1000
+  )
 })