Checking for readiness tests

src/lib/openapi/spec/ready-check-schema.ts

@@ -10,9 +10,9 @@ export const readyCheckSchema = {
     properties: {
         health: {
             description:
-                'The readiness state this Unleash instance is in. GOOD if the server is up and running. It never returns BAD; if the server is unhealthy you will get an unsuccessful http response.',
+                'The readiness state this Unleash instance is in. GOOD if the server is up and running. If the server is unhealthy you will get an unsuccessful http response.',
             type: 'string',
-            enum: ['GOOD', 'BAD'],
+            enum: ['GOOD'],
             example: 'GOOD',
         },
     },

src/lib/routes/ready-check.test.ts

@@ -1,44 +0,0 @@
-import supertest, { type Test } from 'supertest';
-import { createServices } from '../services/index.js';
-import { createTestConfig } from '../../test/config/test-config.js';
-
-import createStores from '../../test/fixtures/store.js';
-import getLogger from '../../test/fixtures/no-logger.js';
-import getApp from '../app.js';
-import type TestAgent from 'supertest/lib/agent.d.ts';
-
-async function getSetup() {
-    const stores = createStores();
-    const config = createTestConfig();
-    const services = createServices(stores, config);
-    const app = await getApp(config, stores, services);
-
-    return {
-        request: supertest(app),
-        stores,
-    };
-}
-let request: TestAgent<Test>;
-beforeEach(async () => {
-    const setup = await getSetup();
-    request = setup.request;
-});
-
-afterEach(() => {
-    getLogger.setMuteError(false);
-});
-
-test('should give 200 when ready', async () => {
-    await request.get('/ready').expect(200);
-});
-
-test('should give health=GOOD when ready', async () => {
-    expect.assertions(2);
-    await request
-        .get('/ready')
-        .expect(200)
-        .expect((res) => {
-            expect(res.status).toBe(200);
-            expect(res.body.health).toBe('GOOD');
-        });
-});

src/lib/routes/ready-check.ts

@@ -1,4 +1,5 @@
 import type { Request, Response } from 'express';
+import type { PoolClient } from 'pg';
 import type { IUnleashConfig } from '../types/option.js';
 import type { IUnleashServices } from '../services/index.js';
 import type { Db } from '../db/db.js';
@@ -8,6 +9,7 @@ import Controller from './controller.js';
 import { NONE } from '../types/permissions.js';
 import { createResponseSchema } from '../openapi/util/create-response-schema.js';
 import type { ReadyCheckSchema } from '../openapi/spec/ready-check-schema.js';
+import { emptyResponse, parseEnvVarNumber } from '../server-impl.js';
 
 export class ReadyCheckController extends Controller {
     private logger: Logger;
@@ -37,7 +39,7 @@ export class ReadyCheckController extends Controller {
                         'This operation returns information about whether this Unleash instance is ready to serve requests or not. Typically used by your deployment orchestrator (e.g. Kubernetes, Docker Swarm, Mesos, et al.).',
                     responses: {
                         200: createResponseSchema('readyCheckSchema'),
-                        500: createResponseSchema('readyCheckSchema'),
+                        503: emptyResponse,
                     },
                 }),
             ],
@@ -47,16 +49,53 @@ export class ReadyCheckController extends Controller {
     async getReady(_: Request, res: Response<ReadyCheckSchema>): Promise<void> {
         if (this.config.checkDbOnReady && this.db) {
             try {
-                await this.db.raw('select 1');
+                const timeoutMs = parseEnvVarNumber(
+                    process.env.DATABASE_STATEMENT_TIMEOUT_MS,
+                    200,
+                );
+
+                await this.runReadinessQuery(timeoutMs);
                 res.status(200).json({ health: 'GOOD' });
                 return;
             } catch (err: any) {
                 this.logger.warn('Database readiness check failed', err);
-                res.status(500).json({ health: 'BAD' });
+                res.status(503).end();
                 return;
             }
         }
 
         res.status(200).json({ health: 'GOOD' });
     }
+
+    private async runReadinessQuery(timeoutMs: number): Promise<void> {
+        const client = getKnexClient(this.db!);
+        const pendingAcquire = client.pool.acquire();
+
+        const abortDelay = setTimeout(() => pendingAcquire.abort(), timeoutMs);
+        let connection: PoolClient | undefined;
+
+        try {
+            connection = (await pendingAcquire.promise) as PoolClient;
+        } finally {
+            clearTimeout(abortDelay);
+        }
+
+        try {
+            connection.query({
+                text: 'SELECT 1',
+            });
+        } finally {
+            if (connection) {
+                await client.releaseConnection(connection);
+            }
+        }
+    }
+}
+
+function getKnexClient(db: Db): Db['client'] {
+    if (db.client?.pool && typeof db.client.releaseConnection === 'function') {
+        return db.client as Db['client'];
+    }
+
+    throw new Error('Unsupported database handle for readiness check');
 }

src/test/e2e/ready.e2e.test.ts

@@ -1,27 +1,62 @@
-import { setupApp } from './helpers/test-helper.js';
+import { setupAppWithCustomConfig } from './helpers/test-helper.js';
 import dbInit, { type ITestDb } from './helpers/database-init.js';
-import getLogger from '../fixtures/no-logger.js';
-import type { IUnleashStores } from '../../lib/types/index.js';
 
-let stores: IUnleashStores;
 let db: ITestDb;
 
-beforeAll(async () => {
-    db = await dbInit('ready_api', getLogger);
-    stores = db.stores;
+describe('DB is up', () => {
+    beforeAll(async () => {
+        db = await dbInit();
+    });
+
+    test('when checkDb is disabled, returns ready', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            undefined,
+            db.rawDatabase,
+        );
+        await request
+            .get('/ready')
+            .expect('Content-Type', /json/)
+            .expect(200)
+            .expect('{"health":"GOOD"}');
+    });
+
+    test('when checkDb is enabled, returns ready', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            { checkDbOnReady: true },
+            db.rawDatabase,
+        );
+        await request.get('/ready').expect(200).expect('{"health":"GOOD"}');
+    });
 });
 
-afterAll(async () => {
-    await db.destroy();
-});
+describe('DB is down', () => {
+    beforeAll(async () => {
+        db = await dbInit();
+    });
 
-test('returns readiness good', async () => {
-    expect.assertions(0);
-    const { request, destroy } = await setupApp(stores);
-    await request
-        .get('/ready')
-        .expect('Content-Type', /json/)
-        .expect(200)
-        .expect('{"health":"GOOD"}');
-    await destroy();
+    test('when checkDb is disabled, returns readiness good', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            undefined,
+            db.rawDatabase,
+        );
+        await db.destroy();
+        await request
+            .get('/ready')
+            .expect('Content-Type', /json/)
+            .expect(200)
+            .expect('{"health":"GOOD"}');
+    });
+
+    test('when checkDb is enabled, fails readiness check', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            { checkDbOnReady: true },
+            db.rawDatabase,
+        );
+        await db.destroy();
+        await request.get('/ready').expect(503);
+    });
 });