diff --git a/packages/backend-core/src/context/mainContext.ts b/packages/backend-core/src/context/mainContext.ts index 0dba2456a9..1cbf1de1c3 100644 --- a/packages/backend-core/src/context/mainContext.ts +++ b/packages/backend-core/src/context/mainContext.ts @@ -147,13 +147,21 @@ export async function doInTenant( export async function doInAppContext( appId: string, task: () => T +): Promise { + return _doInAppContext(appId, task) +} + +async function _doInAppContext( + appId: string, + task: () => T, + extraContextSettings?: ContextMap ): Promise { if (!appId) { throw new Error("appId is required") } const tenantId = getTenantIDFromAppID(appId) - const updates: ContextMap = { appId } + const updates: ContextMap = { appId, ...extraContextSettings } if (tenantId) { updates.tenantId = tenantId } @@ -182,21 +190,15 @@ export async function doInAppMigrationContext( appId: string, task: () => T ): Promise { - if (!appId && !env.isTest()) { - throw new Error("appId is required") + try { + return _doInAppContext(appId, task, { + isMigrating: true, + }) + } finally { + updateContext({ + isMigrating: undefined, + }) } - - let updates: ContextMap - if (!appId) { - updates = { appId: "" } - } else { - const tenantId = getTenantIDFromAppID(appId) - updates = { appId } - if (tenantId) { - updates.tenantId = tenantId - } - } - return newContext(updates, task) } 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 d8bb598af1..ef7cf17e0c 100644 --- a/packages/backend-core/src/context/tests/index.spec.ts +++ b/packages/backend-core/src/context/tests/index.spec.ts @@ -1,6 +1,10 @@ import { testEnv } from "../../../tests/extra" import * as context from "../" import { DEFAULT_TENANT_ID } from "../../constants" +import { structures } from "../../../tests" +import { db } from "../.." +import Context from "../Context" +import { ContextMap } from "../types" describe("context", () => { describe("doInTenant", () => { @@ -144,4 +148,36 @@ describe("context", () => { expect(isScim).toBe(false) }) }) + + describe("doInAppMigrationContext", () => { + it("the context is set correctly", async () => { + const appId = db.generateAppID() + + await context.doInAppMigrationContext(appId, () => { + const context = Context.get() + + const expected: ContextMap = { + appId, + isMigrating: true, + } + expect(context).toEqual(expected) + }) + }) + + it("the context is set correctly when running in a tenant id", async () => { + const tenantId = structures.tenant.id() + const appId = db.generateAppID(tenantId) + + await context.doInAppMigrationContext(appId, () => { + const context = Context.get() + + const expected: ContextMap = { + appId, + isMigrating: true, + tenantId, + } + expect(context).toEqual(expected) + }) + }) + }) }) diff --git a/packages/backend-core/src/context/types.ts b/packages/backend-core/src/context/types.ts index d687a93594..a1606a17b9 100644 --- a/packages/backend-core/src/context/types.ts +++ b/packages/backend-core/src/context/types.ts @@ -8,4 +8,5 @@ export type ContextMap = { environmentVariables?: Record isScim?: boolean automationId?: string + isMigrating?: boolean }