Merge pull request #6257 from Shopify/jb-multi-env-local-storage

Prevent race conditions during multi environment execution

Author: dejmedus

.changeset/calm-ears-add.md

@@ -0,0 +1,5 @@
+---
+'@shopify/theme': minor
+---
+
+Ensure commands run in multiple environments use the correct store URL

.changeset/many-ties-check.md

@@ -0,0 +1,5 @@
+---
+'@shopify/theme': minor
+---
+
+Allow commands run in multiple environments to act on the same store URL

packages/theme/src/cli/services/local-storage.test.ts

@@ -10,6 +10,8 @@ import {
   removeStorefrontPassword,
   ThemeLocalStorageSchema,
   setThemeStore,
+  getThemeStore,
+  useThemeStoreContext,
 } from './local-storage.js'
 import {inTemporaryDirectory} from '@shopify/cli-kit/node/fs'
 import {LocalStorage} from '@shopify/cli-kit/node/local-storage'
@@ -55,4 +57,55 @@ describe('local-storage', () => {
       })
     })
   })
+
+  describe('getThemeStore', () => {
+    test('selects store from context when inside useThemeStoreContext', async () => {
+      await inTemporaryDirectory(async (cwd) => {
+        const storage = new LocalStorage<ThemeLocalStorageSchema>({cwd})
+        setThemeStore('storage-store.myshopify.com', storage)
+
+        const initialStore = getThemeStore(storage)
+        let insideContextStore: string | undefined
+
+        await useThemeStoreContext('context-store.myshopify.com', async () => {
+          insideContextStore = getThemeStore(storage)
+        })
+
+        const outsideContextStore = getThemeStore(storage)
+
+        expect(initialStore).toBe('storage-store.myshopify.com')
+        expect(outsideContextStore).toBe('storage-store.myshopify.com')
+        expect(insideContextStore).toBe('context-store.myshopify.com')
+      })
+    })
+
+    test('ensures concurrently run commands maintain their own store value', async () => {
+      await inTemporaryDirectory(async (cwd) => {
+        const storage = new LocalStorage<ThemeLocalStorageSchema>({cwd})
+        setThemeStore('storage-store.myshopify.com', storage)
+
+        const results: {[key: string]: string | undefined} = {}
+
+        await Promise.all([
+          useThemeStoreContext('store1.myshopify.com', async () => {
+            await new Promise((resolve) => setTimeout(resolve, 10))
+            results.env1 = getThemeStore(storage)
+          }),
+          useThemeStoreContext('store2.myshopify.com', async () => {
+            await new Promise((resolve) => setTimeout(resolve, 5))
+            results.env2 = getThemeStore(storage)
+          }),
+          useThemeStoreContext('store3.myshopify.com', async () => {
+            results.env3 = getThemeStore(storage)
+          }),
+          (results.env4 = getThemeStore(storage)),
+        ])
+
+        expect(results.env1).toBe('store1.myshopify.com')
+        expect(results.env2).toBe('store2.myshopify.com')
+        expect(results.env3).toBe('store3.myshopify.com')
+        expect(results.env4).toBe('storage-store.myshopify.com')
+      })
+    })
+  })
 })

packages/theme/src/cli/services/local-storage.ts

@@ -1,6 +1,7 @@
 import {BugError} from '@shopify/cli-kit/node/error'
 import {LocalStorage} from '@shopify/cli-kit/node/local-storage'
 import {outputDebug, outputContent} from '@shopify/cli-kit/node/output'
+import {AsyncLocalStorage} from 'node:async_hooks'
 
 type DevelopmentThemeId = string
 
@@ -16,6 +17,9 @@ interface ThemeStorePasswordSchema {
   [themeStore: string]: string
 }
 
+/** Preserves the theme store a command is acting on during multi environment execution */
+const themeStoreContext = new AsyncLocalStorage<{store: string}>()
+
 let _themeLocalStorageInstance: LocalStorage<ThemeLocalStorageSchema> | undefined
 let _developmentThemeLocalStorageInstance: LocalStorage<DevelopmentThemeLocalStorageSchema> | undefined
 let _replThemeLocalStorageInstance: LocalStorage<DevelopmentThemeLocalStorageSchema> | undefined
