feat: add beforePoolAcquire and afterPoolAcquire hooks. (#15859)

sequelize · Mar 31, 2023 · 6711351 · 6711351
1 parent 24000e9
commit 6711351
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 1 deletion.
diff --git a/packages/core/src/dialects/abstract/connection-manager.ts b/packages/core/src/dialects/abstract/connection-manager.ts
@@ -175,8 +175,13 @@ export class AbstractConnectionManager<TConnection extends Connection = Connecti
     await this._initDatabaseVersion();
 
     try {
+
+      await this.sequelize.hooks.runAsync('beforePoolAcquire', options);
+
       const result = await this.pool.acquire(options?.type, options?.useMaster);
 
+      await this.sequelize.hooks.runAsync('afterPoolAcquire', result, options);
+
       debug('connection acquired');
 
       return result;

diff --git a/packages/core/src/sequelize-typescript.ts b/packages/core/src/sequelize-typescript.ts
@@ -62,6 +62,16 @@ export interface SequelizeHooks extends ModelHooks {
    * A hook that is run at the end of {@link Sequelize#sync}
    */
   afterBulkSync(options: SyncOptions): AsyncHookReturn;
+
+  /**
+   * A hook that is run before a connection to the pool
+   */
+  beforePoolAcquire(options?: GetConnectionOptions): AsyncHookReturn;
+
+  /**
+   * A hook that is run after a connection to the pool
+   */
+  afterPoolAcquire(connection: Connection, options?: GetConnectionOptions): AsyncHookReturn;
 }
 
 export interface StaticSequelizeHooks {
@@ -95,6 +105,7 @@ const instanceSequelizeHooks = new HookHandlerBuilder<SequelizeHooks>([
   'beforeConnect', 'afterConnect',
   'beforeDisconnect', 'afterDisconnect',
   'beforeDefine', 'afterDefine',
+  'beforePoolAcquire', 'afterPoolAcquire',
   ...validModelHooks,
 ]);
 
@@ -147,6 +158,9 @@ export abstract class SequelizeTypeScript {
   beforeDefine = legacyBuildAddHook(instanceSequelizeHooks, 'beforeDefine');
   afterDefine = legacyBuildAddHook(instanceSequelizeHooks, 'afterDefine');
 
+  beforePoolAcquire = legacyBuildAddHook(instanceSequelizeHooks, 'beforePoolAcquire');
+  afterPoolAcquire = legacyBuildAddHook(instanceSequelizeHooks, 'afterPoolAcquire');
+
   beforeValidate = legacyBuildAddHook(instanceSequelizeHooks, 'beforeValidate');
   afterValidate = legacyBuildAddHook(instanceSequelizeHooks, 'afterValidate');
   validationFailed = legacyBuildAddHook(instanceSequelizeHooks, 'validationFailed');

diff --git a/packages/core/test/integration/hooks/hooks.test.js b/packages/core/test/integration/hooks/hooks.test.js
@@ -355,4 +355,25 @@ describe(Support.getTestDialectTeaser('Hooks'), () => {
       });
     });
   });
+
+  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;
+
+    });
+  });
+
 });
diff --git a/packages/core/test/types/hooks.ts b/packages/core/test/types/hooks.ts
@@ -12,6 +12,7 @@ import type {
   AfterAssociateEventData,
   AssociationOptions,
 } from '@sequelize/core/_non-semver-use-at-your-own-risk_/associations';
+import type { Connection, GetConnectionOptions } from '@sequelize/core/_non-semver-use-at-your-own-risk_/dialects/abstract/connection-manager.js';
 import type { AbstractQuery } from '@sequelize/core/_non-semver-use-at-your-own-risk_/dialects/abstract/query.js';
 import type { ValidationOptions } from '@sequelize/core/_non-semver-use-at-your-own-risk_/instance-validator';
 import type { ModelHooks } from '@sequelize/core/_non-semver-use-at-your-own-risk_/model-hooks.js';
@@ -158,6 +159,32 @@ sequelize.addHook('beforeConnect', (...args) => {
   expectTypeOf(args).toMatchTypeOf<[ConnectionOptions]>();
 });
 
+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]>();
+});
+
 sequelize.beforeQuery((options, query) => {
   expectTypeOf(options).toEqualTypeOf<QueryOptions>();
   expectTypeOf(query).toEqualTypeOf<AbstractQuery>();

diff --git a/packages/core/test/unit/hooks.test.js b/packages/core/test/unit/hooks.test.js
@@ -16,7 +16,7 @@ describe(Support.getTestDialectTeaser('Hooks'), () => {
   });
 
   it('does not expose non-model hooks', function () {
-    for (const badHook of ['beforeDefine', 'afterDefine', 'beforeConnect', 'afterConnect', 'beforeDisconnect', 'afterDisconnect', 'beforeInit', 'afterInit']) {
+    for (const badHook of ['beforeDefine', 'afterDefine', 'beforeConnect', 'afterConnect', 'beforePoolAcquire', 'afterPoolAcquire', 'beforeDisconnect', 'afterDisconnect', 'beforeInit', 'afterInit']) {
       expect(this.Model).to.not.have.property(badHook);
     }
   });