From c54ed9bbb5ff47524620a34c5307fd24d487a592 Mon Sep 17 00:00:00 2001
From: Vaisshnavi <vaisshv@amazon.com>
Date: Tue, 14 Jan 2025 20:20:04 -0800
Subject: [PATCH] Fix: Generate nested CustomTypes in schema output

---
 packages/data-schema/src/SchemaProcessor.ts | 34 ++++++++++++++++-----
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/packages/data-schema/src/SchemaProcessor.ts b/packages/data-schema/src/SchemaProcessor.ts
index 966dbdf59..6718b5449 100644
--- a/packages/data-schema/src/SchemaProcessor.ts
+++ b/packages/data-schema/src/SchemaProcessor.ts
@@ -1397,7 +1397,7 @@ function generateInputTypes(
 
   implicitTypes.forEach(([typeName, typeDef]) => {
     if (isCustomType(typeDef)) {
-      const { gqlFields } = processFields(
+      const { gqlFields, implicitTypes: nestedTypes } = processFields(
         typeName,
         typeDef.data.fields,
         {},
@@ -1416,6 +1416,20 @@ function generateInputTypes(
       const typeKeyword = generateInputType ? 'input' : 'type';
       const customType = `${typeKeyword} ${typeName}${authString ? ` ${authString}` : ''}\n{\n  ${gqlFields.join('\n  ')}\n}`;
       generatedTypes.add(customType);
+
+      // Process nested types
+      if (nestedTypes.length > 0) {
+        const nestedGeneratedTypes = generateInputTypes(
+          nestedTypes,
+          generateInputType,
+          getRefType,
+          authRules,
+          true,
+        );
+        nestedGeneratedTypes.forEach((type) => {
+          generatedTypes.add(type);
+        });
+      }
     } else if (typeDef.type === 'enum') {
       const enumDefinition = `enum ${typeName} {\n  ${typeDef.values.join('\n  ')}\n}`;
       generatedTypes.add(enumDefinition);
@@ -1441,7 +1455,7 @@ const schemaPreprocessor = (
 } => {
   const enumTypes = new Set<string>();
   const gqlModels: string[] = [];
-  const inputTypes: string[] = [];
+  const inputTypes = new Map<string, string>();
 
   const customQueries = [];
   const customMutations = [];
@@ -1578,7 +1592,10 @@ const schemaPreprocessor = (
           databaseType,
           getRefType,
         );
-        inputTypes.push(...operationInputTypes);
+        operationInputTypes.forEach((type) => {
+          const typeName = type.split(' ')[1];
+          inputTypes.set(typeName, type);
+        });
         gqlModels.push(...operationReturnTypes);
 
         /**
@@ -1601,13 +1618,16 @@ const schemaPreprocessor = (
             const shouldGenerateInputType = !operationReturnTypes.some(
               (returnType) => returnType.includes(typeName),
             );
-            const generatedTypeDefinition = generateInputTypes(
+            const generatedTypeDefinitions = generateInputTypes(
               [[typeName, typeDef]],
               shouldGenerateInputType,
               getRefType,
-            )[0];
+            );
             if (shouldGenerateInputType) {
-              inputTypes.push(generatedTypeDefinition);
+              generatedTypeDefinitions.forEach((def) => {
+                const defName = def.split(' ')[1];
+                inputTypes.set(defName, def);
+              });
             }
           }
         });
@@ -1801,7 +1821,7 @@ const schemaPreprocessor = (
     schemaComponents.push(CONVERSATION_SCHEMA_GRAPHQL_TYPES);
   }
 
-  schemaComponents.push(...inputTypes);
+  schemaComponents.push(...inputTypes.values());
 
   const processedSchema = schemaComponents.join('\n\n');