feat: include all envs in gen 2 data codegen

packages/amplify-cli-core/API.md

@@ -2204,7 +2204,7 @@ export function validateExportDirectoryPath(directoryPath: any, defaultPath: str
 export class ViewResourceTableParams {
     constructor(cliParams: CLIParams);
     // (undocumented)
-    get categoryList(): string[] | [];
+    get categoryList(): [] | string[];
     // (undocumented)
     get command(): string;
     // (undocumented)

packages/amplify-gen1-codegen-data-adapter/API.md

@@ -4,11 +4,11 @@
 
 ```ts
 
-import { DataDefinition } from '@aws-amplify/amplify-gen2-codegen';
+import { DataTableMapping } from '@aws-amplify/amplify-gen2-codegen';
 import { Stack } from '@aws-sdk/client-cloudformation';
 
 // @public (undocumented)
-export const getDataDefinition: (dataStack: Stack) => DataDefinition;
+export const getDataDefinition: (dataStack: Stack) => DataTableMapping;
 
 // (No @packageDocumentation comment for this package)
 

packages/amplify-gen1-codegen-data-adapter/src/get_data_definition.test.ts

@@ -13,6 +13,6 @@ describe('Data definition', () => {
       ],
     } as Stack;
     const result = getDataDefinition(stack);
-    assert.equal(result.tableMapping.hello, 'world');
+    assert.equal(result.hello, 'world');
   });
 });

packages/amplify-gen1-codegen-data-adapter/src/get_data_definition.ts

@@ -1,12 +1,12 @@
 import assert from 'node:assert';
-import { DataDefinition } from '@aws-amplify/amplify-gen2-codegen';
+import { DataTableMapping } from '@aws-amplify/amplify-gen2-codegen';
 import { Stack } from '@aws-sdk/client-cloudformation';
 
 export const tableMappingKey = 'DataSourceMappingOutput';
 
-export const getDataDefinition = (dataStack: Stack): DataDefinition => {
+export const getDataDefinition = (dataStack: Stack): DataTableMapping => {
   const rawTableMapping = dataStack.Outputs?.find((o) => o.OutputKey === tableMappingKey)?.OutputValue;
   assert(rawTableMapping);
   const tableMapping = JSON.parse(rawTableMapping);
-  return { tableMapping };
+  return tableMapping;
 };

packages/amplify-gen2-codegen/API.md

@@ -81,9 +81,12 @@ export type CustomAttributes = Partial<Record<`custom:${string}`, CustomAttribut
 
 // @public (undocumented)
 export type DataDefinition = {
-    tableMapping: Record<string, string>;
+    tableMappings: Record<string, DataTableMapping | undefined>;
 };
 
+// @public (undocumented)
+export type DataTableMapping = Record<string, string>;
+
 // @public (undocumented)
 export type EmailOptions = {
     emailVerificationBody: string;

packages/amplify-gen2-codegen/src/data/source_builder.test.ts

@@ -16,23 +16,19 @@ describe('Data Category code generation', () => {
   });
   describe('import map', () => {
     it('is rendered', () => {
-      const tableMapping = { Todo: 'my-todo-mapping' };
-      const source = printNodeArray(generateDataSource({ tableMapping }));
-      assert.match(source, /importedAmplifyDynamoDBTableMap: \{\s+Todo: ['"]my-todo-mapping['"]/);
-    });
-    it('shows each key in the mapping table in the `importedModels` array', () => {
-      const tables = ['Todo', 'Foo', 'Bar'];
-      const tableMapping = tables.reduce((prev, curr) => ({ ...prev, [curr]: 'baz' }), {});
-      const source = printNodeArray(generateDataSource({ tableMapping }));
-      const array = source.match(/importedModels:\s+\[(.*?)\]/);
-      assert.deepEqual(tables, array?.[1].replaceAll('"', '').split(', '));
+      const tableMappings = { dev: { Todo: 'my-todo-mapping' } };
+      const source = printNodeArray(generateDataSource({ tableMappings }));
+      assert.match(
+        source,
+        /importedAmplifyDynamoDBTableMap: \{\s+\n\s+\/\/ replace the environment name \(dev\) with the corresponding branch name\n\s+dev: { Todo: ['"]my-todo-mapping['"] } }/,
+      );
     });
     it('has each each key in defineData', () => {
-      const tableMapping = { Todo: 'my-todo-mapping' };
-      const source = printNodeArray(generateDataSource({ tableMapping }));
+      const tableMappings = { dev: { Todo: 'my-todo-mapping' } };
+      const source = printNodeArray(generateDataSource({ tableMappings }));
       assert.match(
         source,
-        /defineData\({\n\s+importedAmplifyDynamoDBTableMap: \{\s+Todo: ['"]my-todo-mapping['"] },\n\s+importedModels:\s+\[.*?\],\n\s+schema: "TODO: Add your existing graphql schema here"\n}\)/,
+        /defineData\({\n\s+importedAmplifyDynamoDBTableMap: \{\s+\n\s+\/\/ replace the environment name \(dev\) with the corresponding branch name\n\s+dev: { Todo: ['"]my-todo-mapping['"] } },\n\s+schema: "TODO: Add your existing graphql schema here"\n}\)/,
       );
     });
   });

packages/amplify-gen2-codegen/src/data/source_builder.ts

@@ -1,10 +1,12 @@
 import ts, { ObjectLiteralElementLike } from 'typescript';
 import { renderResourceTsFile } from '../resource/resource';
 import { createTodoError } from '../todo_error';
+import { table } from 'console';
 const factory = ts.factory;
 
+export type DataTableMapping = Record<string, string>;
 export type DataDefinition = {
-  tableMapping: Record<string, string>;
+  tableMappings: Record<string, DataTableMapping | undefined>;
 };
 
 const importedAmplifyDynamoDBTableMapKeyName = 'importedAmplifyDynamoDBTableMap';
@@ -17,25 +19,43 @@ export const generateDataSource = (dataDefinition?: DataDefinition): ts.NodeArra
   const namedImports: Record<string, Set<string>> = { '@aws-amplify/backend': new Set() };
   namedImports['@aws-amplify/backend'].add('defineData');
 
-  if (dataDefinition?.tableMapping) {
-    const tableMappingProperties: ObjectLiteralElementLike[] = [];
-    for (const [tableName, tableId] of Object.entries(dataDefinition.tableMapping)) {
-      tableMappingProperties.push(
-        factory.createPropertyAssignment(factory.createIdentifier(tableName), factory.createStringLiteral(tableId)),
+  if (dataDefinition?.tableMappings) {
+    const tableMappingEnvironments: ObjectLiteralElementLike[] = [];
+    for (const [environmentName, tableMapping] of Object.entries(dataDefinition.tableMappings)) {
+      const tableMappingProperties: ObjectLiteralElementLike[] = [];
+      if (tableMapping) {
+        for (const [tableName, tableId] of Object.entries(tableMapping)) {
+          tableMappingProperties.push(
+            factory.createPropertyAssignment(factory.createIdentifier(tableName), factory.createStringLiteral(tableId)),
+          );
+        }
+      }
+
+      let tableMappingExpression = factory.createObjectLiteralExpression(tableMappingProperties);
+      if (tableMappingProperties.length === 0) {
+        tableMappingExpression = ts.addSyntheticLeadingComment(
+          ts.addSyntheticLeadingComment(tableMappingExpression, ts.SyntaxKind.SingleLineCommentTrivia, '', true),
+          ts.SyntaxKind.MultiLineCommentTrivia,
+          ' Unable to find the table mapping for this environment.\n' +
+            ' This could be due the enableGen2Migration feature flag not being set to true for this environment.\n' +
+            ' Please enable the feature flag and push the backend resources.\n' +
+            ' If you are not planning to migrate this environment, you can remove this key',
+          true,
+        );
+      }
+      tableMappingEnvironments.push(
+        ts.addSyntheticLeadingComment(
+          factory.createPropertyAssignment(factory.createIdentifier(environmentName), tableMappingExpression),
+          ts.SyntaxKind.SingleLineCommentTrivia,
+          ` replace the environment name (${environmentName}) with the corresponding branch name`,
+          true,
+        ),
       );
     }
     dataRenderProperties.push(
       factory.createPropertyAssignment(
         importedAmplifyDynamoDBTableMapKeyName,
-        factory.createObjectLiteralExpression(tableMappingProperties),
-      ),
-    );
-    dataRenderProperties.push(
-      factory.createPropertyAssignment(
-        importedModelsKey,
-        factory.createArrayLiteralExpression(
-          Object.keys(dataDefinition.tableMapping).map((tableName) => factory.createStringLiteral(tableName)),
-        ),
+        factory.createObjectLiteralExpression(tableMappingEnvironments),
       ),
     );
   }

packages/amplify-gen2-codegen/src/index.ts

@@ -43,7 +43,7 @@ import {
   ServerSideEncryptionConfiguration,
 } from './storage/source_builder.js';
 
-import { DataDefinition, generateDataSource } from './data/source_builder';
+import { DataDefinition, DataTableMapping, generateDataSource } from './data/source_builder';
 
 import { FunctionDefinition, renderFunctions } from './function/source_builder';
 
@@ -242,6 +242,7 @@ export {
   AuthLambdaTriggers,
   StorageTriggerEvent,
   DataDefinition,
+  DataTableMapping,
   SamlOptions,
   OidcEndPoints,
   MetadataOptions,

packages/amplify-migration/src/backend_environment_selector.ts

@@ -1,5 +1,5 @@
 import assert from 'node:assert';
-import { AmplifyClient, BackendEnvironment, GetBackendEnvironmentCommand } from '@aws-sdk/client-amplify';
+import { AmplifyClient, BackendEnvironment, GetBackendEnvironmentCommand, ListBackendEnvironmentsCommand } from '@aws-sdk/client-amplify';
 import { getEnvInfo } from '@aws-amplify/cli-internal/lib/extensions/amplify-helpers/get-env-info';
 
 export class BackendEnvironmentResolver {
@@ -19,4 +19,16 @@ export class BackendEnvironmentResolver {
     this.selectedEnvironment = backendEnvironment;
     return this.selectedEnvironment;
   };
+
+  getAllBackendEnvironments = async (): Promise<BackendEnvironment[]> => {
+    const envInfo = getEnvInfo();
+    assert(envInfo);
+    const { backendEnvironments } = await this.amplifyClient.send(
+      new ListBackendEnvironmentsCommand({
+        appId: this.appId,
+      }),
+    );
+    assert(backendEnvironments, 'No backend environments found');
+    return backendEnvironments;
+  };
 }

packages/amplify-migration/src/data_definition_fetcher.test.ts

@@ -11,10 +11,13 @@ describe('DataDefinitionFetcher', () => {
       it('maps cloudformation stack output to table mapping', async () => {
         const mapping = { hello: 'world' };
         const mockBackendEnvResolver: BackendEnvironmentResolver = {
-          selectBackendEnvironment: async () => {
-            return {
-              stackName: 'asdf',
-            } as BackendEnvironment;
+          getAllBackendEnvironments: async () => {
+            return [
+              {
+                environmentName: 'dev',
+                stackName: 'asdf',
+              },
+            ] as BackendEnvironment[];
           },
         } as BackendEnvironmentResolver;
         const mockAmplifyStackParser: AmplifyStackParser = {
@@ -32,14 +35,17 @@ describe('DataDefinitionFetcher', () => {
         } as unknown as AmplifyStackParser;
         const dataDefinitionFetcher = new DataDefinitionFetcher(mockBackendEnvResolver, mockAmplifyStackParser);
         const results = await dataDefinitionFetcher.getDefinition();
-        assert(results?.tableMapping);
+        assert(results?.tableMappings);
       });
-      it('throws an error if the json cannot be parsed', async () => {
+      it('return undefined for mapping if json cannot be parsed', async () => {
         const mockBackendEnvResolver: BackendEnvironmentResolver = {
-          selectBackendEnvironment: async () => {
-            return {
-              stackName: 'asdf',
-            } as BackendEnvironment;
+          getAllBackendEnvironments: async () => {
+            return [
+              {
+                environmentName: 'dev',
+                stackName: 'asdf',
+              },
+            ] as BackendEnvironment[];
           },
         } as BackendEnvironmentResolver;
         const mockAmplifyStackParser: AmplifyStackParser = {
@@ -56,16 +62,21 @@ describe('DataDefinitionFetcher', () => {
             } as AmplifyStacks),
         } as unknown as AmplifyStackParser;
         const dataDefinitionFetcher = new DataDefinitionFetcher(mockBackendEnvResolver, mockAmplifyStackParser);
-        await assert.rejects(() => dataDefinitionFetcher.getDefinition(), { message: 'Could not parse the Amplify Data table mapping' });
+        const results = await dataDefinitionFetcher.getDefinition();
+        assert(results?.tableMappings);
+        assert.equal(JSON.stringify(results?.tableMappings), JSON.stringify({ dev: undefined }));
       });
     });
     describe('table mapping is not defined', () => {
-      it('reject with table mapping assertion', async () => {
+      it('return undefined for table mapping', async () => {
         const mockBackendEnvResolver: BackendEnvironmentResolver = {
-          selectBackendEnvironment: async () => {
-            return {
-              stackName: 'asdf',
-            } as BackendEnvironment;
+          getAllBackendEnvironments: async () => {
+            return [
+              {
+                environmentName: 'dev',
+                stackName: 'asdf',
+              },
+            ] as BackendEnvironment[];
           },
         } as BackendEnvironmentResolver;
         const mockAmplifyStackParser: AmplifyStackParser = {
@@ -75,17 +86,22 @@ describe('DataDefinitionFetcher', () => {
             } as AmplifyStacks),
         } as unknown as AmplifyStackParser;
         const dataDefinitionFetcher = new DataDefinitionFetcher(mockBackendEnvResolver, mockAmplifyStackParser);
-        await assert.rejects(dataDefinitionFetcher.getDefinition);
+        const results = await dataDefinitionFetcher.getDefinition();
+        assert(results?.tableMappings);
+        assert.equal(JSON.stringify(results?.tableMappings), JSON.stringify({ dev: undefined }));
       });
     });
   });
   describe('if data stack is undefined', () => {
     it('does not reject with table mapping assertion', async () => {
       const mockBackendEnvResolver: BackendEnvironmentResolver = {
-        selectBackendEnvironment: async () => {
-          return {
-            stackName: 'asdf',
-          } as BackendEnvironment;
+        getAllBackendEnvironments: async () => {
+          return [
+            {
+              environmentName: 'dev',
+              stackName: 'asdf',
+            },
+          ] as BackendEnvironment[];
         },
       } as BackendEnvironmentResolver;
       const mockAmplifyStackParser: AmplifyStackParser = {
@@ -97,12 +113,15 @@ describe('DataDefinitionFetcher', () => {
       const dataDefinitionFetcher = new DataDefinitionFetcher(mockBackendEnvResolver, mockAmplifyStackParser);
       await assert.doesNotReject(dataDefinitionFetcher.getDefinition);
     });
-    it('returns undefined', async () => {
+    it('returns undefined for table mapping', async () => {
       const mockBackendEnvResolver: BackendEnvironmentResolver = {
-        selectBackendEnvironment: async () => {
-          return {
-            stackName: 'asdf',
-          } as BackendEnvironment;
+        getAllBackendEnvironments: async () => {
+          return [
+            {
+              environmentName: 'dev',
+              stackName: 'asdf',
+            },
+          ] as BackendEnvironment[];
         },
       } as BackendEnvironmentResolver;
       const mockAmplifyStackParser: AmplifyStackParser = {
@@ -112,8 +131,9 @@ describe('DataDefinitionFetcher', () => {
           } as AmplifyStacks),
       } as unknown as AmplifyStackParser;
       const dataDefinitionFetcher = new DataDefinitionFetcher(mockBackendEnvResolver, mockAmplifyStackParser);
-      const definition = await dataDefinitionFetcher.getDefinition();
-      assert.equal(definition, undefined);
+      const results = await dataDefinitionFetcher.getDefinition();
+      assert(results?.tableMappings);
+      assert.equal(JSON.stringify(results?.tableMappings), JSON.stringify({ dev: undefined }));
     });
   });
 });

packages/amplify-migration/src/data_definition_fetcher.ts

@@ -1,4 +1,3 @@
-import assert from 'node:assert';
 import { DataDefinition } from '@aws-amplify/amplify-gen2-codegen';
 import { AmplifyStackParser } from './amplify_stack_parser.js';
 import { BackendEnvironmentResolver } from './backend_environment_selector.js';
@@ -8,20 +7,29 @@ const dataSourceMappingOutputKey = 'DataSourceMappingOutput';
 export class DataDefinitionFetcher {
   constructor(private backendEnvironmentResolver: BackendEnvironmentResolver, private amplifyStackClient: AmplifyStackParser) {}
   getDefinition = async (): Promise<DataDefinition | undefined> => {
-    const backendEnvironment = await this.backendEnvironmentResolver.selectBackendEnvironment();
-    assert(backendEnvironment?.stackName);
-    const amplifyStacks = await this.amplifyStackClient.getAmplifyStacks(backendEnvironment?.stackName);
-    if (amplifyStacks.dataStack) {
-      const tableMappingText = amplifyStacks.dataStack?.Outputs?.find((o) => o.OutputKey === dataSourceMappingOutputKey)?.OutputValue;
-      assert(tableMappingText, 'Amplify Data table mapping not found.');
-      try {
-        return {
-          tableMapping: JSON.parse(tableMappingText),
-        };
-      } catch (e) {
-        throw new Error('Could not parse the Amplify Data table mapping');
-      }
-    }
-    return undefined;
+    const backendEnvironments = await this.backendEnvironmentResolver.getAllBackendEnvironments();
+    const tableMappings = await Promise.all(
+      backendEnvironments.map(async (backendEnvironment) => {
+        if (!backendEnvironment?.stackName) {
+          return [backendEnvironment.environmentName, undefined];
+        }
+        const amplifyStacks = await this.amplifyStackClient.getAmplifyStacks(backendEnvironment?.stackName);
+        if (amplifyStacks.dataStack) {
+          const tableMappingText = amplifyStacks.dataStack?.Outputs?.find((o) => o.OutputKey === dataSourceMappingOutputKey)?.OutputValue;
+          if (!tableMappingText) {
+            return [backendEnvironment.environmentName, undefined];
+          }
+          try {
+            return [backendEnvironment.environmentName, JSON.parse(tableMappingText)];
+          } catch (e) {
+            return [backendEnvironment.environmentName, undefined];
+          }
+        }
+        return [backendEnvironment.environmentName, undefined];
+      }),
+    );
+    return {
+      tableMappings: Object.fromEntries(tableMappings),
+    };
   };
 }