feat: add proxy topology mode for SPQR-compatible routing

README.md

@@ -58,13 +58,25 @@ export const {db, CoreBaseModel, helpers} = initDB({
 
 - `connectionString` is a set of [postgres connection strings](https://stackoverflow.com/questions/3582552/postgresql-connection-url) separated by a comma, with at least one host required, for example: `'postgresql://user:password@dbHost1:5432/dbName,postgresql://user:password@dbHost2:5432/dbName'`
 - `logger`: Provide the `info` and `error` callbacks for logging messages
-- `dispatcherOptions`: Settings for the primary/replica balancing (none of the options are required)
+- `dispatcherOptions`: Settings for connection balancing (none of the options are required)
   - `healthcheckInterval`: Health check interval in milliseconds, default value: 5000ms
   - `healthcheckTimeout`: Health check interval in milliseconds, default value: 700ms
-  - `suppressStatusLogs`: Boolean value that disables database health checks (useful for developers)
+  - `suppressStatusLogs`: Boolean value that disables database health-check status logs (useful for developers)
   - `beforeTerminate`: Function called before terminating a connection, must return a Promise
+  - `topologyMode`: Connection topology, either `primary-replica` (the default) or `proxy`
 - `knexOptions`: Non-required additional options that will be passed to Knex before initialization
 
+When all connection strings point to equivalent proxy or router instances, such as SPQR routers, use `proxy` mode. Healthy endpoints are then eligible for both primary and replica queries, and the endpoint with the lowest latest health-check latency is selected:
+
+```typescript
+initDB({
+  connectionString: process.env.POSTGRES_DSN_LIST,
+  dispatcherOptions: {
+    topologyMode: 'proxy',
+  },
+});
+```
+
 Here is the recommended project structure (we only list the directories that have to do with working with the database):
 
 ```
@@ -82,7 +94,7 @@ The `initDB` constructor exports three elements: `db`, `CoreBaseModel`, and `hel
 
 ### db
 
-`db`: Instance of the [PGDispatcher](https://github.com/gravity-ui/postgreskit/blob/main/lib/dispatcher.ts) module responsible for primary/replica connection balancing. Under the hood, this module creates N instances of knex (N is the number of hosts passed in `connectionString`). From these instances, it polls the database hosts every `healthcheckInterval` milliseconds, requesting if they are primary hosts or replica hosts (using the `SELECT pg_is_in_recovery()` query).
+`db`: Instance of the [PGDispatcher](https://github.com/gravity-ui/postgreskit/blob/main/lib/dispatcher.ts) module responsible for connection balancing. Under the hood, this module creates N instances of knex (N is the number of hosts passed in `connectionString`). By default, it polls the database hosts every `healthcheckInterval` milliseconds, requesting if they are primary hosts or replica hosts (using the `SELECT pg_is_in_recovery()` query). In `proxy` topology mode, it uses `SELECT 1` instead and routes both `db.primary` and `db.replica` to the fastest healthy endpoint.
 
 Public `db` methods:
 

jest/unit.config.js

@@ -0,0 +1,6 @@
+module.exports = {
+    clearMocks: true,
+    rootDir: '..',
+    testEnvironment: 'node',
+    testMatch: ['<rootDir>/tests/**/*.test.js'],
+};

lib/constants.ts

@@ -20,6 +20,7 @@ export const defaultDispatcherOptions: PDOptions = {
     healthcheckTimeout: 700,
     suppressStatusLogs: false,
     beforeTerminate: () => Promise.resolve(),
+    topologyMode: 'primary-replica',
 };
 
 export const defaultExLogger: ExLogger = {

lib/core.ts

@@ -4,13 +4,16 @@ import {type Constructor, Model} from 'objection';
 
 import {defaultDispatcherOptions, defaultExLogger, defaultKnexOptions} from './constants';
 import {PGDispatcher} from './dispatcher';
-import type {BaseModel, ExLogger} from './types';
+import type {BaseModel, ExLogger, TopologyMode} from './types';
+
+export type {TopologyMode} from './types';
 
 export interface CoreDBDispatcherOptions {
     healthcheckInterval?: number;
     healthcheckTimeout?: number;
     suppressStatusLogs?: boolean;
     beforeTerminate?: () => Promise<void>;
+    topologyMode?: TopologyMode;
 }
 
 export type GetModelParams = {cancelOnTimeout?: boolean; useLimitInFirst?: boolean};

lib/dispatcher.ts

@@ -138,6 +138,10 @@ export class PGDispatcher {
     get primary() {
         this.checkConnectionsAvailability();
 
+        if (this.options.topologyMode === 'proxy') {
+            return this.fastestHealthyConnection.knex;
+        }
+
         const primaryConnections = this.connections.filter((c) => c.primary);
         if (primaryConnections.length > 1) {
             this.logger.error({
@@ -160,6 +164,10 @@ export class PGDispatcher {
     get replica() {
         this.checkConnectionsAvailability();
 
+        if (this.options.topologyMode === 'proxy') {
+            return this.fastestHealthyConnection.knex;
+        }
+
         const replicaConnections = this.healthyConnections.filter((c) => !c.primary);
 
         if (replicaConnections.length) {
@@ -213,6 +221,12 @@ export class PGDispatcher {
     }
 
     private async performCheckupQuery(knex: Knex): Promise<PDCheckupResult> {
+        if (this.options.topologyMode === 'proxy') {
+            await knex.raw('SELECT 1;').timeout(this.options.healthcheckTimeout);
+
+            return {pingOk: true, primary: false};
+        }
+
         const result = await knex
             .raw('SELECT pg_is_in_recovery();')
             .timeout(this.options.healthcheckTimeout);
@@ -278,6 +292,10 @@ export class PGDispatcher {
         return this.connections.filter((c) => c.healthy);
     }
 
+    private get fastestHealthyConnection() {
+        return this.healthyConnections.sort((a, b) => a.latency - b.latency)[0];
+    }
+
     private checkConnectionsAvailability() {
         if (!this.healthyConnections.length) {
             const error = new PDError('No connections available');

lib/types.ts

@@ -1,11 +1,14 @@
 import type {Model} from 'objection';
 import type {PGDispatcher} from './dispatcher';
 
+export type TopologyMode = 'primary-replica' | 'proxy';
+
 export interface PDOptions {
     healthcheckInterval: number;
     healthcheckTimeout: number;
     suppressStatusLogs: boolean;
     beforeTerminate: () => Promise<void>;
+    topologyMode: TopologyMode;
 }
 
 export type Dict = {[key: string]: unknown};

tests/dispatcher.test.js

@@ -0,0 +1,122 @@
+/* eslint-env jest */
+
+jest.mock('knex', () => jest.fn());
+
+const knexBuilder = require('knex');
+
+const {defaultDispatcherOptions} = require('../build/constants');
+const {PGDispatcher} = require('../build/dispatcher');
+
+const activeDispatchers = [];
+
+function createKnex(checkup) {
+    const timeout = jest.fn(() => checkup());
+
+    return {
+        destroy: jest.fn(() => Promise.resolve()),
+        raw: jest.fn(() => ({timeout})),
+        timeout,
+    };
+}
+
+function successfulCheckup(row, latency = 0) {
+    return () =>
+        new Promise((resolve) => {
+            setTimeout(() => resolve({rows: [row]}), latency);
+        });
+}
+
+function failedCheckup() {
+    return Promise.reject(new Error('Proxy unavailable'));
+}
+
+function createDispatcher(clients, options = {}) {
+    clients.forEach((client) => knexBuilder.mockImplementationOnce(() => client));
+
+    const logger = {
+        error: jest.fn(),
+        info: jest.fn(),
+    };
+    const dispatcher = new PGDispatcher({
+        connections: clients.map(
+            (_, index) => `postgresql://user:password@database-${index}.example/db`,
+        ),
+        logger,
+        options: {
+            ...defaultDispatcherOptions,
+            healthcheckInterval: 60_000,
+            healthcheckTimeout: 100,
+            ...options,
+        },
+    });
+
+    activeDispatchers.push(dispatcher);
+
+    return {dispatcher, logger};
+}
+
+function loggedErrorMessages(logger) {
+    return logger.error.mock.calls.map(([, error]) => error.message);
+}
+
+afterEach(async () => {
+    await Promise.all(activeDispatchers.splice(0).map((dispatcher) => dispatcher.terminate()));
+});
+
+describe('PGDispatcher topology modes', () => {
+    test('the default mode keeps the primary/replica health check and routing behavior', async () => {
+        const primary = createKnex(successfulCheckup({pg_is_in_recovery: false}));
+        const replica = createKnex(successfulCheckup({pg_is_in_recovery: true}));
+        const {dispatcher} = createDispatcher([primary, replica]);
+
+        await dispatcher.ready();
+
+        expect(defaultDispatcherOptions.topologyMode).toBe('primary-replica');
+        expect(primary.raw).toHaveBeenCalledWith('SELECT pg_is_in_recovery();');
+        expect(replica.raw).toHaveBeenCalledWith('SELECT pg_is_in_recovery();');
+        expect(dispatcher.primary).toBe(primary);
+        expect(dispatcher.replica).toBe(replica);
+    });
+
+    test('proxy mode routes both roles to the fastest healthy endpoint without topology warnings', async () => {
+        const slowerProxy = createKnex(successfulCheckup({value: 1}, 30));
+        const fasterProxy = createKnex(successfulCheckup({value: 1}, 5));
+        const unhealthyProxy = createKnex(failedCheckup);
+        const {dispatcher, logger} = createDispatcher([slowerProxy, fasterProxy, unhealthyProxy], {
+            topologyMode: 'proxy',
+        });
+
+        await dispatcher.ready();
+
+        for (const proxy of [slowerProxy, fasterProxy, unhealthyProxy]) {
+            expect(proxy.raw).toHaveBeenCalledWith('SELECT 1;');
+            expect(proxy.raw).not.toHaveBeenCalledWith('SELECT pg_is_in_recovery();');
+        }
+        expect(dispatcher.primary).toBe(fasterProxy);
+        expect(dispatcher.replica).toBe(fasterProxy);
+        expect(loggedErrorMessages(logger)).not.toEqual(
+            expect.arrayContaining([
+                'Multiple primary connections detected, something is wrong',
+                'No alive replica available, using master for read',
+            ]),
+        );
+    });
+
+    test('the existing unavailable-database error is preserved when all proxies are unhealthy', async () => {
+        const {dispatcher} = createDispatcher(
+            [createKnex(failedCheckup), createKnex(failedCheckup)],
+            {topologyMode: 'proxy'},
+        );
+
+        await dispatcher.ready();
+
+        for (const connection of ['primary', 'replica']) {
+            expect(() => dispatcher[connection]).toThrow(
+                expect.objectContaining({
+                    code: 'ERR_DB_NOT_AVAILABLE',
+                    message: 'No connections available',
+                }),
+            );
+        }
+    });
+});