feat: add beforePoolAcquire and afterPoolAcquire hooks (#15874)

Co-authored-by: Filip Havel <filip.havel@goodrequest.com>
Co-authored-by: Rik Smale <13023439+WikiRik@users.noreply.github.com>

src/dialects/abstract/connection-manager.js

@@ -283,7 +283,13 @@ class ConnectionManager {
     let result;
 
     try {
+
+      await this.sequelize.runHooks('beforePoolAcquire', options);
+
       result = await this.pool.acquire(options.type, options.useMaster);
+
+      await this.sequelize.runHooks('afterPoolAcquire', result, options);
+
     } catch (error) {
       if (error instanceof TimeoutError) throw new errors.ConnectionAcquireTimeoutError(error);
       throw error;

src/hooks.d.ts

@@ -15,6 +15,7 @@ import { AbstractQuery } from './dialects/abstract/query';
 import { QueryOptions } from './dialects/abstract/query-interface';
 import { Config, Options, Sequelize, SyncOptions } from './sequelize';
 import { DeepWriteable } from './utils';
+import { Connection, GetConnectionOptions } from './dialects/abstract/connection-manager';
 
 export type HookReturn = Promise<void> | void;
 
@@ -76,6 +77,8 @@ export interface SequelizeHooks<
   afterInit(sequelize: Sequelize): void;
   beforeConnect(config: DeepWriteable<Config>): HookReturn;
   afterConnect(connection: unknown, config: Config): HookReturn;
+  beforePoolAcquire(config: GetConnectionOptions): HookReturn;
+  afterPoolAcquire(connection: Connection, config: GetConnectionOptions): HookReturn;
   beforeDisconnect(connection: unknown): HookReturn;
   afterDisconnect(connection: unknown): HookReturn;
 }

src/hooks.js

@@ -43,6 +43,8 @@ const hookTypes = {
   afterConnect: { params: 2, noModel: true },
   beforeDisconnect: { params: 1, noModel: true },
   afterDisconnect: { params: 1, noModel: true },
+  beforePoolAcquire: { params: 1, noModel: true },
+  afterPoolAcquire: { params: 2, noModel: true },
   beforeSync: { params: 1 },
   afterSync: { params: 1 },
   beforeBulkSync: { params: 1 },
@@ -541,6 +543,24 @@ exports.applyTo = applyTo;
  * @memberof Sequelize
  */
 
+/**
+ *  A hook that is run before a connection to the pool
+ *
+ * @param {string}   name
+ * @param {Function} fn   A callback function that is called with config passed to connection
+ * @name beforePoolAcquire
+ * @memberof Sequelize
+ */
+
+/**
+ * A hook that is run after a connection to the pool
+ *
+ * @param {string}   name
+ * @param {Function} fn   A callback function that is called with the connection object and the config passed to connection
+ * @name afterPoolAcquire
+ * @memberof Sequelize
+ */
+
 /**
  * A hook that is run before a connection is disconnected
  *

src/sequelize.d.ts

@@ -30,7 +30,7 @@ import QueryTypes = require('./query-types');
 import { Transaction, TransactionOptions } from './transaction';
 import { Op } from './index';
 import { Cast, Col, DeepWriteable, Fn, Json, Literal, Where } from './utils';
-import { ConnectionManager } from './dialects/abstract/connection-manager';
+import { Connection, ConnectionManager, GetConnectionOptions } from './dialects/abstract/connection-manager';
 
 export type RetryOptions = RetryAsPromisedOptions;
 
@@ -720,6 +720,26 @@ export class Sequelize extends Hooks {
   public static afterDisconnect(name: string, fn: (connection: unknown) => void): void;
   public static afterDisconnect(fn: (connection: unknown) => void): void;
 
+
+  /**
+   * A hook that is run before attempting to acquire a connection from the pool
+   *
+   * @param name
+   * @param fn   A callback function that is called with options
+   */
+  public static beforePoolAcquire(name: string, fn: (options: GetConnectionOptions) => void): void;
+  public static beforePoolAcquire(fn: (options: GetConnectionOptions) => void): void;
+
+  /**
+   * A hook that is run after successfully acquiring a connection from the pool
+   *
+   * @param name
+   * @param fn   A callback function that is called with options
+   */
+  public static afterPoolAcquire(name: string, fn: (connection: Connection, options: GetConnectionOptions) => void): void;
+  public static afterPoolAcquire(fn: (connection: Connection, options: GetConnectionOptions) => void): void;
+
+
   /**
    * A hook that is run before a find (select) query, after any { include: {all: ...} } options are expanded
    *

test/integration/hooks/hooks.test.js

@@ -403,4 +403,24 @@ describe(Support.getTestDialectTeaser('Hooks'), () => {
       expect(narutoHook).to.have.been.calledTwice;
     });
   });
+
+  describe('Sequelize hooks', () => {
+    it('should run before/afterPoolAcquire hooks', async function() {
+      if (dialect === 'sqlite') {
+        return this.skip();
+      }
+
+      const beforeHook = sinon.spy();
+      const afterHook = sinon.spy();
+
+      this.sequelize.addHook('beforePoolAcquire', beforeHook);
+      this.sequelize.addHook('afterPoolAcquire', afterHook);
+
+      await this.sequelize.authenticate();
+
+      expect(beforeHook).to.have.been.calledOnce;
+      expect(afterHook).to.have.been.calledOnce;
+
+    });
+  });
 });

test/types/hooks.ts

@@ -1,5 +1,6 @@
 import { expectTypeOf } from "expect-type";
 import { FindOptions, Model, QueryOptions, SaveOptions, Sequelize, UpsertOptions, Config, Utils } from "sequelize";
+import { Connection, GetConnectionOptions } from "sequelize/lib/dialects/abstract/connection-manager";
 import { ModelHooks } from "sequelize/lib/hooks";
 import { AbstractQuery } from "sequelize/lib/query";
 import { SemiDeepWritable } from "./type-helpers/deep-writable";
@@ -85,4 +86,29 @@ import { SemiDeepWritable } from "./type-helpers/deep-writable";
   Sequelize.beforeConnect('name', config => expectTypeOf(config).toEqualTypeOf<Utils.DeepWriteable<Config>>());
   Sequelize.beforeConnect(config => expectTypeOf(config).toEqualTypeOf<Utils.DeepWriteable<Config>>());
   Sequelize.addHook('beforeConnect', (...args) => { expectTypeOf(args).toEqualTypeOf<[Utils.DeepWriteable<Config>]>(); });
+  Sequelize.beforePoolAcquire('name', (options?: GetConnectionOptions) => {
+    expectTypeOf(options).toMatchTypeOf<GetConnectionOptions | undefined>();
+  });
+
+  Sequelize.beforePoolAcquire((options?: GetConnectionOptions) => {
+    expectTypeOf(options).toMatchTypeOf<GetConnectionOptions | undefined>();
+  });
+
+  Sequelize.addHook('beforePoolAcquire', (...args: [GetConnectionOptions | undefined]) => {
+    expectTypeOf(args).toMatchTypeOf<[GetConnectionOptions | undefined]>();
+  });
+
+  Sequelize.afterPoolAcquire('name', (connection: Connection, options?: GetConnectionOptions) => {
+    expectTypeOf(connection).toMatchTypeOf<Connection>();
+    expectTypeOf(options).toMatchTypeOf<GetConnectionOptions | undefined>();
+  });
+
+  Sequelize.afterPoolAcquire((connection: Connection, options?: GetConnectionOptions) => {
+    expectTypeOf(connection).toMatchTypeOf<Connection>();
+    expectTypeOf(options).toMatchTypeOf<GetConnectionOptions | undefined>();
+  });
+
+  Sequelize.addHook('afterPoolAcquire', (...args: [Connection | GetConnectionOptions | undefined]) => {
+    expectTypeOf(args).toMatchTypeOf<[Connection | GetConnectionOptions | undefined]>();
+  });
 }

test/unit/esm-named-exports.test.js

@@ -91,6 +91,8 @@ describe('ESM module', () => {
       'afterAssociate',
       'beforeConnect',
       'afterConnect',
+      'beforePoolAcquire',
+      'afterPoolAcquire',
       'beforeDisconnect',
       'afterDisconnect',
       'beforeSync',