Allow Partners API for extension migration queries

gonzaloriestra · gonzaloriestra · commit 9b9a51f01b32 · 2025-12-10T14:10:52.000+01:00
diff --git a/packages/app/src/cli/utilities/developer-platform-client/partners-client.test.ts b/packages/app/src/cli/utilities/developer-platform-client/partners-client.test.ts
@@ -1,5 +1,8 @@
 import {PartnersClient} from './partners-client.js'
 import {CreateAppQuery} from '../../api/graphql/create_app.js'
+import {MigrateToUiExtensionQuery} from '../../api/graphql/extension_migrate_to_ui_extension.js'
+import {MigrateFlowExtensionMutation} from '../../api/graphql/extension_migrate_flow_extension.js'
+import {MigrateAppModuleMutation} from '../../api/graphql/extension_migrate_app_module.js'
 import {AppInterface, WebType} from '../../models/app/app.js'
 import {Organization, OrganizationSource, OrganizationStore} from '../../models/organization.js'
 import {
@@ -110,7 +113,7 @@ describe('createApp', () => {
     expect(partnersRequest).toHaveBeenCalledWith(CreateAppQuery, 'token', variables, undefined, undefined, {
       type: 'token_refresh',
       handler: expect.any(Function),
-    })
+    }, false)
   })
 
   test('creates an app with non-launchable defaults', async () => {
@@ -142,7 +145,7 @@ describe('createApp', () => {
     expect(partnersRequest).toHaveBeenCalledWith(CreateAppQuery, 'token', variables, undefined, undefined, {
       type: 'token_refresh',
       handler: expect.any(Function),
-    })
+    }, false)
   })
 
   test('throws error if requests has a user error', async () => {
@@ -176,7 +179,7 @@ describe('fetchApp', async () => {
     expect(partnersRequest).toHaveBeenCalledWith(FindOrganizationQuery, 'token', {id: ORG1.id}, undefined, undefined, {
       type: 'token_refresh',
       handler: expect.any(Function),
-    })
+    }, false)
   })
 
   test('throws if there are no organizations', async () => {
@@ -192,7 +195,7 @@ describe('fetchApp', async () => {
     expect(partnersRequest).toHaveBeenCalledWith(FindOrganizationQuery, 'token', {id: ORG1.id}, undefined, undefined, {
       type: 'token_refresh',
       handler: expect.any(Function),
-    })
+    }, false)
   })
 })
 
@@ -230,3 +233,101 @@ describe('singleton pattern', () => {
     expect(instance1).not.toBe(instance2)
   })
 })
+
+describe('request with forceUsePartnersApi', () => {
+  test('passes forceUsePartnersApi=true for MigrateToUiExtensionQuery', async () => {
+    // Given
+    const partnersClient = PartnersClient.getInstance(testPartnersUserSession)
+    const mockResponse = {migrateToUiExtension: {userErrors: [], migratedToUiExtension: true}}
+    vi.mocked(partnersRequest).mockResolvedValueOnce(mockResponse)
+
+    // When
+    await partnersClient.request(MigrateToUiExtensionQuery, {extensionId: '123'})
+
+    // Then
+    expect(partnersRequest).toHaveBeenCalledWith(
+      MigrateToUiExtensionQuery,
+      testPartnersUserSession.token,
+      {extensionId: '123'},
+      undefined,
+      undefined,
+      expect.objectContaining({
+        type: 'token_refresh',
+        handler: expect.any(Function),
+      }),
+      true, // forceUsePartnersApi should be true
+    )
+  })
+
+  test('passes forceUsePartnersApi=true for MigrateFlowExtensionMutation', async () => {
+    // Given
+    const partnersClient = PartnersClient.getInstance(testPartnersUserSession)
+    const mockResponse = {migrateFlowExtension: {userErrors: [], migratedFlowExtension: true}}
+    vi.mocked(partnersRequest).mockResolvedValueOnce(mockResponse)
+
+    // When
+    await partnersClient.request(MigrateFlowExtensionMutation, {extensionId: '123'})
+
+    // Then
+    expect(partnersRequest).toHaveBeenCalledWith(
+      MigrateFlowExtensionMutation,
+      testPartnersUserSession.token,
+      {extensionId: '123'},
+      undefined,
+      undefined,
+      expect.objectContaining({
+        type: 'token_refresh',
+        handler: expect.any(Function),
+      }),
+      true, // forceUsePartnersApi should be true
+    )
+  })
+
+  test('passes forceUsePartnersApi=true for MigrateAppModuleMutation', async () => {
+    // Given
+    const partnersClient = PartnersClient.getInstance(testPartnersUserSession)
+    const mockResponse = {migrateAppModule: {userErrors: [], migratedAppModule: true}}
+    vi.mocked(partnersRequest).mockResolvedValueOnce(mockResponse)
+
+    // When
+    await partnersClient.request(MigrateAppModuleMutation, {appModuleId: '123'})
+
+    // Then
+    expect(partnersRequest).toHaveBeenCalledWith(
+      MigrateAppModuleMutation,
+      testPartnersUserSession.token,
+      {appModuleId: '123'},
+      undefined,
+      undefined,
+      expect.objectContaining({
+        type: 'token_refresh',
+        handler: expect.any(Function),
+      }),
+      true, // forceUsePartnersApi should be true
+    )
+  })
+
+  test('passes forceUsePartnersApi=false for other queries', async () => {
+    // Given
+    const partnersClient = PartnersClient.getInstance(testPartnersUserSession)
+    const mockResponse = {organizations: {nodes: []}}
+    vi.mocked(partnersRequest).mockResolvedValueOnce(mockResponse)
+
+    // When
+    await partnersClient.request(FindOrganizationQuery, {id: '1'})
+
+    // Then
+    expect(partnersRequest).toHaveBeenCalledWith(
+      FindOrganizationQuery,
+      testPartnersUserSession.token,
+      {id: '1'},
+      undefined,
+      undefined,
+      expect.objectContaining({
+        type: 'token_refresh',
+        handler: expect.any(Function),
+      }),
+      false, // forceUsePartnersApi should be false for non-migration queries
+    )
+  })
+})
diff --git a/packages/app/src/cli/utilities/developer-platform-client/partners-client.ts b/packages/app/src/cli/utilities/developer-platform-client/partners-client.ts
@@ -262,13 +262,16 @@ export class PartnersClient implements DeveloperPlatformClient {
     cacheOptions?: CacheOptions,
     preferredBehaviour?: RequestModeInput,
   ): Promise<T> {
+    const allowedQueries = [MigrateToUiExtensionQuery, MigrateFlowExtensionMutation, MigrateAppModuleMutation]
+    const forceUsePartnersApi = allowedQueries.some((allowedQuery) => query === allowedQuery)
     return partnersRequest(
       query,
       await this.token(),
       variables,
       cacheOptions,
       preferredBehaviour,
       this.createUnauthorizedHandler(),
+      forceUsePartnersApi,
     )
   }
 
