feat(server): introduce maxResultSize limit and fix pg errors handling

Author: avallete

src/lib/db.ts

@@ -1,7 +1,6 @@
-import pg, { PoolConfig } from 'pg'
-import { DatabaseError } from 'pg-protocol'
+import pg from 'pg'
 import { parse as parseArray } from 'postgres-array'
-import { PostgresMetaResult } from './types.js'
+import { PostgresMetaResult, PoolConfig } from './types.js'
 
 pg.types.setTypeParser(pg.types.builtins.INT8, (x) => {
   const asNumber = Number(x)
@@ -21,6 +20,38 @@ pg.types.setTypeParser(1185, parseArray) // _timestamptz
 pg.types.setTypeParser(600, (x) => x) // point
 pg.types.setTypeParser(1017, (x) => x) // _point
 
+// Ensure any query will have an appropriate error handler on the pool to prevent connections errors
+// to bubble up all the stack eventually killing the server
+const poolerQueryHandleError = (pool: pg.Pool, sql: string): Promise<pg.QueryResult<any>> => {
+  return new Promise((resolve, reject) => {
+    let rejected = false
+    const connectionErrorHandler = (err: any) => {
+      // If the error hasn't already be propagated to the catch
+      if (!rejected) {
+        rejected = true
+        reject(err)
+      }
+    }
+    // This listened avoid getting uncaught exceptions for errors happening at connection level within the stream
+    // such as parse or RESULT_SIZE_EXCEEDED errors instead, handle the error gracefully by bubbling in up to the caller
+    pool.once('error', connectionErrorHandler)
+    pool
+      .query(sql)
+      .then((results: pg.QueryResult<any>) => {
+        if (!rejected) {
+          return resolve(results)
+        }
+      })
+      .catch((err: any) => {
+        // If the error hasn't already be handled within the error listener
+        if (!rejected) {
+          rejected = true
+          return reject(err)
+        }
+      })
+  })
+}
+
 export const init: (config: PoolConfig) => {
   query: (sql: string) => Promise<PostgresMetaResult<any>>
   end: () => Promise<void>
@@ -60,26 +91,27 @@ export const init: (config: PoolConfig) => {
   // compromise: if we run `query` after `pool.end()` is called (i.e. pool is
   // `null`), we temporarily create a pool and close it right after.
   let pool: pg.Pool | null = new pg.Pool(config)
+
   return {
     async query(sql) {
       try {
         if (!pool) {
           const pool = new pg.Pool(config)
-          let res = await pool.query(sql)
+          let res = await poolerQueryHandleError(pool, sql)
           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 pool.query(sql)
+        let res = await poolerQueryHandleError(pool, sql)
         if (Array.isArray(res)) {
           res = res.reverse().find((x) => x.rows.length !== 0) ?? { rows: [] }
         }
         return { data: res.rows, error: null }
       } catch (error: any) {
-        if (error instanceof DatabaseError) {
+        if (error.constructor.name === 'DatabaseError') {
           // Roughly based on:
           // - https://github.com/postgres/postgres/blob/fc4089f3c65a5f1b413a3299ba02b66a8e5e37d0/src/interfaces/libpq/fe-protocol3.c#L1018
           // - https://github.com/brianc/node-postgres/blob/b1a8947738ce0af004cb926f79829bb2abc64aa6/packages/pg/lib/native/query.js#L33
@@ -147,7 +179,60 @@ ${' '.repeat(5 + lineNumber.toString().length + 2 + lineOffset)}^
           }
         }
 
-        return { data: null, error: { message: error.message } }
+        // Handle stream errors and result size exceeded errors
+        if (error.code === 'RESULT_SIZE_EXCEEDED') {
+          // Force kill the connection without waiting for graceful shutdown
+          const _pool = pool
+          pool = null
+          try {
+            if (_pool) {
+              // Force kill the connection by destroying the socket
+              const client = (_pool as any)._clients?.[0]
+              if (client?.connection?.stream) {
+                client.connection.stream.destroy()
+              }
+            }
+          } catch (endError) {
+            // Ignore any errors during cleanup
+          }
+          return {
+            data: null,
+            error: {
+              message: `Query result size (${error.resultSize} bytes) exceeded the configured limit (${error.maxResultSize} bytes)`,
+              code: error.code,
+              resultSize: error.resultSize,
+              maxResultSize: error.maxResultSize,
+            },
+          }
+        }
+
+        // Handle other stream errors
+        if (error.code === 'STREAM_ERROR') {
+          // Force kill the connection without waiting for graceful shutdown
+          const _pool = pool
+          pool = null
+          try {
+            if (_pool) {
+              // Force kill the connection by destroying the socket
+              const client = (_pool as any)._clients?.[0]
+              if (client?.connection?.stream) {
+                client.connection.stream.destroy()
+              }
+            }
+          } catch (endError) {
+            // Ignore any errors during cleanup
+          }
+          return {
+            data: null,
+            error: {
+              message: 'Stream error occurred while processing query',
+              code: error.code,
+              details: error.message,
+            },
+          }
+        }
+
+        return { data: null, error: { ...error, message: error.message } }
       }
     },
 
@@ -156,7 +241,11 @@ ${' '.repeat(5 + lineNumber.toString().length + 2 + lineOffset)}^
       pool = null
       // Gracefully wait for active connections to be idle, then close all
       // connections in the pool.
-      if (_pool) await _pool.end()
+      if (_pool) {
+        // Remove all listeners before ending to prevent memory leaks
+        _pool.removeAllListeners()
+        await _pool.end()
+      }
     },
   }
 }

src/lib/types.ts

@@ -1,6 +1,7 @@
 import { Static, Type } from '@sinclair/typebox'
 import { DatabaseError } from 'pg-protocol'
 import type { Options as PrettierOptions } from 'prettier'
+import { PoolConfig as PgPoolConfig } from 'pg'
 
 export interface FormatterOptions extends PrettierOptions {}
 
@@ -251,13 +252,7 @@ export const postgresPublicationSchema = Type.Object({
   publish_delete: Type.Boolean(),
   publish_truncate: Type.Boolean(),
   tables: Type.Union([
-    Type.Array(
-      Type.Object({
-        id: Type.Integer(),
-        name: Type.String(),
-        schema: Type.String(),
-      })
-    ),
+    Type.Array(Type.Object({ id: Type.Integer(), name: Type.String(), schema: Type.String() })),
     Type.Null(),
   ]),
 })
@@ -445,12 +440,7 @@ export const postgresTypeSchema = Type.Object({
   schema: Type.String(),
   format: Type.String(),
   enums: Type.Array(Type.String()),
-  attributes: Type.Array(
-    Type.Object({
-      name: Type.String(),
-      type_id: Type.Integer(),
-    })
-  ),
+  attributes: Type.Array(Type.Object({ name: Type.String(), type_id: Type.Integer() })),
   comment: Type.Union([Type.String(), Type.Null()]),
 })
 export type PostgresType = Static<typeof postgresTypeSchema>
@@ -596,3 +586,7 @@ export const postgresColumnPrivilegesRevokeSchema = Type.Object({
   ]),
 })
 export type PostgresColumnPrivilegesRevoke = Static<typeof postgresColumnPrivilegesRevokeSchema>
+
+export interface PoolConfig extends PgPoolConfig {
+  maxResultSize?: number
+}

src/server/app.ts

@@ -5,14 +5,10 @@ import { PG_META_REQ_HEADER } from './constants.js'
 import routes from './routes/index.js'
 import { extractRequestForLogging } from './utils.js'
 // Pseudo package declared only for this module
-import pkg from '#package.json' assert { type: 'json' }
+import pkg from '#package.json' with { type: 'json' }
 
 export const build = (opts: FastifyServerOptions = {}): FastifyInstance => {
-  const app = fastify({
-    disableRequestLogging: true,
-    requestIdHeader: PG_META_REQ_HEADER,
-    ...opts,
-  })
+  const app = fastify({ disableRequestLogging: true, requestIdHeader: PG_META_REQ_HEADER, ...opts })
 
   app.setErrorHandler((error, request, reply) => {
     app.log.error({ error: error.toString(), request: extractRequestForLogging(request) })

src/server/constants.ts

@@ -1,8 +1,8 @@
 import crypto from 'crypto'
-import { PoolConfig } from 'pg'
+import { PoolConfig } from '../lib/types.js'
 import { getSecret } from '../lib/secrets.js'
 import { AccessControl } from './templates/swift.js'
-import pkg from '#package.json' assert { type: 'json' }
+import pkg from '#package.json' with { type: 'json' }
 
 export const PG_META_HOST = process.env.PG_META_HOST || '0.0.0.0'
 export const PG_META_PORT = Number(process.env.PG_META_PORT || 1337)
@@ -49,11 +49,16 @@ export const GENERATE_TYPES_SWIFT_ACCESS_CONTROL = process.env
   ? (process.env.PG_META_GENERATE_TYPES_SWIFT_ACCESS_CONTROL as AccessControl)
   : 'internal'
 
+export const PG_META_MAX_RESULT_SIZE = process.env.PG_META_MAX_RESULT_SIZE
+  ? parseInt(process.env.PG_META_MAX_RESULT_SIZE, 10)
+  : 2 * 1024 * 1024 * 1024 // default to 2GB max query size result
+
 export const DEFAULT_POOL_CONFIG: PoolConfig = {
   max: 1,
   connectionTimeoutMillis: PG_CONN_TIMEOUT_SECS * 1000,
   ssl: PG_META_DB_SSL_ROOT_CERT ? { ca: PG_META_DB_SSL_ROOT_CERT } : undefined,
   application_name: `postgres-meta ${pkg.version}`,
+  maxResultSize: PG_META_MAX_RESULT_SIZE,
 }
 
 export const PG_META_REQ_HEADER = process.env.PG_META_REQ_HEADER || 'request-id'

test/index.test.ts

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

test/lib/secrets.ts

@@ -6,10 +6,7 @@ vi.mock('node:fs/promises', async (): Promise<typeof import('node:fs/promises')>
   const originalModule =
     await vi.importActual<typeof import('node:fs/promises')>('node:fs/promises')
   const readFile = vi.fn()
-  return {
-    ...originalModule,
-    readFile,
-  }
+  return { ...originalModule, readFile }
 })
 
 describe('getSecret', () => {
@@ -57,6 +54,6 @@ describe('getSecret', () => {
     const e: NodeJS.ErrnoException = new Error('permission denied')
     e.code = 'EACCES'
     vi.mocked(readFile).mockRejectedValueOnce(e)
-    expect(getSecret('SECRET')).rejects.toThrow()
+    await expect(getSecret('SECRET')).rejects.toThrow()
   })
 })

test/server/result-size-limit.ts

@@ -0,0 +1,71 @@
+import { expect, test, beforeAll, afterAll, describe } from 'vitest'
+import { app } from './utils'
+import { pgMeta } from '../lib/utils'
+
+describe('test max-result-size limit', () => {
+  // Create a table with large data for testing
+  beforeAll(async () => {
+    // Create a table with a large text column
+    await pgMeta.query(`
+      CREATE TABLE large_data (
+        id SERIAL PRIMARY KEY,
+        data TEXT
+      );
+    `)
+    await pgMeta.query(`
+      CREATE TABLE small_data (
+        id SERIAL PRIMARY KEY,
+        data TEXT
+      );
+    `)
+
+    // Insert data that will exceed our limit in tests it's set around ~20MB
+    await pgMeta.query(`
+      INSERT INTO large_data (data)
+      SELECT repeat('x', 1024 * 1024) -- 1MB of data per row
+      FROM generate_series(1, 50);
+    `)
+    await pgMeta.query(`
+      INSERT INTO small_data (data)
+      SELECT repeat('x', 10 * 1024) -- 10KB per row
+      FROM generate_series(1, 50);
+    `)
+  })
+
+  afterAll(async () => {
+    // Clean up the test table
+    await pgMeta.query('DROP TABLE large_data;')
+    await pgMeta.query('DROP TABLE small_data;')
+  })
+
+  test('query exceeding result size limit', async () => {
+    // Set a small maxResultSize (50MB)
+    const res = await app.inject({
+      method: 'POST',
+      path: '/query',
+      headers: { pg: 'postgresql://postgres:postgres@localhost:5432/postgres' },
+      payload: { query: 'SELECT * FROM large_data;' },
+    })
+
+    // Check that we get the proper error response
+    expect(res.statusCode).toBe(400)
+    expect(res.json()).toMatchObject({
+      error: expect.stringContaining('Query result size'),
+      code: 'RESULT_SIZE_EXCEEDED',
+      resultSize: expect.any(Number),
+      maxResultSize: 20 * 1024 * 1024,
+    })
+
+    // Verify that subsequent queries still work and the server isn't killed
+    const nextRes = await app.inject({
+      method: 'POST',
+      path: '/query',
+      headers: { pg: 'postgresql://postgres:postgres@localhost:5432/postgres' },
+      payload: { query: 'SELECT * FROM small_data;' },
+    })
+
+    expect(nextRes.statusCode).toBe(200)
+    // Should have retrieve the 50 rows as expected
+    expect(nextRes.json()).toHaveLength(50)
+  })
+})