fix: handle complex types inside inline fragments in gRPC execution plan

v2/pkg/engine/datasource/grpc_datasource/execution_plan.go

@@ -102,8 +102,8 @@ type RPCMessage struct {
 	Name string
 	// Fields is a list of fields in the message
 	Fields RPCFields
-	// FieldSelectionSet are field selections based on inline fragments
-	FieldSelectionSet RPCFieldSelectionSet
+	// FragmentFields are field selections based on inline fragments
+	FragmentFields RPCFieldSelectionSet
 	// OneOfType indicates the type of the oneof field
 	OneOfType OneOfType
 	// MemberTypes provides the names of the types that are implemented by the Interface or Union
@@ -453,6 +453,57 @@ func (r *rpcPlanningContext) resolveRPCMethodMapping(operationType ast.Operation
 	return rpcConfig, nil
 }
 
+// descendIntoResponseField handles descending into a nested response message
+// when entering a selection set. It branches on whether we are inside an inline
+// fragment or not. Returns true if it descended into a nested message,
+// false if there was nothing to descend into.
+func (r *rpcPlanningContext) descendIntoResponseField(info *planningInfo, enclosingTypeNode ast.Node, selectionSetRef int) bool {
+	if info.inlineFragmentRef == ast.InvalidRef {
+		if len(info.currentResponseMessage.Fields) == 0 {
+			return false
+		}
+
+		lastIndex := len(info.currentResponseMessage.Fields) - 1
+
+		if info.currentResponseMessage.Fields[lastIndex].Message == nil {
+			info.currentResponseMessage.Fields[lastIndex].Message = r.newMessageFromSelectionSet(enclosingTypeNode, selectionSetRef)
+		}
+
+		info.inlineFragmentRefAncestors.push(info.inlineFragmentRef)
+		info.responseMessageAncestors = append(info.responseMessageAncestors, info.currentResponseMessage)
+		info.currentResponseMessage = info.currentResponseMessage.Fields[lastIndex].Message
+		info.inlineFragmentRef = ast.InvalidRef
+	} else {
+		inlineFragmentName := r.operation.InlineFragmentTypeConditionNameString(info.inlineFragmentRef)
+		fragmentFields := info.currentResponseMessage.FragmentFields[inlineFragmentName]
+		if len(fragmentFields) == 0 {
+			return false
+		}
+		lastIndex := len(fragmentFields) - 1
+
+		if fragmentFields[lastIndex].Message == nil {
+			info.currentResponseMessage.FragmentFields[inlineFragmentName][lastIndex].Message = r.newMessageFromSelectionSet(enclosingTypeNode, selectionSetRef)
+		}
+
+		info.inlineFragmentRefAncestors.push(info.inlineFragmentRef)
+		info.responseMessageAncestors = append(info.responseMessageAncestors, info.currentResponseMessage)
+		info.currentResponseMessage = info.currentResponseMessage.FragmentFields[inlineFragmentName][lastIndex].Message
+		info.inlineFragmentRef = ast.InvalidRef
+	}
+	return true
+}
+
+// ascendFromResponseField pops the response message ancestors and restores the
+// inline fragment ref when leaving a selection set.
+func (r *rpcPlanningContext) ascendFromResponseField(info *planningInfo) {
+	if len(info.responseMessageAncestors) > 0 {
+		info.currentResponseMessage = info.responseMessageAncestors[len(info.responseMessageAncestors)-1]
+		info.responseMessageAncestors = info.responseMessageAncestors[:len(info.responseMessageAncestors)-1]
+		info.inlineFragmentRef = info.inlineFragmentRefAncestors.peek()
+		info.inlineFragmentRefAncestors.pop()
+	}
+}
+
 // newMessageFromSelectionSet creates a new message from the enclosing type node and the selection set reference.
 func (r *rpcPlanningContext) newMessageFromSelectionSet(enclosingTypeNode ast.Node, selectSetRef int) *RPCMessage {
 	message := &RPCMessage{
@@ -769,11 +820,11 @@ func (r *rpcPlanningContext) buildFieldMessage(fieldTypeNode ast.Node, fieldRef
 			return nil, err
 		}
 
-		if message.FieldSelectionSet == nil {
-			message.FieldSelectionSet = make(RPCFieldSelectionSet)
+		if message.FragmentFields == nil {
+			message.FragmentFields = make(RPCFieldSelectionSet)
 		}
 
-		message.FieldSelectionSet.Add(typeName, fields...)
+		message.FragmentFields.Add(typeName, fields...)
 	}
 
 	for _, fieldRef := range fieldRefs {
@@ -1081,7 +1132,7 @@ func (r *rpcPlanningContext) buildFieldResolverTypeMessage(typeName string, reso
 
 	// If the resolved field returns a composite type we need to handle the selection set for the inline fragment.
 	if len(resolverField.fragmentSelections) > 0 {
-		message.FieldSelectionSet = make(RPCFieldSelectionSet, len(resolverField.fragmentSelections))
+		message.FragmentFields = make(RPCFieldSelectionSet, len(resolverField.fragmentSelections))
 
 		for _, fragmentSelection := range resolverField.fragmentSelections {
 			inlineFragmentTypeNode, found := r.definition.NodeByNameStr(fragmentSelection.typeName)
@@ -1094,7 +1145,7 @@ func (r *rpcPlanningContext) buildFieldResolverTypeMessage(typeName string, reso
 				return nil, err
 			}
 
-			message.FieldSelectionSet[fragmentSelection.typeName] = fields
+			message.FragmentFields[fragmentSelection.typeName] = fields
 		}
 	}