diff --git a/packages/cli-kit/src/public/node/api/partners.test.ts b/packages/cli-kit/src/public/node/api/partners.test.ts
@@ -1,12 +1,15 @@
 import {partnersRequest, handleDeprecations} from './partners.js'
 import {graphqlRequest, GraphQLResponse} from './graphql.js'
 import {partnersFqdn} from '../context/fqdn.js'
+import {blockPartnersAccess} from '../environment.js'
+import {BugError} from '../error.js'
 import {setNextDeprecationDate} from '../../../private/node/context/deprecations-store.js'
 import {test, vi, expect, describe, beforeEach, beforeAll} from 'vitest'
 
 vi.mock('./graphql.js')
 vi.mock('../../../private/node/context/deprecations-store.js')
 vi.mock('../context/fqdn.js')
+vi.mock('../environment.js')
 
 const mockedResult = 'OK'
 const partnersFQDN = 'partners.shopify.com'
@@ -16,6 +19,7 @@ const mockedToken = 'token'
 
 beforeEach(() => {
   vi.mocked(partnersFqdn).mockResolvedValue(partnersFQDN)
+  vi.mocked(blockPartnersAccess).mockReturnValue(false)
 })
 
 describe('partnersRequest', () => {
@@ -36,6 +40,50 @@ describe('partnersRequest', () => {
       responseOptions: {onResponse: handleDeprecations},
     })
   })
