@aws-amplify/amplify-util-uibuilder 1.4.1-beta.0 → 1.4.2-rc.cf54accdc.0

package/src/__tests__/getAmplifyDataSchema.test.ts

@@ -0,0 +1,102 @@
+import aws from 'aws-sdk'; // eslint-disable-line import/no-extraneous-dependencies
+import { AmplifyCategories, AmplifySupportedService } from 'amplify-cli-core'; // eslint-disable-line import/no-extraneous-dependencies
+import { printer } from 'amplify-prompts'; // eslint-disable-line import/no-extraneous-dependencies
+import { AmplifyStudioClient } from '../clients';
+import { getAmplifyDataSchema } from '../commands/utils';
+
+const awsMock = aws as any;
+const printerMock = printer as any;
+
+jest.mock('amplify-cli-core', () => ({
+  ...jest.requireActual('amplify-cli-core'),
+  stateManager: {
+    getMeta: jest.fn(() => ({
+      [AmplifyCategories.API]: {
+        MyResourceName: {
+          service: AmplifySupportedService.APPSYNC,
+        },
+      },
+    })),
+  },
+  FeatureFlags: {
+    getBoolean: () => false,
+    getNumber: () => 0,
+  },
+}));
+jest.mock('amplify-prompts');
+
+const getMockedSchema = jest.fn(
+  // eslint-disable-next-line no-useless-escape, spellcheck/spell-checker
+  () => `export const schema = {\n    \"models\": {\n        \"Blog\": {\n            \"name\": \"Blog\",\n            \"fields\": {\n                \"id\": {\n                    \"name\": \"id\",\n                    \"isArray\": false,\n                    \"type\": \"ID\",\n                    \"isRequired\": true,\n                    \"attributes\": []\n                },\n                \"title\": {\n                    \"name\": \"title\",\n                    \"isArray\": false,\n                    \"type\": \"String\",\n                    \"isRequired\": false,\n                    \"attributes\": []\n                },\n                \"content\": {\n                    \"name\": \"content\",\n                    \"isArray\": false,\n                    \"type\": \"String\",\n                    \"isRequired\": false,\n                    \"attributes\": []\n                },\n                \"status\": {\n                    \"name\": \"status\",\n                    \"isArray\": false,\n                    \"type\": {\n                        \"enum\": \"Status\"\n                    },\n                    \"isRequired\": false,\n                    \"attributes\": []\n                },\n                \"createdAt\": {\n                    \"name\": \"createdAt\",\n                    \"isArray\": false,\n                    \"type\": \"AWSDateTime\",\n                    \"isRequired\": false,\n                    \"attributes\": [],\n                    \"isReadOnly\": true\n                },\n                \"updatedAt\": {\n                    \"name\": \"updatedAt\",\n                    \"isArray\": false,\n                    \"type\": \"AWSDateTime\",\n                    \"isRequired\": false,\n                    \"attributes\": [],\n                    \"isReadOnly\": true\n                }\n            },\n            \"syncable\": true,\n            \"pluralName\": \"Blogs\",\n            \"attributes\": [\n                {\n                    \"type\": \"model\",\n                    \"properties\": {}\n                },\n                {\n                    \"type\": \"auth\",\n                    \"properties\": {\n                        \"rules\": [\n                            {\n                                \"allow\": \"public\",\n                                \"operations\": [\n                                    \"create\",\n                                    \"update\",\n                                    \"delete\",\n                                    \"read\"\n                                ]\n                            }\n                        ]\n                    }\n                }\n            ]\n        }\n    },\n    \"enums\": {\n        \"Status\": {\n            \"name\": \"Status\",\n            \"values\": [\n                \"ENABLED\",\n                \"DISABLED\"\n            ]\n        }\n    },\n    \"nonModels\": {},\n    \"version\": \"demo\"\n};`,
+);
+
+describe('should sync amplify backend models', () => {
+  let context: any;
+
+  beforeEach(() => {
+    context = {
+      exeInfo: {
+        projectConfig: {
+          javascript: {
+            config: {
+              SourceDir: 'src',
+            },
+          },
+        },
+      },
+      amplify: {
+        invokePluginMethod: () => ({}),
+      },
+      parameters: {
+        argv: [],
+      },
+      input: {
+        options: {
+          appId: 'testAppId',
+          envName: 'testEnvName',
+        },
+      },
+    };
+  });
+
+  it('should getAmplifyBackendModels', async () => {
+    awsMock.AmplifyBackend = jest.fn(() => ({
+      getBackendAPIModels: jest.fn(() => ({
+        promise: () => ({
+          Models: getMockedSchema(),
+        }),
+      })),
+    }));
+    const client = await AmplifyStudioClient.setClientInfo(context);
+    const dataSchema = await getAmplifyDataSchema(client);
+    expect(getMockedSchema).toBeCalled();
+    expect(dataSchema).toBeDefined();
+    expect(Object.keys(dataSchema!.models)).toContain('Blog');
+  });
+
+  it('should handle Models not found', async () => {
+    awsMock.AmplifyBackend = jest.fn(() => ({
+      getBackendAPIModels: jest.fn(() => ({
+        promise: jest.fn(() => ({
+          Models: undefined,
+        })),
+      })),
+    }));
+    const client = await AmplifyStudioClient.setClientInfo(context);
+    const dataSchema = await getAmplifyDataSchema(client);
+    expect(dataSchema).toBeUndefined();
+    expect(printerMock.debug).toHaveBeenCalledWith('Provided ResourceName: MyResourceName did not yield Models.');
+  });
+
+  it('should handle network error', async () => {
+    awsMock.AmplifyBackend = jest.fn(() => ({
+      getBackendAPIModels: jest.fn(() => ({
+        promise: jest.fn().mockRejectedValue(new Error('Network Error')),
+      })),
+    }));
+    const client = await AmplifyStudioClient.setClientInfo(context);
+    const dataSchema = await getAmplifyDataSchema(client);
+    expect(dataSchema).toBeUndefined();
+    expect(printerMock.debug).toHaveBeenCalledWith('Error: Models not found in AmplifyBackend:GetBackendAPIModels response: Network Error');
+  });
+});

