feat(migrations): add Migration.execute() method

This method can be used to execute queries inside the same transaction as the queries added via `Migration.addSql()` method. Closes #770
mikro-orm · Aug 24, 2020 · 5c1f60a · 5c1f60a
1 parent b357686
commit 5c1f60a
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 20 deletions.
diff --git a/packages/knex/src/AbstractSqlConnection.ts b/packages/knex/src/AbstractSqlConnection.ts
@@ -24,7 +24,7 @@ export abstract class AbstractSqlConnection extends Connection {
     }
   }
 
-  async transactional<T>(cb: (trx: Transaction) => Promise<T>, ctx?: Transaction): Promise<T> {
+  async transactional<T>(cb: (trx: Transaction<KnexTransaction>) => Promise<T>, ctx?: Transaction<KnexTransaction>): Promise<T> {
     return (ctx || this.client).transaction(cb);
   }
 

diff --git a/packages/migrations/src/Migration.ts b/packages/migrations/src/Migration.ts
@@ -1,9 +1,10 @@
-import { Configuration } from '@mikro-orm/core';
-import { AbstractSqlDriver } from '@mikro-orm/knex';
+import { Configuration, Transaction } from '@mikro-orm/core';
+import { AbstractSqlDriver, Knex } from '@mikro-orm/knex';
 
 export abstract class Migration {
 
   private readonly queries: string[] = [];
+  protected ctx?: Transaction<Knex.Transaction>;
 
   constructor(protected readonly driver: AbstractSqlDriver,
               protected readonly config: Configuration) { }
@@ -24,6 +25,15 @@ export abstract class Migration {
 
   reset(): void {
     this.queries.length = 0;
+    this.ctx = undefined;
+  }
+
+  setTransactionContext(ctx: Transaction): void {
+    this.ctx = ctx;
+  }
+
+  async execute(sql: string) {
+    return this.driver.execute(sql, undefined, 'all', this.ctx);
   }
 
   getQueries(): string[] {

diff --git a/packages/migrations/src/MigrationRunner.ts b/packages/migrations/src/MigrationRunner.ts
@@ -1,4 +1,4 @@
-import { Configuration, MigrationsOptions, Utils, Transaction } from '@mikro-orm/core';
+import { Configuration, MigrationsOptions, Transaction, Utils } from '@mikro-orm/core';
 import { AbstractSqlDriver } from '@mikro-orm/knex';
 import { Migration } from './Migration';
 
@@ -14,25 +14,17 @@ export class MigrationRunner {
 
   async run(migration: Migration, method: 'up' | 'down'): Promise<void> {
     migration.reset();
-    await migration[method]();
-    let queries = migration.getQueries();
-
-    if (this.options.disableForeignKeys) {
-      const charset = this.config.get('charset')!;
-      queries.unshift(...this.helper.getSchemaBeginning(charset).split('\n'));
-      queries.push(...this.helper.getSchemaEnd().split('\n'));
-    }
-
-    queries = queries.filter(sql => sql.trim().length > 0);
 
     if (!this.options.transactional || !migration.isTransactional()) {
+      const queries = await this.getQueries(migration, method);
       await Utils.runSerial(queries, sql => this.driver.execute(sql));
-      return;
+    } else {
+      await this.connection.transactional(async tx => {
+        migration.setTransactionContext(tx);
+        const queries = await this.getQueries(migration, method);
+        await Utils.runSerial(queries, sql => this.driver.execute(tx.raw(sql)));
+      }, this.masterTransaction);
     }
-
-    await this.connection.transactional(async tx => {
-      await Utils.runSerial(queries, sql => this.driver.execute(tx.raw(sql)));
-    }, this.masterTransaction);
   }
 
   setMasterMigration(trx: Transaction) {
@@ -43,4 +35,18 @@ export class MigrationRunner {
     delete this.masterTransaction;
   }
 
+  private async getQueries(migration: Migration, method: 'up' | 'down') {
+    await migration[method]();
+    let queries = migration.getQueries();
+
+    if (this.options.disableForeignKeys) {
+      const charset = this.config.get('charset')!;
+      queries.unshift(...this.helper.getSchemaBeginning(charset).split('\n'));
+      queries.push(...this.helper.getSchemaEnd().split('\n'));
+    }
+
+    queries = queries.filter(sql => sql.trim().length > 0);
+    return queries;
+  }
+
 }
diff --git a/tests/Migrator.test.ts b/tests/Migrator.test.ts
@@ -18,6 +18,8 @@ class MigrationTest2 extends Migration {
 
   async up(): Promise<void> {
     this.addSql('select 1 + 1');
+    const res = await this.execute('select 1 + 1 as count');
+    expect(res).toEqual([{ count: 2 }]);
   }
 
   isTransactional(): boolean {
@@ -147,8 +149,9 @@ describe('Migrator', () => {
     migrator.options.disableForeignKeys = false;
     const migration2 = new MigrationTest2(orm.em.getDriver(), orm.config);
     await runner.run(migration2, 'up');
-    expect(mock.mock.calls.length).toBe(1);
+    expect(mock.mock.calls.length).toBe(2);
     expect(mock.mock.calls[0][0]).toMatch('select 1 + 1');
+    expect(mock.mock.calls[0][0]).toMatch('select 1 + 1 as count');
   });
 
   test('up/down params [all or nothing enabled]', async () => {