diff --git a/lib/core.ts b/lib/core.ts index 5a34bae..e076748 100644 --- a/lib/core.ts +++ b/lib/core.ts @@ -1,6 +1,6 @@ import type {Knex} from 'knex'; import _ from 'lodash'; -import {Model} from 'objection'; +import {type Constructor, Model} from 'objection'; import {defaultDispatcherOptions, defaultExLogger, defaultKnexOptions} from './constants'; import {PGDispatcher} from './dispatcher'; @@ -13,16 +13,22 @@ export interface CoreDBDispatcherOptions { beforeTerminate?: () => Promise; } +export type GetModelParams = {cancelOnTimeout?: boolean; useLimitInFirst?: boolean}; + export interface CoreDBConstructorArgs { connectionString: string; dispatcherOptions?: CoreDBDispatcherOptions; knexOptions?: Knex.Config; logger?: ExLogger; + modelParams?: GetModelParams; } -export function getModel(): typeof BaseModel { +export function getModel(params: GetModelParams = {}): typeof BaseModel { let _db: PGDispatcher; + const cancelOnTimeout = Boolean(params.cancelOnTimeout); + const useLimitInFirst = Boolean(params.useLimitInFirst); + class CoreBaseModel extends Model { static set db(value: PGDispatcher) { if (!_db) { @@ -45,6 +51,28 @@ export function getModel(): typeof BaseModel { get replica() { return _db.replica; } + + static query( + this: Constructor, + ...args: Parameters> + ): ReturnType> { + const query = super.query(...args); + + if (cancelOnTimeout) { + const originalTimeout = query.timeout; + + query.timeout = (ms, options) => { + const optionsWithCancel = {cancel: true, ...(options ?? {})}; + return originalTimeout.apply(query, [ms, optionsWithCancel]); + }; + } + + return query; + } + + static get useLimitInFirst() { + return useLimitInFirst; + } } return CoreBaseModel; @@ -55,6 +83,7 @@ export function initDB({ dispatcherOptions, knexOptions = {}, logger = defaultExLogger, + modelParams, }: CoreDBConstructorArgs) { if (!connectionString) { throw new Error('Empty connection string'); @@ -83,7 +112,7 @@ export function initDB({ process.on('SIGINT', terminate); - const CoreBaseModel = getModel(); + const CoreBaseModel = getModel(modelParams); CoreBaseModel.db = db; const helpers = {