From 5c1f60aeeb128c2dfedfa08cc9b98d05c1582434 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Ad=C3=A1mek?= <banan23@gmail.com>
Date: Mon, 24 Aug 2020 22:46:41 +0200
Subject: [PATCH] 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
---
 packages/knex/src/AbstractSqlConnection.ts |  2 +-
 packages/migrations/src/Migration.ts       | 14 ++++++--
 packages/migrations/src/MigrationRunner.ts | 38 +++++++++++++---------
 tests/Migrator.test.ts                     |  5 ++-
 4 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/packages/knex/src/AbstractSqlConnection.ts b/packages/knex/src/AbstractSqlConnection.ts
index 1d7ec6a4c8e8..15edeca14cba 100644
--- 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
index 80b3a9d30b4c..cc503576ee97 100644
--- 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
index b38ee1cec7a8..bba0d0dd06dc 100644
--- 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
index 54c640a14599..ab5bdba190f4 100644
--- 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 () => {