diff --git a/packages/amplify-gen2-codegen/src/data/source_builder.test.ts b/packages/amplify-gen2-codegen/src/data/source_builder.test.ts
index 49d4060d3a..74ac6df952 100644
--- a/packages/amplify-gen2-codegen/src/data/source_builder.test.ts
+++ b/packages/amplify-gen2-codegen/src/data/source_builder.test.ts
@@ -27,5 +27,13 @@ describe('Data Category code generation', () => {
       const array = source.match(/importedModels:\s+\[(.*?)\]/);
       assert.deepEqual(tables, array?.[1].replaceAll('"', '').split(', '));
     });
+    it('has each each key in defineData', () => {
+      const tableMapping = { Todo: 'my-todo-mapping' };
+      const source = printNodeArray(generateDataSource({ tableMapping }));
+      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}\)/,
+      );
+    });
   });
 });
diff --git a/packages/amplify-gen2-codegen/src/data/source_builder.ts b/packages/amplify-gen2-codegen/src/data/source_builder.ts
index 0fc8b960ad..23cdaf0de8 100644
--- a/packages/amplify-gen2-codegen/src/data/source_builder.ts
+++ b/packages/amplify-gen2-codegen/src/data/source_builder.ts
@@ -23,14 +23,6 @@ export const generateDataSource = (dataDefinition?: DataDefinition): ts.NodeArra
       tableMappingProperties.push(
         factory.createPropertyAssignment(factory.createIdentifier(tableName), factory.createStringLiteral(tableId)),
       );
-      tableMappingProperties.push(
-        factory.createPropertyAssignment(
-          factory.createIdentifier(importedModelsKey),
-          factory.createArrayLiteralExpression(
-            Object.keys(dataDefinition.tableMapping).map((tableName) => factory.createStringLiteral(tableName)),
-          ),
-        ),
-      );
     }
     dataRenderProperties.push(
       factory.createPropertyAssignment(
@@ -38,6 +30,14 @@ export const generateDataSource = (dataDefinition?: DataDefinition): ts.NodeArra
         factory.createObjectLiteralExpression(tableMappingProperties),
       ),
     );
+    dataRenderProperties.push(
+      factory.createPropertyAssignment(
+        importedModelsKey,
+        factory.createArrayLiteralExpression(
+          Object.keys(dataDefinition.tableMapping).map((tableName) => factory.createStringLiteral(tableName)),
+        ),
+      ),
+    );
   }
   dataRenderProperties.push(
     factory.createPropertyAssignment(factory.createIdentifier('schema'), factory.createStringLiteral(schemaPlaceholderComment)),