Merge pull request #946 from supabase/fix/query-timeout-connection-closing

fix(query): ensure that open connection are killed after 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 vitest run --coverage",
-    "test:update": "run-s db:clean db:run && PG_META_MAX_RESULT_SIZE_MB=20 vitest run --update && run-s db:clean"
+    "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"
   },
   "engines": {
     "node": ">=20",

src/server/constants.ts

@@ -59,6 +59,9 @@ 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/index.test.ts

@@ -23,3 +23,4 @@ import './server/ssl'
 import './server/table-privileges'
 import './server/typegen'
 import './server/result-size-limit'
+import './server/query-timeout'

test/server/query-timeout.ts

@@ -0,0 +1,33 @@
+import { expect, test, describe } from 'vitest'
+import { app } from './utils'
+import { pgMeta } from '../lib/utils'
+
+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(`
+      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)
+})