From ec8cdf1b0daa88d97e9d03aacf6c099427a17d11 Mon Sep 17 00:00:00 2001
From: Ariel Caplan <ariel.caplan@shopify.com>
Date: Mon, 5 Aug 2024 23:31:39 +0300
Subject: [PATCH] Check organization beta flags for gated template
 specifications

---
 .../app-management-client.ts                  | 28 ++++++++---
 .../src/public/node/api/business-platform.ts  | 49 ++++++++++++++++---
 2 files changed, 62 insertions(+), 15 deletions(-)

diff --git a/packages/app/src/cli/utilities/developer-platform-client/app-management-client.ts b/packages/app/src/cli/utilities/developer-platform-client/app-management-client.ts
index 3053d3c86a..ccf9582fa0 100644
--- a/packages/app/src/cli/utilities/developer-platform-client/app-management-client.ts
+++ b/packages/app/src/cli/utilities/developer-platform-client/app-management-client.ts
@@ -38,6 +38,11 @@ import {
   AppVersionByIdQueryVariables,
   AppModule as AppModuleReturnType,
 } from './app-management-client/graphql/app-version-by-id.js'
+import {
+  OrganizationBetaFlagsQuerySchema,
+  OrganizationBetaFlagsQueryVariables,
+  organizationBetaFlagsQuery,
+} from './app-management-client/graphql/organization_beta_flags.js'
 import {RemoteSpecification} from '../../api/graphql/extension_specifications.js'
 import {
   DeveloperPlatformClient,
@@ -127,6 +132,7 @@ import {appManagementRequest} from '@shopify/cli-kit/node/api/app-management'
 import {appDevRequest} from '@shopify/cli-kit/node/api/app-dev'
 import {
   businessPlatformOrganizationsRequest,
+  businessPlatformOrganizationsRequestDoc,
   businessPlatformRequest,
   businessPlatformRequestDoc,
 } from '@shopify/cli-kit/node/api/business-platform'
@@ -381,7 +387,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
   // partners-client and app-management-client. Since we need transferDisabled and convertableToPartnerTest values
   // from the Partners OrganizationStore schema, we will return this type for now
   async devStoresForOrg(orgId: string): Promise<OrganizationStore[]> {
-    const storesResult = await businessPlatformOrganizationsRequest<ListAppDevStoresQuery>(
+    const storesResult = await businessPlatformOrganizationsRequestDoc<ListAppDevStoresQuery>(
       ListAppDevStores,
       await this.businessPlatformToken(),
       orgId,
@@ -731,7 +737,7 @@ export class AppManagementClient implements DeveloperPlatformClient {
   // from the Partners FindByStoreDomainSchema, we will return this type for now
   async storeByDomain(orgId: string, shopDomain: string): Promise<FindStoreByDomainSchema> {
     const queryVariables: FetchDevStoreByDomainQueryVariables = {domain: shopDomain}
-    const storesResult = await businessPlatformOrganizationsRequest(
+    const storesResult = await businessPlatformOrganizationsRequestDoc(
       FetchDevStoreByDomain,
       await this.businessPlatformToken(),
       orgId,
@@ -865,15 +871,21 @@ export class AppManagementClient implements DeveloperPlatformClient {
   }
 
   private async organizationBetaFlags(
-    _organizationId: string,
+    organizationId: string,
     allBetaFlags: string[],
-  ): Promise<{[key: string]: boolean}> {
-    // For now, stub everything as false
-    const stub: {[flag: string]: boolean} = {}
+  ): Promise<{[flag: (typeof allBetaFlags)[number]]: boolean}> {
+    const variables: OrganizationBetaFlagsQueryVariables = {organizationId: encodedGidFromId(organizationId)}
+    const flagsResult = await businessPlatformOrganizationsRequest<OrganizationBetaFlagsQuerySchema>(
+      organizationBetaFlagsQuery(allBetaFlags),
+      await this.businessPlatformToken(),
+      organizationId,
+      variables,
+    )
+    const result: {[flag: (typeof allBetaFlags)[number]]: boolean} = {}
     allBetaFlags.forEach((flag) => {
-      stub[flag] = false
+      result[flag] = Boolean(flagsResult.organization[`flag_${flag}`])
     })
-    return stub
+    return result
   }
 }
 
diff --git a/packages/cli-kit/src/public/node/api/business-platform.ts b/packages/cli-kit/src/public/node/api/business-platform.ts
index 62b3353e82..dca22d7a96 100644
--- a/packages/cli-kit/src/public/node/api/business-platform.ts
+++ b/packages/cli-kit/src/public/node/api/business-platform.ts
@@ -61,6 +61,23 @@ export async function businessPlatformRequestDoc<TResult, TVariables extends Var
   })
 }
 
+/**
+ * Sets up the request to the Business Platform Organizations API.
+ *
+ * @param token - Business Platform token.
+ */
+async function setupOrganizationsRequest(token: string, organizationId: string) {
+  const api = 'BusinessPlatform'
+  const fqdn = await businessPlatformFqdn()
+  const url = `https://${fqdn}/organizations/api/unstable/organization/${organizationId}/graphql`
+  return {
+    token,
+    api,
+    url,
+    responseOptions: {onResponse: handleDeprecations},
+  }
+}
+
 /**
  * Executes a GraphQL query against the Business Platform Organizations API.
  *
@@ -70,20 +87,38 @@ export async function businessPlatformRequestDoc<TResult, TVariables extends Var
  * @param variables - GraphQL variables to pass to the query.
  * @returns The response of the query of generic type <T>.
  */
-export async function businessPlatformOrganizationsRequest<TResult>(
+export async function businessPlatformOrganizationsRequest<T>(
+  query: string,
+  token: string,
+  organizationId: string,
+  variables?: GraphQLVariables,
+): Promise<T> {
+  return graphqlRequest<T>({
+    query,
+    ...(await setupOrganizationsRequest(token, organizationId)),
+    variables,
+    responseOptions: {onResponse: handleDeprecations},
+  })
+}
+
+/**
+ * Executes a GraphQL query against the Business Platform Organizations API. Uses typed documents.
+ *
+ * @param query - GraphQL query to execute.
+ * @param token - Business Platform token.
+ * @param organizationId - Organization ID as a numeric value.
+ * @param variables - GraphQL variables to pass to the query.
+ * @returns The response of the query of generic type <T>.
+ */
+export async function businessPlatformOrganizationsRequestDoc<TResult>(
   query: TypedDocumentNode<TResult, GraphQLVariables> | TypedDocumentNode<TResult, Exact<{[key: string]: never}>>,
   token: string,
   organizationId: string,
   variables?: GraphQLVariables,
 ): Promise<TResult> {
-  const api = 'BusinessPlatform'
-  const fqdn = await businessPlatformFqdn()
-  const url = `https://${fqdn}/organizations/api/unstable/organization/${organizationId}/graphql`
   return graphqlRequestDoc({
     query,
-    api,
-    url,
-    token,
+    ...(await setupOrganizationsRequest(token, organizationId)),
     variables,
     responseOptions: {onResponse: handleDeprecations},
   })