fix(api): build selection set bottom up (#2763)

* fix(api): build selection set bottom up

* chore(doc): add more documentation

Author: 5d

AmplifyPlugins/API/Tests/APIHostApp/AWSAPIPluginLazyLoadTests/LL1/GraphQLLazyLoadPostComment4V2Tests.swift

@@ -19,8 +19,8 @@ final class GraphQLLazyLoadPostComment4V2Tests: GraphQLLazyLoadBaseTest {
         await setup(withModels: PostComment4V2Models())
         let post = Post(title: "title")
         let comment = Comment(content: "content", post: post)
-        let createdPost = try await mutate(.create(post))
-        let createdComment = try await mutate(.create(comment))
+        try await mutate(.create(post))
+        try await mutate(.create(comment))
     }
 
     // Without `includes` and latest codegenerated types with the model path, the post should be lazy loaded
@@ -117,7 +117,16 @@ final class GraphQLLazyLoadPostComment4V2Tests: GraphQLLazyLoadBaseTest {
         assertLazyReference(comments.first!._post, state: .notLoaded(identifiers: [.init(name: "id", value: post.identifier)]))
     }
 
-    // This looks broken
+
+    /*
+     - Given: Api plugin is cleared
+     - When:
+        - Initialize with PostComment4V2Models
+        - Create post
+        - Create comment with [comment.post.comments.post] association path
+     - Then:
+        - The comment creation request GraphQL Selection Set represents the assocation path
+     */
     func testCommentWithEagerLoadPostAndPostCommentsAndPostCommentsPost() async throws {
         await setup(withModels: PostComment4V2Models())
         let post = Post(title: "title")
@@ -145,25 +154,21 @@ final class GraphQLLazyLoadPostComment4V2Tests: GraphQLLazyLoadBaseTest {
                   updatedAt
                   post {
                     id
+                    createdAt
+                    title
+                    updatedAt
                     __typename
                   }
                   __typename
                 }
               }
-              post {
-                id
-                createdAt
-                title
-                updatedAt
-                __typename
-              }
             }
             __typename
           }
         }
         """
         XCTAssertEqual(request.document, expectedDocument)
-        let createdComment = try await mutate(request)
+        try await mutate(request)
     }
 
     // Without `includes` and latest codegenerated types with the model path, the post's comments should be lazy loaded

AmplifyPlugins/Core/AWSPluginsCore/Model/Decorator/IncludeAssociationDecorator.swift

@@ -35,8 +35,9 @@ public struct IncludeAssociationDecorator: ModelBasedGraphQLDocumentDecorator {
         includedAssociations.forEach { association in
             // we don't include the root reference because it refers to the root model
             // fields in the selection set, only the nested/included ones are needed
-            let associationSelectionSet = association.asSelectionSet(includeRoot: false)
-            selectionSet.merge(with: associationSelectionSet)
+            if let associationSelectionSet = association.asSelectionSet(includeRoot: false) {
+                selectionSet.merge(with: associationSelectionSet)
+            }
         }
 
         return document.copy(selectionSet: selectionSet)
@@ -45,24 +46,48 @@ public struct IncludeAssociationDecorator: ModelBasedGraphQLDocumentDecorator {
 }
 
 extension PropertyContainerPath {
-    
-    func asSelectionSet(includeRoot: Bool = true) -> SelectionSet {
-        let metadata = getMetadata()
-        let modelName = getModelType().modelName
-        guard let schema = ModelRegistry.modelSchema(from: modelName) else {
-            fatalError("Schema for model \(modelName) could not be found.")
+    /// Build GraphQL Selection Set based on Model Associations
+    /// `PropertyContainerPath` is a tree path with leaf node pointer that
+    /// represents the associations from bottom to the top.
+    ///
+    /// - Returns: A Optional<SelectionSet> represents GraphQL Selection Set from top to bottom.
+    func asSelectionSet(includeRoot: Bool = true) -> SelectionSet? {
+        func getSelectionSet(node: PropertyContainerPath) -> SelectionSet {
+            let metadata = node.getMetadata()
+            let modelName = node.getModelType().modelName
+
+            guard let schema = ModelRegistry.modelSchema(from: modelName) else {
+                fatalError("Schema for model \(modelName) could not be found.")
+            }
+            let fieldType: SelectionSetFieldType = metadata.isCollection ? .collection : .model
+
+            let selectionSet = SelectionSet(value: .init(name: metadata.name, fieldType: fieldType))
+            selectionSet.withModelFields(schema.graphQLFields, primaryKeysOnly: true)
+            return selectionSet
+        }
+
+        func shouldProcessNode(node: PropertyContainerPath) -> Bool {
+            includeRoot || node.getMetadata().name != "root"
         }
-        let fieldType: SelectionSetFieldType = metadata.isCollection ? .collection : .model
-        
-        var selectionSet = SelectionSet(value: .init(name: metadata.name, fieldType: fieldType))
-        selectionSet.withModelFields(schema.graphQLFields, primaryKeysOnly: true)
-        if let parent = metadata.parent as? PropertyContainerPath,
-           parent.getMetadata().parent != nil || includeRoot {
-            let parentSelectionSet = parent.asSelectionSet(includeRoot: includeRoot)
-            parentSelectionSet.replaceChild(selectionSet)
-            selectionSet = parentSelectionSet
+
+        func nodesInPath(node: PropertyContainerPath) -> [PropertyContainerPath] {
+            var parent: PropertyContainerPath? = node
+            var path = [PropertyContainerPath]()
+            while let currentNode = parent, shouldProcessNode(node: currentNode) {
+                path.append(currentNode)
+                parent = currentNode.getMetadata().parent as? PropertyContainerPath
+            }
+            return path
+        }
+
+        let selectionSets = nodesInPath(node: self).map(getSelectionSet(node:))
+        return selectionSets.dropFirst().reduce(selectionSets.first) { partialResult, selectionSet in
+            guard let partialResult = partialResult else {
+                return selectionSet
+            }
+            selectionSet.replaceChild(partialResult)
+            return selectionSet
         }
-        return selectionSet
     }
 
 }