Merge pull request #947 from supabase/fix/only-add-statement-timeout-on-direct-connection

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/PostgresMeta.ts

@@ -22,7 +22,10 @@ import { init } from './db.js'
 import { PostgresMetaResult, PoolConfig } from './types.js'
 
 export default class PostgresMeta {
-  query: (sql: string, trackQueryInSentry?: boolean) => Promise<PostgresMetaResult<any>>
+  query: (
+    sql: string,
+    opts?: { statementQueryTimeout?: number; trackQueryInSentry?: boolean }
+  ) => Promise<PostgresMetaResult<any>>
   end: () => Promise<void>
   columnPrivileges: PostgresMetaColumnPrivileges
   columns: PostgresMetaColumns

src/lib/db.ts

@@ -62,7 +62,10 @@ const poolerQueryHandleError = (pgpool: pg.Pool, sql: string): Promise<pg.QueryR
 }
 
 export const init: (config: PoolConfig) => {
-  query: (sql: string, trackQueryInSentry?: boolean) => Promise<PostgresMetaResult<any>>
+  query: (
+    sql: string,
+    opts?: { statementQueryTimeout?: number; trackQueryInSentry?: boolean }
+  ) => Promise<PostgresMetaResult<any>>
   end: () => Promise<void>
 } = (config) => {
   return Sentry.startSpan({ op: 'db', name: 'db.init' }, () => {
@@ -103,7 +106,10 @@ export const init: (config: PoolConfig) => {
     let pool: pg.Pool | null = new pg.Pool(config)
 
     return {
-      async query(sql, trackQueryInSentry = true) {
+      async query(
+        sql,
+        { statementQueryTimeout, trackQueryInSentry } = { trackQueryInSentry: true }
+      ) {
         return Sentry.startSpan(
           // For metrics purposes, log the query that will be run if it's not an user provided query (with possibly sentitives infos)
           {
@@ -112,18 +118,26 @@ export const init: (config: PoolConfig) => {
             attributes: { sql: trackQueryInSentry ? sql : 'custom' },
           },
           async () => {
+            const statementTimeoutQueryPrefix = statementQueryTimeout
+              ? `SET statement_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
+            // allow to set this parameter at connection time
+            const sqlWithStatementTimeout = `${statementTimeoutQueryPrefix}${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 +161,10 @@ 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 - statementTimeoutQueryPrefix.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

@@ -59,9 +59,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}`,

src/server/routes/query.ts

@@ -19,11 +19,18 @@ export default async (fastify: FastifyInstance) => {
     Body: {
       query: string
     }
+    Querystring: {
+      statementTimeoutSecs?: number
+    }
   }>('/', async (request, reply) => {
+    const statementTimeoutSecs = request.query.statementTimeoutSecs
     errorOnEmptyQuery(request)
     const config = createConnectionConfig(request)
     const pgMeta = new PostgresMeta(config)
-    const { data, error } = await pgMeta.query(request.body.query, false)
+    const { data, error } = await pgMeta.query(request.body.query, {
+      trackQueryInSentry: true,
+      statementQueryTimeout: statementTimeoutSecs,
+    })
     await pgMeta.end()
     if (error) {
       request.log.error({ error, request: extractRequestForLogging(request) })

test/server/query-timeout.ts

@@ -2,32 +2,72 @@ 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
+
 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,
-      },
-    })
-
-    // 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(`
+  test(
+    `query timeout after ${TIMEOUT}s 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',
+        query: `statementTimeoutSecs=${STATEMENT_TIMEOUT}`,
+        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))
+
+      // 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',
+    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',
+        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
+      await new Promise((resolve) => setTimeout(resolve, 1000))
+
+      // Verify that the connection has not been cleaned up sinice there is no statementTimetout
+      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(1)
+    },
+    TIMEOUT * 1000
+  )
 })