diff --git a/src/lib/routes/ready-check.ts b/src/lib/routes/ready-check.ts index 860c020732..381e90e1d4 100644 --- a/src/lib/routes/ready-check.ts +++ b/src/lib/routes/ready-check.ts @@ -81,9 +81,22 @@ export class ReadyCheckController extends Controller { } try { - connection.query({ - text: 'SELECT 1', + await connection.query('BEGIN'); + await connection.query({ + text: `SET LOCAL statement_timeout = ${timeoutMs}`, }); + await connection.query('SELECT 1'); + await connection.query('COMMIT'); + } catch (error) { + try { + await connection.query('ROLLBACK'); + } catch (rollbackError) { + this.logger.debug( + 'Failed to rollback readiness timeout transaction', + rollbackError, + ); + } + throw error; } finally { if (connection) { await client.releaseConnection(connection); diff --git a/src/lib/types/option.ts b/src/lib/types/option.ts index 44de4e6240..ef0df45051 100644 --- a/src/lib/types/option.ts +++ b/src/lib/types/option.ts @@ -302,6 +302,5 @@ export interface IUnleashConfig { userInactivityThresholdInDays: number; buildDate?: string; unleashFrontendToken?: string; - /** If true, the readiness endpoint will attempt a simple Postgres query to verify DB availability */ checkDbOnReady?: boolean; } diff --git a/src/test/e2e/ready.e2e.test.ts b/src/test/e2e/ready.e2e.test.ts index 7ddb95fc4e..f84ef6c388 100644 --- a/src/test/e2e/ready.e2e.test.ts +++ b/src/test/e2e/ready.e2e.test.ts @@ -29,6 +29,46 @@ describe('DB is up', () => { ); await request.get('/ready').expect(200).expect('{"health":"GOOD"}'); }); + + test('fails fast when readiness query hangs', async () => { + const { request } = await setupAppWithCustomConfig( + db.stores, + { checkDbOnReady: true }, + db.rawDatabase, + ); + + const pool = db.rawDatabase.client.pool; + const originalAcquire = pool.acquire.bind(pool); + + pool.acquire = () => { + const pending = originalAcquire(); + pending.promise = pending.promise.then((conn: any) => { + const originalQuery = conn.query; + conn.query = (queryConfig: any, ...args: any[]) => { + const isSelectOne = + queryConfig?.toUpperCase() === 'SELECT 1'; + + if (isSelectOne) { + return originalQuery.call( + conn, + 'SELECT pg_sleep(1)', + ...args, + ); + } + + return originalQuery.call(conn, queryConfig, ...args); + }; + return conn; + }); + return pending; + }; + + try { + await request.get('/ready').expect(503); + } finally { + pool.acquire = originalAcquire; + } + }); }); describe('DB is down', () => {