@@ -56,7 +60,8 @@ function themeStorePasswordStorage() {
 }
 
 export function getThemeStore(storage: LocalStorage<ThemeLocalStorageSchema> = themeLocalStorage()) {
-  return storage.get('themeStore')
+  const context = themeStoreContext.getStore()
+  return context?.store ? context.store : storage.get('themeStore')
 }
 
 export function setThemeStore(store: string, storage: LocalStorage<ThemeLocalStorageSchema> = themeLocalStorage()) {
@@ -144,3 +149,8 @@ function assertThemeStoreExists(storage: LocalStorage<ThemeLocalStorageSchema> =
   }
   return themeStore
 }
+
+/** Provides theme store context to each environment during concurrent multi environment execution */
+export async function useThemeStoreContext<T>(store: string, callback: () => Promise<T>): Promise<T> {
+  return themeStoreContext.run({store}, callback)
+}

packages/theme/src/cli/utilities/theme-command.test.ts

@@ -135,8 +135,9 @@ describe('ThemeCommand', () => {
 
     test('multiple environments provided - uses renderConcurrent for parallel execution', async () => {
       // Given
-      const environmentConfig = {store: 'store.myshopify.com'}
-      vi.mocked(loadEnvironment).mockResolvedValue(environmentConfig)
+      vi.mocked(loadEnvironment)
+        .mockResolvedValueOnce({store: 'store1.myshopify.com', development: true})
+        .mockResolvedValueOnce({store: 'store2.myshopify.com', theme: 'staging'})
       vi.mocked(ensureAuthenticatedThemes).mockResolvedValue(mockSession)
 
       vi.mocked(renderConcurrent).mockResolvedValue(undefined)
@@ -150,7 +151,6 @@ describe('ThemeCommand', () => {
       // Then
       expect(loadEnvironment).toHaveBeenCalledWith('development', 'shopify.theme.toml', {from: undefined, silent: true})
       expect(loadEnvironment).toHaveBeenCalledWith('staging', 'shopify.theme.toml', {from: undefined, silent: true})
-      expect(ensureAuthenticatedThemes).toHaveBeenCalledTimes(2)
 
       expect(renderConcurrent).toHaveBeenCalledOnce()
       expect(renderConcurrent).toHaveBeenCalledWith(
@@ -166,10 +166,40 @@ describe('ThemeCommand', () => {
   })
 
   describe('multi environment', () => {
+    test('commands that act on the same store are run in groups to prevent conflicts', async () => {
+      // Given
+      vi.mocked(loadEnvironment)
+        .mockResolvedValueOnce({store: 'store1.myshopify.com', theme: 'wow a theme'})
+        .mockResolvedValueOnce({store: 'store1.myshopify.com', development: true})
+        .mockResolvedValueOnce({store: 'store2.myshopify.com', theme: 'another theme'})
+
+      vi.mocked(renderConfirmationPrompt).mockResolvedValue(true)
+      vi.mocked(renderConcurrent).mockResolvedValue(undefined)
+
+      await CommandConfig.load()
+      const command = new TestThemeCommandWithUnionFlags(
+        ['--environment', 'store1-theme', '--environment', 'store1-development', '--environment', 'store2-theme'],
+        CommandConfig,
+      )
+
+      // When
+      await command.run()
+
+      // Then
+      const runGroupOneProcesses = vi.mocked(renderConcurrent).mock.calls[0]?.[0]?.processes
+      expect(runGroupOneProcesses).toHaveLength(2)
+      expect(runGroupOneProcesses?.map((process) => process.prefix)).toEqual(['store1-theme', 'store2-theme'])
+
+      const runGroupTwoProcesses = vi.mocked(renderConcurrent).mock.calls[1]?.[0]?.processes
+      expect(runGroupTwoProcesses).toHaveLength(1)
+      expect(runGroupTwoProcesses?.map((process) => process.prefix)).toEqual(['store1-development'])
+    })
+
     test('commands with --force flag should not prompt for confirmation', async () => {
       // Given
-      const environmentConfig = {store: 'store.myshopify.com'}
-      vi.mocked(loadEnvironment).mockResolvedValue(environmentConfig)
+      vi.mocked(loadEnvironment)
+        .mockResolvedValueOnce({store: 'store1.myshopify.com', development: true})
+        .mockResolvedValueOnce({store: 'store2.myshopify.com', theme: 'staging'})
       vi.mocked(renderConfirmationPrompt).mockResolvedValue(true)
       vi.mocked(renderConcurrent).mockResolvedValue(undefined)
 
@@ -198,8 +228,9 @@ describe('ThemeCommand', () => {
 
     test('commands that do not allow --force flag should not prompt for confirmation', async () => {
       // Given
-      const environmentConfig = {store: 'store.myshopify.com'}
-      vi.mocked(loadEnvironment).mockResolvedValue(environmentConfig)
+      vi.mocked(loadEnvironment)
+        .mockResolvedValueOnce({store: 'store1.myshopify.com', development: true})
+        .mockResolvedValueOnce({store: 'store2.myshopify.com', theme: 'staging'})
       vi.mocked(renderConfirmationPrompt).mockResolvedValue(true)
       vi.mocked(renderConcurrent).mockResolvedValue(undefined)
 
@@ -216,8 +247,9 @@ describe('ThemeCommand', () => {
 
     test('commands without --force flag that allow it should prompt for confirmation', async () => {
       // Given
-      const environmentConfig = {store: 'store.myshopify.com'}
-      vi.mocked(loadEnvironment).mockResolvedValue(environmentConfig)
+      vi.mocked(loadEnvironment)
+        .mockResolvedValueOnce({store: 'store1.myshopify.com', development: true})
+        .mockResolvedValueOnce({store: 'store2.myshopify.com', theme: 'staging'})
       vi.mocked(renderConfirmationPrompt).mockResolvedValue(true)
       vi.mocked(renderConcurrent).mockResolvedValue(undefined)
 
@@ -244,8 +276,9 @@ describe('ThemeCommand', () => {
 
     test('should not execute command if confirmation is cancelled', async () => {
       // Given
-      const environmentConfig = {store: 'store.myshopify.com'}
-      vi.mocked(loadEnvironment).mockResolvedValue(environmentConfig)
+      vi.mocked(loadEnvironment)
+        .mockResolvedValueOnce({store: 'store1.myshopify.com', development: true})
+        .mockResolvedValueOnce({store: 'store2.myshopify.com', theme: 'staging'})
       vi.mocked(renderConfirmationPrompt).mockResolvedValue(false)
       vi.mocked(renderConcurrent).mockResolvedValue(undefined)
 
@@ -340,8 +373,10 @@ describe('ThemeCommand', () => {
 
     test('commands error gracefully and continue with other environments', async () => {
       // Given
-      vi.mocked(loadEnvironment).mockResolvedValue({store: 'store.myshopify.com'})
-
+      vi.mocked(loadEnvironment)
+        .mockResolvedValueOnce({store: 'store1.myshopify.com', development: true})
+        .mockResolvedValueOnce({store: 'store2.myshopify.com', theme: 'staging'})
+        .mockResolvedValueOnce({store: 'store3.myshopify.com', live: true})
       vi.mocked(renderConfirmationPrompt).mockResolvedValue(true)
       vi.mocked(renderConcurrent).mockImplementation(async ({processes}) => {
         for (const process of processes) {

packages/theme/src/cli/utilities/theme-command.ts

@@ -1,6 +1,7 @@
 import {ensureThemeStore} from './theme-store.js'
 import {configurationFileName} from '../constants.js'
 import metadata from '../metadata.js'
+import {useThemeStoreContext} from '../services/local-storage.js'
 import {hashString} from '@shopify/cli-kit/node/crypto'
 import {Input} from '@oclif/core/interfaces'
 import Command, {ArgOutput, FlagOutput} from '@shopify/cli-kit/node/base-command'
@@ -153,9 +154,8 @@ export default abstract class ThemeCommand extends Command {
     environmentMap: Map<EnvironmentName, FlagValues>,
     requiredFlags: Exclude<RequiredFlags, null>,
   ) {
-    const valid: {environment: EnvironmentName; flags: FlagValues; session: AdminSession}[] = []
+    const valid: {environment: EnvironmentName; flags: FlagValues}[] = []
     const invalid: {environment: EnvironmentName; reason: string}[] = []
-    const commandName = this.constructor.name.toLowerCase()
 
     for (const [environmentName, environmentFlags] of environmentMap) {
       const validationResult = this.validConfig(environmentFlags, requiredFlags, environmentName)
@@ -165,12 +165,7 @@ export default abstract class ThemeCommand extends Command {
         continue
       }
 
-      // eslint-disable-next-line no-await-in-loop
-      const session = await this.createSession(environmentFlags)
-
-      recordEvent(`theme-command:${commandName}:multi-env:authenticated`)
-
-      valid.push({environment: environmentName, flags: environmentFlags, session})
+      valid.push({environment: environmentName, flags: environmentFlags})
     }
 
     return {valid, invalid}
@@ -227,30 +222,61 @@ export default abstract class ThemeCommand extends Command {
    * Run the command in each valid environment concurrently
    * @param validEnvironments - The valid environments to run the command in
    */
-  private async runConcurrent(
-    validEnvironments: {environment: EnvironmentName; flags: FlagValues; session: AdminSession}[],
-  ) {
+  private async runConcurrent(validEnvironments: {environment: EnvironmentName; flags: FlagValues}[]) {
     const abortController = new AbortController()
 
-    await renderConcurrent({
-      processes: validEnvironments.map(({environment, flags, session}) => ({
-        prefix: environment,
-        action: async (stdout: Writable, stderr: Writable, _signal) => {
-          try {
-            await this.command(flags, session, true, {stdout, stderr})
-
-            // eslint-disable-next-line no-catch-all/no-catch-all
-          } catch (error) {
-            if (error instanceof Error) {
-              error.message = `Environment ${environment} failed: \n\n${error.message}`
-              renderError({body: [error.message]})
+    const stores = validEnvironments.map((env) => env.flags.store as string)
+    const uniqueStores = new Set(stores)
+    const runGroups =
+      stores.length === uniqueStores.size ? [validEnvironments] : this.createSequentialGroups(validEnvironments)
+
+    for (const runGroup of runGroups) {
+      // eslint-disable-next-line no-await-in-loop
+      await renderConcurrent({
+        processes: runGroup.map(({environment, flags}) => ({
+          prefix: environment,
+          action: async (stdout: Writable, stderr: Writable, _signal) => {
+            try {
+              const store = flags.store as string
+              await useThemeStoreContext(store, async () => {
+                const session = await this.createSession(flags)
+
+                const commandName = this.constructor.name.toLowerCase()
+                recordEvent(`theme-command:${commandName}:multi-env:authenticated`)
+
+                await this.command(flags, session, true, {stdout, stderr})
+              })
+
+              // eslint-disable-next-line no-catch-all/no-catch-all
+            } catch (error) {
+              if (error instanceof Error) {
+                error.message = `Environment ${environment} failed: \n\n${error.message}`
+                renderError({body: [error.message]})
+              }
             }
-          }
-        },
-      })),
-      abortSignal: abortController.signal,
-      showTimestamps: true,
+          },
+        })),
+        abortSignal: abortController.signal,
+        showTimestamps: true,
+      })
+    }
+  }
+
+  /**
+   * Create groups of environments with unique flags.store values to run sequentially
+   * to prevent conflicts between environments acting on the same store
+   * @param environments - The environments to group
+   * @returns The environment groups
+   */
+  private createSequentialGroups(environments: {environment: EnvironmentName; flags: FlagValues}[]) {
+    const groups: {environment: EnvironmentName; flags: FlagValues}[][] = []
+
+    environments.forEach((environment) => {
+      const groupWithoutStore = groups.find((arr) => !arr.some((env) => env.flags.store === environment.flags.store))
+      groupWithoutStore ? groupWithoutStore.push(environment) : groups.push([environment])
     })
+
+    return groups
   }
 
   /**