diff --git a/src/lib/db.ts b/src/lib/db.ts index 900e0fb8..8b5322cb 100644 --- a/src/lib/db.ts +++ b/src/lib/db.ts @@ -32,16 +32,32 @@ const poolerQueryHandleError = ( { op: 'db', name: 'poolerQuery' }, () => new Promise((resolve, reject) => { - let rejected = false + let settled = false + let connectionErrorTimer: NodeJS.Timeout | undefined + const cleanup = () => { + if (connectionErrorTimer) { + clearTimeout(connectionErrorTimer) + } + pgpool.removeListener('error', connectionErrorHandler) + } + const resolveOnce = (results: pg.QueryResult) => { + if (settled) return + settled = true + cleanup() + resolve(results) + } + const rejectOnce = (err: any) => { + if (settled) return + settled = true + cleanup() + reject(err) + } const connectionErrorHandler = (err: any) => { // If the error hasn't already be propagated to the catch - if (!rejected) { + if (!settled) { // This is a trick to wait for the next tick, leaving a chance for handled errors such as // RESULT_SIZE_LIMIT to take over other stream errors such as `unexpected commandComplete message` - setTimeout(() => { - rejected = true - return reject(err) - }) + connectionErrorTimer = setTimeout(() => rejectOnce(err)) } } // This listened avoid getting uncaught exceptions for errors happening at connection level within the stream @@ -50,16 +66,11 @@ const poolerQueryHandleError = ( pgpool .query(sql, parameters) .then((results: pg.QueryResult) => { - if (!rejected) { - return resolve(results) - } + resolveOnce(results) }) .catch((err: any) => { // If the error hasn't already be handled within the error listener - if (!rejected) { - rejected = true - return reject(err) - } + rejectOnce(err) }) }) ) diff --git a/test/db.test.ts b/test/db.test.ts new file mode 100644 index 00000000..743022df --- /dev/null +++ b/test/db.test.ts @@ -0,0 +1,60 @@ +import pg from 'pg' +import { afterEach, expect, test, vi } from 'vitest' +import { init } from '../src/lib/db.js' + +afterEach(() => { + vi.restoreAllMocks() +}) + +test('successful queries remove their temporary pool error listeners', async () => { + let pool: pg.Pool | undefined + vi.spyOn(pg.Pool.prototype, 'query').mockImplementation(function (this: pg.Pool) { + pool = this + return Promise.resolve({ rows: [] }) as ReturnType + }) + + const db = init({}) + for (let i = 0; i < 12; i++) { + await db.query('select 1') + } + + expect(pool).toBeDefined() + expect(pool!.listenerCount('error')).toBe(0) + + await db.end() +}) + +test('query rejections remove their temporary pool error listener', async () => { + let pool: pg.Pool | undefined + vi.spyOn(pg.Pool.prototype, 'query').mockImplementation(function (this: pg.Pool) { + pool = this + return Promise.reject(new Error('query failed')) as ReturnType + }) + + const db = init({}) + const result = await db.query('select 1') + + expect(result.error?.message).toBe('query failed') + expect(pool!.listenerCount('error')).toBe(0) + + await db.end() +}) + +test('connection-level pool errors still reject the active query', async () => { + let pool: pg.Pool | undefined + const pending = Promise.withResolvers() + vi.spyOn(pg.Pool.prototype, 'query').mockImplementation(function (this: pg.Pool) { + pool = this + return pending.promise as ReturnType + }) + + const db = init({}) + const resultPromise = db.query('select 1') + pool!.emit('error', new Error('connection failed')) + const result = await resultPromise + + expect(result.error?.message).toBe('connection failed') + expect(pool!.listenerCount('error')).toBe(0) + + await db.end() +})