+
+  test('throws BugError when blockPartnersAccess returns true without forceUsePartnersApi', async () => {
+    // Given
+    vi.mocked(blockPartnersAccess).mockReturnValue(true)
+
+    // When/Then
+    await expect(partnersRequest('query', mockedToken, {variables: 'variables'})).rejects.toThrow(BugError)
+    expect(blockPartnersAccess).toHaveBeenCalledWith(false)
+  })
+
+  test('throws BugError when blockPartnersAccess returns true even with forceUsePartnersApi=false', async () => {
+    // Given
+    vi.mocked(blockPartnersAccess).mockReturnValue(true)
+
+    // When/Then
+    await expect(
+      partnersRequest('query', mockedToken, {variables: 'variables'}, undefined, undefined, undefined, false),
+    ).rejects.toThrow(BugError)
+    expect(blockPartnersAccess).toHaveBeenCalledWith(false)
+  })
+
+  test('does not throw when forceUsePartnersApi=true even if blockPartnersAccess would normally block', async () => {
+    // Given
+    vi.mocked(blockPartnersAccess).mockReturnValue(false)
+    vi.mocked(graphqlRequest).mockResolvedValue(mockedResult)
+
+    // When
+    await partnersRequest('query', mockedToken, {variables: 'variables'}, undefined, undefined, undefined, true)
+
+    // Then
+    expect(blockPartnersAccess).toHaveBeenCalledWith(true)
+    expect(graphqlRequest).toHaveBeenCalled()
+  })
+
+  test('passes forceUsePartnersApi to blockPartnersAccess', async () => {
+    // Given
+    vi.mocked(graphqlRequest).mockResolvedValue(mockedResult)
+
+    // When
+    await partnersRequest('query', mockedToken, {variables: 'variables'}, undefined, undefined, undefined, true)
+
+    // Then
+    expect(blockPartnersAccess).toHaveBeenCalledWith(true)
+  })
 })
 
 describe('handleDeprecations', () => {
diff --git a/packages/cli-kit/src/public/node/api/partners.ts b/packages/cli-kit/src/public/node/api/partners.ts
@@ -11,9 +11,10 @@ import {partnersFqdn} from '../context/fqdn.js'
 import {setNextDeprecationDate} from '../../../private/node/context/deprecations-store.js'
 import {getPackageManager} from '../node-package-manager.js'
 import {cwd} from '../path.js'
-import {AbortError} from '../error.js'
+import {AbortError, BugError} from '../error.js'
 import {formatPackageManagerCommand} from '../output.js'
 import {RequestModeInput} from '../http.js'
+import {blockPartnersAccess} from '../environment.js'
 import Bottleneck from 'bottleneck'
 import {Variables} from 'graphql-request'
 import {TypedDocumentNode} from '@graphql-typed-document-node/core'
@@ -30,8 +31,13 @@ const limiter = new Bottleneck({
  * Sets up the request to the Partners API.
  *
  * @param token - Partners token.
+ * @param forceUsePartnersApi - If true, the request will be forced to use the Partners API.
  */
-async function setupRequest(token: string) {
+async function setupRequest(token: string, forceUsePartnersApi = false) {
+  if (blockPartnersAccess(forceUsePartnersApi)) {
+    throw new BugError('Partners API is no longer available.')
+  }
+
   const api = 'Partners'
   const fqdn = await partnersFqdn()
   const url = `https://${fqdn}/api/cli/graphql`
@@ -52,6 +58,7 @@ async function setupRequest(token: string) {
  * @param cacheOptions - Cache options.
  * @param preferredBehaviour - Preferred behaviour for the request.
  * @param unauthorizedHandler - Optional handler for unauthorized requests.
+ * @param forceUsePartnersApi - If true, the request will be forced to use the Partners API.
  * @returns The response of the query of generic type <T>.
  */
 export async function partnersRequest<T>(
@@ -61,8 +68,9 @@ export async function partnersRequest<T>(
   cacheOptions?: CacheOptions,
   preferredBehaviour?: RequestModeInput,
   unauthorizedHandler?: UnauthorizedHandler,
+  forceUsePartnersApi = false,
 ): Promise<T> {
-  const opts = await setupRequest(token)
+  const opts = await setupRequest(token, forceUsePartnersApi)
   const result = limiter.schedule(() =>
     graphqlRequest<T>({
       ...opts,
diff --git a/packages/cli-kit/src/public/node/context/fqdn.test.ts b/packages/cli-kit/src/public/node/context/fqdn.test.ts
@@ -8,11 +8,9 @@ import {
   adminFqdn,
 } from './fqdn.js'
 import {Environment, serviceEnvironment} from '../../../private/node/context/service.js'
-import {blockPartnersAccess} from '../environment.js'
 import {expect, describe, test, vi} from 'vitest'
 
 vi.mock('../../../private/node/context/service.js')
-vi.mock('../environment.js')
 
 vi.mock('../vendor/dev_server/index.js', () => {
   return {
@@ -34,7 +32,6 @@ describe('partners', () => {
   test('returns the local fqdn when the environment is local', async () => {
     // Given
     vi.mocked(serviceEnvironment).mockReturnValue(Environment.Local)
-    vi.mocked(blockPartnersAccess).mockReturnValue(false)
 
     // When
     const got = await partnersFqdn()
@@ -46,7 +43,6 @@ describe('partners', () => {
   test('returns the production fqdn when the environment is production', async () => {
     // Given
     vi.mocked(serviceEnvironment).mockReturnValue(Environment.Production)
-    vi.mocked(blockPartnersAccess).mockReturnValue(false)
 
     // When
     const got = await partnersFqdn()
diff --git a/packages/cli-kit/src/public/node/context/fqdn.ts b/packages/cli-kit/src/public/node/context/fqdn.ts
@@ -1,7 +1,6 @@
-import {AbortError, BugError} from '../error.js'
+import {AbortError} from '../error.js'
 import {serviceEnvironment} from '../../../private/node/context/service.js'
 import {DevServer, DevServerCore} from '../vendor/dev_server/index.js'
-import {blockPartnersAccess} from '../environment.js'
 
 export const NotProvidedStoreFQDNError = new AbortError(
   "Couldn't obtain the Shopify FQDN because the store FQDN was not provided.",
@@ -13,9 +12,6 @@ export const NotProvidedStoreFQDNError = new AbortError(
  * @returns Fully-qualified domain of the partners service we should interact with.
  */
 export async function partnersFqdn(): Promise<string> {
-  if (blockPartnersAccess()) {
-    throw new BugError('Partners API is is no longer available.')
-  }
   const environment = serviceEnvironment()
   const productionFqdn = 'partners.shopify.com'
   switch (environment) {
diff --git a/packages/cli-kit/src/public/node/environment.test.ts b/packages/cli-kit/src/public/node/environment.test.ts
diff --git a/packages/cli-kit/src/public/node/environment.ts b/packages/cli-kit/src/public/node/environment.ts
