fix: use query level statement_timeout

package.json

@@ -30,8 +30,8 @@
     "test": "run-s db:clean db:run test:run db:clean",
     "db:clean": "cd test/db && docker compose down",
     "db:run": "cd test/db && docker compose up --detach --wait",
-    "test:run": "PG_META_MAX_RESULT_SIZE_MB=20 PG_QUERY_TIMEOUT_SECS=3 PG_CONN_TIMEOUT_SECS=30 vitest run --coverage",
-    "test:update": "run-s db:clean db:run && PG_META_MAX_RESULT_SIZE_MB=20 PG_QUERY_TIMEOUT_SECS=3 PG_CONN_TIMEOUT_SECS=30 vitest run --update && run-s db:clean"
+    "test:run": "PG_META_MAX_RESULT_SIZE_MB=20 PG_QUERY_TIMEOUT_SECS=5 PG_CONN_TIMEOUT_SECS=30 vitest run --coverage",
+    "test:update": "run-s db:clean db:run && PG_META_MAX_RESULT_SIZE_MB=20 PG_QUERY_TIMEOUT_SECS=5 PG_CONN_TIMEOUT_SECS=30 vitest run --update && run-s db:clean"
   },
   "engines": {
     "node": ">=20",

src/lib/db.ts

@@ -2,6 +2,9 @@ import pg from 'pg'
 import * as Sentry from '@sentry/node'
 import { parse as parseArray } from 'postgres-array'
 import { PostgresMetaResult, PoolConfig } from './types.js'
+import { PG_STATEMENT_TIMEOUT_SECS } from '../server/constants.js'
+
+const STATEMENT_TIMEOUT_QUERY_PREFIX = `SET statement_timeout='${PG_STATEMENT_TIMEOUT_SECS}s';`
 
 pg.types.setTypeParser(pg.types.builtins.INT8, (x) => {
   const asNumber = Number(x)
@@ -112,18 +115,23 @@ export const init: (config: PoolConfig) => {
             attributes: { sql: trackQueryInSentry ? sql : 'custom' },
           },
           async () => {
+            // 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
+            // allow to set this parameter at connection time
+            const sqlWithStatementTimeout = `${STATEMENT_TIMEOUT_QUERY_PREFIX}${sql}`
             try {
               if (!pool) {
                 const pool = new pg.Pool(config)
-                let res = await poolerQueryHandleError(pool, sql)
+                let res = await poolerQueryHandleError(pool, sqlWithStatementTimeout)
                 if (Array.isArray(res)) {
                   res = res.reverse().find((x) => x.rows.length !== 0) ?? { rows: [] }
                 }
                 await pool.end()
                 return { data: res.rows, error: null }
               }
 
-              let res = await poolerQueryHandleError(pool, sql)
+              let res = await poolerQueryHandleError(pool, sqlWithStatementTimeout)
               if (Array.isArray(res)) {
                 res = res.reverse().find((x) => x.rows.length !== 0) ?? { rows: [] }
               }
@@ -147,7 +155,11 @@ export const init: (config: PoolConfig) => {
                   formattedError += '\n'
                   if (error.position) {
                     // error.position is 1-based
-                    const position = Number(error.position) - 1
+                    // we also remove our `SET statement_timeout = 'XXs';\n` from the position
+                    const position =
+                      Number(error.position) - 1 - STATEMENT_TIMEOUT_QUERY_PREFIX.length
+                    // we set the new error position
+                    error.position = `${position + 1}`
 
                     let line = ''
                     let lineNumber = 0

src/lib/secrets.ts

@@ -1,6 +1,3 @@
-// Use dynamic import to support module mock
-const fs = await import('node:fs/promises')
-
 export const getSecret = async (key: string) => {
   if (!key) {
     return ''
@@ -15,6 +12,8 @@ export const getSecret = async (key: string) => {
   if (!file) {
     return ''
   }
+  // Use dynamic import to support module mock
+  const fs = await import('node:fs/promises')
 
   return await fs.readFile(file, { encoding: 'utf8' }).catch((e) => {
     if (e.code == 'ENOENT') {

src/server/constants.ts

@@ -17,6 +17,7 @@ const PG_META_DB_SSL_MODE = process.env.PG_META_DB_SSL_MODE || 'disable'
 
 const PG_CONN_TIMEOUT_SECS = Number(process.env.PG_CONN_TIMEOUT_SECS || 15)
 const PG_QUERY_TIMEOUT_SECS = Number(process.env.PG_QUERY_TIMEOUT_SECS || 55)
+export const PG_STATEMENT_TIMEOUT_SECS = PG_QUERY_TIMEOUT_SECS + 1
 
 export let PG_CONNECTION = process.env.PG_META_DB_URL
 if (!PG_CONNECTION) {
@@ -59,9 +60,6 @@ export const PG_META_MAX_RESULT_SIZE = process.env.PG_META_MAX_RESULT_SIZE_MB
 export const DEFAULT_POOL_CONFIG: PoolConfig = {
   max: 1,
   connectionTimeoutMillis: PG_CONN_TIMEOUT_SECS * 1000,
-  // node-postgrest 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
-  statement_timeout: (PG_QUERY_TIMEOUT_SECS + 1) * 1000,
   query_timeout: PG_QUERY_TIMEOUT_SECS * 1000,
   ssl: PG_META_DB_SSL_ROOT_CERT ? { ca: PG_META_DB_SSL_ROOT_CERT } : undefined,
   application_name: `postgres-meta ${pkg.version}`,

test/server/query-timeout.ts

@@ -2,32 +2,38 @@ 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
+
 describe('test query timeout', () => {
-  test('query timeout after 3s and connection cleanup', async () => {
-    const query = `SELECT pg_sleep(10);`
-    // Execute a query that will sleep for 10 seconds
-    const res = await app.inject({
-      method: 'POST',
-      path: '/query',
-      payload: {
-        query,
-      },
-    })
+  test(
+    `query timeout after ${TIMEOUT}s and connection cleanup`,
+    async () => {
+      const query = `SELECT pg_sleep(${TIMEOUT});`
+      // Execute a query that will sleep for 10 seconds
+      const res = await app.inject({
+        method: 'POST',
+        path: '/query',
+        payload: {
+          query,
+        },
+      })
 
-    // Check that we get the proper timeout error response
-    expect(res.statusCode).toBe(408) // Request Timeout
-    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))
+      // Check that we get the proper timeout error response
+      expect(res.statusCode).toBe(408) // Request Timeout
+      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(`
+      // 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)
-  }, 5000)
+      // Should have no active connections except for our current query
+      expect(connectionsRes.data).toHaveLength(0)
+    },
+    TIMEOUT * 1000
+  )
 })