From 7f52a1e28ce622b3c2259213f26397641dbe7ed3 Mon Sep 17 00:00:00 2001
From: Adria Navarro <adria@budibase.com>
Date: Fri, 1 Dec 2023 10:39:44 +0100
Subject: [PATCH] Guard migration

---
 .../backend-core/src/context/mainContext.ts   | 17 ++++--
 .../src/context/tests/index.spec.ts           | 54 +++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/packages/backend-core/src/context/mainContext.ts b/packages/backend-core/src/context/mainContext.ts
index 1a4d6db454..d2fdd040df 100644
--- a/packages/backend-core/src/context/mainContext.ts
+++ b/packages/backend-core/src/context/mainContext.ts
@@ -99,6 +99,8 @@ function updateContext(updates: ContextMap): ContextMap {
 }
 
 async function newContext<T>(updates: ContextMap, task: () => T) {
+  guardMigration()
+
   // see if there already is a context setup
   let context: ContextMap = updateContext(updates)
   return Context.run(context, task)
@@ -186,13 +188,22 @@ export async function doInIdentityContext<T>(
   return newContext(context, task)
 }
 
+function guardMigration() {
+  const context = Context.get()
+  if (context?.isMigrating) {
+    throw new Error(
+      "The context cannot be change, a migration is currently running"
+    )
+  }
+}
+
 export async function doInAppMigrationContext<T>(
   appId: string,
   task: () => T
 ): Promise<T> {
-    return _doInAppContext(appId, task, {
-      isMigrating: true,
-    })
+  return _doInAppContext(appId, task, {
+    isMigrating: true,
+  })
 }
 
 export function getIdentity(): IdentityContext | undefined {
diff --git a/packages/backend-core/src/context/tests/index.spec.ts b/packages/backend-core/src/context/tests/index.spec.ts
index 4bfbea74ce..183e09a509 100644
--- a/packages/backend-core/src/context/tests/index.spec.ts
+++ b/packages/backend-core/src/context/tests/index.spec.ts
@@ -5,6 +5,7 @@ import { structures } from "../../../tests"
 import { db } from "../.."
 import Context from "../Context"
 import { ContextMap } from "../types"
+import { IdentityType } from "@budibase/types"
 
 describe("context", () => {
   describe("doInTenant", () => {
@@ -197,5 +198,58 @@ describe("context", () => {
 
       expect(Context.get()).toBeUndefined()
     })
+
+    it.each([
+      [
+        "doInAppMigrationContext",
+        () => context.doInAppMigrationContext(db.generateAppID(), () => {}),
+      ],
+      [
+        "doInAppContext",
+        () => context.doInAppContext(db.generateAppID(), () => {}),
+      ],
+      [
+        "doInAutomationContext",
+        () =>
+          context.doInAutomationContext({
+            appId: db.generateAppID(),
+            automationId: structures.generator.guid(),
+            task: () => {},
+          }),
+      ],
+      ["doInContext", () => context.doInContext(db.generateAppID(), () => {})],
+      [
+        "doInEnvironmentContext",
+        () => context.doInEnvironmentContext({}, () => {}),
+      ],
+      [
+        "doInIdentityContext",
+        () =>
+          context.doInIdentityContext(
+            {
+              account: undefined,
+              type: IdentityType.USER,
+              _id: structures.users.user()._id!,
+            },
+            () => {}
+          ),
+      ],
+      ["doInScimContext", () => context.doInScimContext(() => {})],
+      [
+        "doInTenant",
+        () => context.doInTenant(structures.tenant.id(), () => {}),
+      ],
+    ])(
+      "a nested context.%s function cannot run",
+      async (_, otherContextCall: () => Promise<void>) => {
+        await expect(
+          context.doInAppMigrationContext(db.generateAppID(), async () => {
+            await otherContextCall()
+          })
+        ).rejects.toThrowError(
+          "The context cannot be change, a migration is currently running"
+        )
+      }
+    )
   })
 })