fix: use query level statement timeout

Author: avallete

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

@@ -4,6 +4,8 @@ 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)
   if (Number.isSafeInteger(asNumber)) {
@@ -117,7 +119,7 @@ export const init: (config: PoolConfig) => {
             // 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 = `SET statement_timeout='${PG_STATEMENT_TIMEOUT_SECS}s';\n${sql}`
+            const sqlWithStatementTimeout = `${STATEMENT_TIMEOUT_QUERY_PREFIX}${sql}`
             try {
               if (!pool) {
                 const pool = new pg.Pool(config)
@@ -153,13 +155,17 @@ 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
                     let lineOffset = 0
 
-                    const lines = sqlWithStatementTimeout.split('\n')
+                    const lines = sql.split('\n')
                     let currentOffset = 0
                     for (let i = 0; i < lines.length; i++) {
                       if (currentOffset + lines[i].length > position) {

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
+  )
 })

test/server/result-size-limit.ts

@@ -72,23 +72,29 @@ describe('test js parser error max result', () => {
   // Create a table with large data for testing
   beforeAll(async () => {
     // Create a table with a large text column
-    await pgMeta.query(`
+    await pgMeta.query(
+      `
       CREATE TABLE very_large_data (
         id SERIAL PRIMARY KEY,
         data TEXT
       );
-    `)
+    `,
+      false
+    )
 
     // Insert data that will exceed our limit in tests it's set around ~20MB
-    await pgMeta.query(`
+    await pgMeta.query(
+      `
       INSERT INTO very_large_data (data)
       VALUES (repeat('x', 710 * 1024 * 1024)) -- 700+MB string will raise a JS exception at parse time
-    `)
+    `,
+      false
+    )
   })
 
   afterAll(async () => {
     // Clean up the test table
-    await pgMeta.query('DROP TABLE very_large_data;')
+    await pgMeta.query('DROP TABLE very_large_data;', false)
   })
 
   test(