Add migration lifecycle hooks (#5541)

Co-authored-by: Raz Luvaton <16746759+rluvaton@users.noreply.github.com>
knex · Jan 13, 2024 · 3764ded · 3764ded
1 parent 4ca3dd5
commit 3764ded
Show file tree

Hide file tree

Showing 3 changed files with 377 additions and 22 deletions.
diff --git a/lib/migrations/migrate/Migrator.js b/lib/migrations/migrate/Migrator.js
@@ -379,12 +379,18 @@ class Migrator {
 
       let batchNo = await this._latestBatchNumber(trx);
       if (direction === 'up') batchNo++;
-      const res = await this._waterfallBatch(
-        batchNo,
-        migrations,
-        direction,
-        trx
-      );
+
+      // Run any hooks before/after this batch
+      const beforeAll = this.config.beforeAll || (() => {});
+      const afterAll = this.config.afterAll || (() => {});
+
+      let res = [];
+      if (migrations.length > 0) {
+        await beforeAll(trx || this.knex, migrations);
+        res = await this._waterfallBatch(batchNo, migrations, direction, trx);
+        await afterAll(trx || this.knex, migrations);
+      }
+
       await this._freeLock(canGetLockInTransaction ? trx : undefined);
       return res;
     } catch (error) {
@@ -496,30 +502,40 @@ class Migrator {
       const migrationContent =
         this.config.migrationSource.getMigration(migration);
 
+      const beforeEach = this.config.beforeEach || (() => {});
+      const afterEach = this.config.afterEach || (() => {});
+
       // We're going to run each of the migrations in the current "up".
       current = current
         .then(async () => await migrationContent) //maybe promise
-        .then((migrationContent) => {
+        .then(async (migrationContent) => {
           this._activeMigration.fileName = name;
           if (
             !trx &&
             this._useTransaction(migrationContent, disableTransactions)
           ) {
             this.knex.enableProcessing();
-            return this._transaction(
-              this.knex,
-              migrationContent,
-              direction,
-              name
-            );
+            return await this.knex.transaction(async (trx) => {
+              await beforeEach(trx, [migration]);
+              const migrationResult = await checkPromise(
+                this.knex.client.logger,
+                migrationContent[direction](trx),
+                name
+              );
+              await afterEach(trx, [migration]);
+              return migrationResult;
+            });
           }
 
           trxOrKnex.enableProcessing();
-          return checkPromise(
+          await beforeEach(trxOrKnex, [migration]);
+          const migrationResult = await checkPromise(
             this.knex.client.logger,
             migrationContent[direction](trxOrKnex),
             name
           );
+          await afterEach(trxOrKnex, [migration]);
+          return migrationResult;
         })
         .then(() => {
           trxOrKnex.disableProcessing();
@@ -584,12 +600,9 @@ function getNewMigrations(migrationSource, all, completed) {
   });
 }
 
-function checkPromise(logger, migrationPromise, name, commitFn) {
+function checkPromise(logger, migrationPromise, name) {
   if (!migrationPromise || typeof migrationPromise.then !== 'function') {
     logger.warn(`migration ${name} did not return a promise`);
-    if (commitFn) {
-      commitFn();
-    }
   }
   return migrationPromise;
 }