package/src/__tests__/syncAmplifyUiComponents.test.ts

@@ -1,24 +1,29 @@
-import { $TSContext } from 'amplify-cli-core';
+import { AmplifyCategories, AmplifySupportedService, stateManager } from 'amplify-cli-core'; // eslint-disable-line import/no-extraneous-dependencies
+import aws from 'aws-sdk'; // eslint-disable-line import/no-extraneous-dependencies
 import {
   generateUiBuilderComponents,
+  generateUiBuilderThemes,
   getEnvName,
   getAppId,
-  listUiBuilderComponents,
-  listUiBuilderThemes,
   resolveAppId,
-  generateUiBuilderThemes,
-} from '../commands/utils/syncAmplifyUiBuilderComponents';
-import { getAmplifyUIBuilderService } from '../commands/utils/amplifyUiBuilderService'
-import aws from 'aws-sdk';
-import * as CLICore from 'amplify-cli-core';
-import * as createUiBuilderComponentDependency from '../commands/utils/createUiBuilderComponent';
-const aws_mock = aws as any;
-const createUiBuilderComponentDependency_mock = createUiBuilderComponentDependency as any;
+} from '../commands/utils';
+import { AmplifyStudioClient } from '../clients';
+import * as createUiBuilderComponentDependency from '../commands/utils/codegenResources';
+
+jest.mock('amplify-cli-core', () => ({
+  ...jest.requireActual('amplify-cli-core'),
+  FeatureFlags: {
+    getBoolean: () => false,
+    getNumber: () => 0,
+  },
+}));
+
+const awsMock = aws as any;
+const stateManagerMock = stateManager as any;
+const createUiBuilderComponentDependencyMock = createUiBuilderComponentDependency as any;
 describe('should sync amplify ui builder components', () => {
-  let context: $TSContext | any;
+  let context: any;
   beforeEach(() => {
-    CLICore.FeatureFlags.getBoolean = () => false;
-    CLICore.FeatureFlags.getNumber = () => 0;
     context = {
       exeInfo: {
         projectConfig: {
@@ -43,7 +48,18 @@ describe('should sync amplify ui builder components', () => {
       },
     };
 
-    aws_mock.AmplifyUIBuilder = jest.fn(() => ({
+    stateManagerMock.getMeta = jest.fn(() => ({
+      [AmplifyCategories.API]: {
+        MyResourceName: {
+          service: AmplifySupportedService.APPSYNC,
+        },
+      },
+      providers: {
+        awscloudformation: { AmplifyAppId: 'testAppId' },
+      },
+    }));
+
+    awsMock.AmplifyUIBuilder = jest.fn(() => ({
       exportComponents: jest.fn(() => ({
         promise: jest.fn(() => ({
           entities: [
@@ -68,8 +84,8 @@ describe('should sync amplify ui builder components', () => {
       })),
     }));
 
-    createUiBuilderComponentDependency_mock.createUiBuilderComponent = jest.fn();
-    createUiBuilderComponentDependency_mock.createUiBuilderTheme = jest.fn();
+    createUiBuilderComponentDependencyMock.createUiBuilderComponent = jest.fn();
+    createUiBuilderComponentDependencyMock.createUiBuilderTheme = jest.fn();
   });
 
   it('pulls components from aws-sdk and passes them to createUiBuilderComponent', () => {
@@ -77,14 +93,14 @@ describe('should sync amplify ui builder components', () => {
   });
 
   it('does not throw an error when createUiBuilderComponent fails', () => {
-    createUiBuilderComponentDependency_mock.createUiBuilderComponent = jest.fn(() => {
+    createUiBuilderComponentDependencyMock.createUiBuilderComponent = jest.fn(() => {
       throw new Error('ahhh!');
     });
     expect(async () => generateUiBuilderComponents(context, [])).not.toThrow();
   });
 
   it('does not throw an error when createUiBuilderThemes fails', () => {
-    createUiBuilderComponentDependency_mock.createUiBuilderComponent = jest.fn(() => {
+    createUiBuilderComponentDependencyMock.createUiBuilderComponent = jest.fn(() => {
       throw new Error('ahhh!');
     });
     expect(async () => generateUiBuilderThemes(context, [])).not.toThrow();
@@ -93,20 +109,26 @@ describe('should sync amplify ui builder components', () => {
   it('can getAmplifyUIBuilderService', async () => {
     process.env.UI_BUILDER_ENDPOINT = 'https://mock-endpoint.com';
     process.env.UI_BUILDER_REGION = 'mock-region';
-    const service = await getAmplifyUIBuilderService(context, 'testEnv', 'testAppId');
-    expect(Object.keys(service)).toContain('exportComponents');
-    expect(Object.keys(service)).toContain('exportThemes');
+    const client = await AmplifyStudioClient.setClientInfo(context);
+    expect(Object.keys(client)).toEqual(expect.arrayContaining([
+      'listComponents',
+      'listThemes',
+      'getModels',
+      'createComponent',
+    ]));
   });
-  it('can listUiBuilderThemes', async () => {
-    const themes = await listUiBuilderThemes(context, 'testEnv');
+  it('can list themes', async () => {
+    const client = await AmplifyStudioClient.setClientInfo(context);
+    const themes = await client.listThemes();
     expect(themes.entities).toHaveLength(1);
   });
   it('can listUiBuilderComponents', async () => {
-    const components = await listUiBuilderComponents(context, 'testEnv');
+    const client = await AmplifyStudioClient.setClientInfo(context);
+    const components = await client.listThemes();
     expect(components.entities).toHaveLength(1);
   });
   it('can getAppId', async () => {
-    const appId = await getAppId(context);
+    const appId = getAppId(context);
     expect(appId).toBe('testAppId');
   });
   it('can getEnvName', () => {
@@ -114,34 +136,34 @@ describe('should sync amplify ui builder components', () => {
     expect(envName).toBe('testEnvName');
   });
   it('can resolveAppId', async () => {
-    context.amplify.invokePluginMethod = () => 'testAppId';
-    const appId = await resolveAppId(context);
+    const appId = resolveAppId();
     expect(appId).toBe('testAppId');
   });
   it('can throw on getAppId', async () => {
     context.input.options.appId = null;
     context.amplify.invokePluginMethod = () => null;
-    expect(async () => await getAppId(context)).rejects.toThrowError();
+    stateManagerMock.getMeta = () => ({});
+    expect(() => getAppId(context)).toThrowError();
   });
   it('can generate ui builder components', async () => {
-    createUiBuilderComponentDependency_mock.createUiBuilderComponent = jest.fn().mockImplementation(() => ({}));
+    createUiBuilderComponentDependencyMock.createUiBuilderComponent = jest.fn().mockImplementation(() => ({}));
     const components = generateUiBuilderComponents(context, [{}, {}]);
     expect(components.every(component => component.resultType === 'SUCCESS')).toBeTruthy();
   });
   it('can handle failed generation generate ui builder components', async () => {
-    createUiBuilderComponentDependency_mock.createUiBuilderComponent = jest.fn().mockImplementation(() => {
+    createUiBuilderComponentDependencyMock.createUiBuilderComponent = jest.fn().mockImplementation(() => {
       throw new Error('ahh!');
     });
     const components = generateUiBuilderComponents(context, [{}, {}]);
     expect(components.every(component => component.resultType === 'FAILURE')).toBeTruthy();
   });
   it('can generate ui builder themes', async () => {
-    createUiBuilderComponentDependency_mock.createUiBuilderTheme = jest.fn().mockImplementation(() => ({}));
+    createUiBuilderComponentDependencyMock.createUiBuilderTheme = jest.fn().mockImplementation(() => ({}));
     const themes = generateUiBuilderThemes(context, [{}, {}]);
     expect(themes.every(theme => theme.resultType === 'SUCCESS')).toBeTruthy();
   });
   it('can handle failed generation generate ui builder themes', async () => {
-    createUiBuilderComponentDependency_mock.createUiBuilderTheme = jest.fn().mockImplementation(() => {
+    createUiBuilderComponentDependencyMock.createUiBuilderTheme = jest.fn().mockImplementation(() => {
       throw new Error('ahh!');
     });
     const themes = generateUiBuilderThemes(context, [{}, {}]);

package/src/clients/amplify-studio-client.ts

@@ -0,0 +1,161 @@
+import type { $TSContext } from 'amplify-cli-core';
+import type { ServiceConfigurationOptions } from 'aws-sdk/lib/service';
+import type { CreateComponentData, Component, Theme } from 'aws-sdk/clients/amplifyuibuilder';
+import { AmplifyUIBuilder, AmplifyBackend } from 'aws-sdk';
+import { printer } from 'amplify-prompts';
+import { getAppId, getEnvName } from '../commands/utils/environmentHelpers';
+
+/**
+ * Builds the Amplify Backend Client
+ */
+const buildAmplifyBackendClient = (awsConfigInfo: ServiceConfigurationOptions): AmplifyBackend => {
+  const awsConfig = { ...awsConfigInfo };
+  if (process.env.AMPLIFY_BACKEND_ENDPOINT) {
+    awsConfig.endpoint = process.env.AMPLIFY_BACKEND_ENDPOINT;
+  }
+
+  if (process.env.AMPLIFY_BACKEND_REGION) {
+    awsConfig.region = process.env.AMPLIFY_BACKEND_REGION;
+  }
+
+  return new AmplifyBackend(awsConfig);
+};
+
+/**
+ * Builds the Amplify UIBuilder Client
+ */
+const buildAmplifyUiBuilderClient = (awsConfigInfo: ServiceConfigurationOptions): AmplifyUIBuilder => {
+  const awsConfig = { ...awsConfigInfo };
+  if (process.env.UI_BUILDER_ENDPOINT) {
+    awsConfig.endpoint = process.env.UI_BUILDER_ENDPOINT;
+  }
+
+  if (process.env.UI_BUILDER_REGION) {
+    awsConfig.region = process.env.UI_BUILDER_REGION;
+  }
+
+  return new AmplifyUIBuilder(awsConfig);
+};
+
+/**
+ * used for Amplify Studio related clients
+ */
+export default class AmplifyStudioClient {
+  #amplifyUiBuilder: AmplifyUIBuilder;
+  #amplifyBackend: AmplifyBackend;
+  #appId: string;
+  #envName: string;
+  /**
+   * Used to configure the AWS Amplify clients.
+   */
+  static async setClientInfo(context: $TSContext, envName?: string, appId?: string): Promise<AmplifyStudioClient> {
+    const resolvedEnvName = getEnvName(context, envName);
+    const resolvedAppId = getAppId(context, appId);
+    const awsConfigInfo = (await context.amplify.invokePluginMethod(context, 'awscloudformation', undefined, 'loadConfigurationForEnv', [
+      context,
+      resolvedEnvName,
+      resolvedAppId,
+    ])) as ServiceConfigurationOptions;
+    return new AmplifyStudioClient(awsConfigInfo, resolvedAppId, resolvedEnvName);
+  }
+
+  /*
+  * Builds the clients
+  */
+  constructor(awsConfigInfo: ServiceConfigurationOptions, appId: string, envName: string) {
+    this.#amplifyUiBuilder = buildAmplifyUiBuilderClient(awsConfigInfo);
+    this.#amplifyBackend = buildAmplifyBackendClient(awsConfigInfo);
+    this.#appId = appId;
+    this.#envName = envName;
+  }
+
+  /**
+   * listComponents from Studio
+   */
+  listComponents = async (envName?: string, appId?: string): Promise<{ entities: Component[] }> => {
+    const environmentName = envName || this.#envName;
+    const resolvedAppId = appId || this.#appId;
+    try {
+      let nextToken: string | undefined;
+      const uiBuilderComponents: Component[] = [];
+      do {
+        const response = await this.#amplifyUiBuilder
+          .exportComponents({
+            appId: resolvedAppId,
+            environmentName,
+            nextToken,
+          })
+          .promise();
+        uiBuilderComponents.push(...response.entities);
+        nextToken = response.nextToken;
+      } while (nextToken);
+      printer.debug(JSON.stringify(uiBuilderComponents, null, 2));
+      return { entities: uiBuilderComponents };
+    } catch (e) {
+      printer.debug(e);
+      throw e;
+    }
+  }
+
+  /**
+ * Returns all the UI Builder themes from the app
+ */
+ listThemes = async (envName?: string, appId?: string): Promise<{ entities: Theme[] }> => {
+   const environmentName = envName || this.#envName;
+   const resolvedAppId = appId || this.#appId;
+
+   try {
+     let nextToken: string | undefined;
+     const uiBuilderThemes: Theme[] = [];
+     do {
+       const response = await this.#amplifyUiBuilder
+         .exportThemes({
+           appId: resolvedAppId,
+           environmentName,
+           nextToken,
+         })
+         .promise();
+       uiBuilderThemes.push(...response.entities);
+       nextToken = response.nextToken;
+     } while (nextToken);
+     printer.debug(JSON.stringify(uiBuilderThemes, null, 2));
+     return { entities: uiBuilderThemes };
+   } catch (e) {
+     printer.debug(e);
+     throw e;
+   }
+ }
+
+ createComponent = async (component: CreateComponentData, envName?: string, appId?: string): Promise<Component | undefined> => {
+   const environmentName = envName || this.#envName;
+   const resolvedAppId = appId || this.#appId;
+   try {
+     const response = await this.#amplifyUiBuilder.createComponent({
+       appId: resolvedAppId,
+       environmentName,
+       componentToCreate: component,
+     }).promise();
+     return response.entity;
+   } catch (err) {
+     printer.debug(err.toString());
+     throw err;
+   }
+ }
+
+ getModels = async (resourceName: string, envName?: string, appId?: string): Promise<string | undefined> => {
+   try {
+     const environmentName = envName || this.#envName;
+     const resolvedAppId = appId || this.#appId;
+     const { Models } = await this.#amplifyBackend.getBackendAPIModels({
+       AppId: resolvedAppId,
+       BackendEnvironmentName: environmentName,
+       ResourceName: resourceName,
+     })
+       .promise();
+     return Models;
+   } catch (e) {
+     printer.debug(e);
+     throw new Error(`Models not found in AmplifyBackend:GetBackendAPIModels response: ${e.message}`);
+   }
+ }
+}

package/src/clients/index.ts

@@ -0,0 +1 @@
+export { default as AmplifyStudioClient } from './amplify-studio-client';

package/src/commands/cloneComponentsFromEnv.ts

@@ -1,21 +1,25 @@
-import { extractArgs } from './utils/extractArgs';
-import { listUiBuilderComponents } from './utils/syncAmplifyUiBuilderComponents';
 import { printer } from 'amplify-prompts';
 import { $TSContext } from 'amplify-cli-core';
-import { getAmplifyUIBuilderService } from './utils/amplifyUiBuilderService';
+import { extractArgs } from './utils';
+import { AmplifyStudioClient } from '../clients';
 
-export async function run(context: $TSContext) {
+/**
+ * Clones the components from the source to new environment
+ */
+export const run = async (context: $TSContext): Promise<void> => {
   printer.debug('Running create components command in amplify-util-uibuilder');
   const args = extractArgs(context);
   const sourceEnvName = args.sourceEnvName ? args.sourceEnvName : context.exeInfo.sourceEnvName;
   const newEnvName = args.newEnvName ? args.newEnvName : context.exeInfo.localEnvInfo.envName;
+  const appId = args.appId ?? context.exeInfo.teamProviderInfo[sourceEnvName].awscloudformation.AmplifyAppId;
+  const studioClient = await AmplifyStudioClient.setClientInfo(context, sourceEnvName, appId);
 
-  const existingComponents = await listUiBuilderComponents(context, sourceEnvName);
+  const [existingComponents, existingComponentsNewEnv] = await Promise.all([
+    studioClient.listComponents(sourceEnvName), studioClient.listComponents(newEnvName)]);
   if (existingComponents.entities.length === 0) {
     printer.debug(`${existingComponents.entities.length} components exist in source env. Skipping creation of local components.`);
     return;
   }
-  const existingComponentsNewEnv = await listUiBuilderComponents(context, newEnvName);
 
   if (existingComponentsNewEnv.entities.length > 0) {
     printer.debug(
@@ -24,35 +28,28 @@ export async function run(context: $TSContext) {
     return;
   }
 
-  const environmentName = args.environmentName ?? context.exeInfo.localEnvInfo.envName;
-  const appId = args.appId ?? context.exeInfo.teamProviderInfo[environmentName].awscloudformation.AmplifyAppId;
-  const amplifyUIBuilder = await getAmplifyUIBuilderService(context, environmentName, appId);
   const components = existingComponents.entities;
   if (!components.length) {
     printer.debug(`No UIBuilder components found in app ${appId} from env ${sourceEnvName}. Skipping component clone process.`);
     return;
   }
-  const promises = components.map(async (component: any) => {
-    return await amplifyUIBuilder
-      .createComponent({
-        appId,
-        environmentName,
-        componentToCreate: {
-          bindingProperties: component.bindingProperties,
-          children: component.children,
-          componentType: component.componentType,
-          name: component.name,
-          overrides: component.overrides,
-          properties: component.properties,
-          sourceId: component.sourceId,
-          variants: component.variants,
-        },
-      })
-      .promise();
-  });
-  await Promise.all(promises);
+  for (let i = 0; i < components.length; i++) {
+    const {
+      appId: _appId, // eslint-disable-line @typescript-eslint/no-unused-vars
+      environmentName, // eslint-disable-line @typescript-eslint/no-unused-vars
+      id, // eslint-disable-line @typescript-eslint/no-unused-vars
+      createdAt, // eslint-disable-line @typescript-eslint/no-unused-vars
+      modifiedAt, // eslint-disable-line @typescript-eslint/no-unused-vars
+      ...componentCreateData
+    } = components[i];
+    await studioClient.createComponent(
+      componentCreateData,
+      newEnvName,
+      appId,
+    );
+  }
 
   printer.info(
     `Successfully cloned ${components.length} UIBuilder components in app ${appId} from env ${sourceEnvName} to env ${newEnvName}.`,
   );
-}
+};

package/src/commands/generateComponents.ts

@@ -1,35 +1,38 @@
 import { StudioComponent } from '@aws-amplify/codegen-ui';
-import { printer } from 'amplify-prompts';
 import ora from 'ora';
+import { printer } from 'amplify-prompts';
 import { $TSContext } from 'amplify-cli-core';
-import { notifyMissingPackages } from './utils/notifyMissingPackages';
-import { shouldRenderComponents } from './utils/shouldRenderComponents';
+import { AmplifyStudioClient } from '../clients';
 import {
-  listUiBuilderComponents,
-  listUiBuilderThemes,
+  notifyMissingPackages,
+  shouldRenderComponents,
   generateUiBuilderComponents,
   generateUiBuilderThemes,
-} from './utils/syncAmplifyUiBuilderComponents';
-import { generateAmplifyUiBuilderIndexFile } from './utils/createUiBuilderComponent';
+  getAmplifyDataSchema,
+  generateAmplifyUiBuilderIndexFile,
+} from './utils';
 
 /**
  * Pulls ui components from Studio backend and generates the code in the user's file system
  */
 export const run = async (context: $TSContext): Promise<void> => {
-  printer.debug('Running generate components command in amplify-util-uibuilder');
   if (!(await shouldRenderComponents(context))) {
     return;
   }
   const spinner = ora('');
   try {
-    const [componentSchemas, themeSchemas] = await Promise.all([listUiBuilderComponents(context), listUiBuilderThemes(context)]);
+    const studioClient = await AmplifyStudioClient.setClientInfo(context);
+    const [componentSchemas, themeSchemas, dataSchema] = await Promise.all([
+      studioClient.listComponents(),
+      studioClient.listThemes(),
+      getAmplifyDataSchema(studioClient),
+    ]);
     if (componentSchemas.entities.length === 0 && themeSchemas.entities.length === 0) {
       printer.debug('Skipping UI component generation since none are found.');
       return;
     }
-
     spinner.start('Generating UI components...');
-    const generatedComponentResults = generateUiBuilderComponents(context, componentSchemas.entities);
+    const generatedComponentResults = generateUiBuilderComponents(context, componentSchemas.entities, dataSchema);
     const generatedThemeResults = generateUiBuilderThemes(context, themeSchemas.entities);
 
     generateAmplifyUiBuilderIndexFile(context, [
@@ -50,7 +53,11 @@ export const run = async (context: $TSContext): Promise<void> => {
 
     const invalidComponentNames = componentSchemas.entities.filter(component => !component.schemaVersion).map(component => component.name);
     if (invalidComponentNames.length) {
-      printer.warn(`The components ${invalidComponentNames.join(', ')} were synced with an older version of Amplify Studio. Please re-sync your components with Figma to get latest features and changes.`);
+      printer.warn(
+        `The components ${invalidComponentNames.join(
+          ', ',
+        )} were synced with an older version of Amplify Studio. Please re-sync your components with Figma to get latest features and changes.`, // eslint-disable-line spellcheck/spell-checker
+      );
     }
 
     notifyMissingPackages(context);