Merge branch 'master' into jesse/eng-8566-new-subscription-client

Author: endigma

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

@@ -0,0 +1,74 @@
+package grpcdatasource
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/tidwall/gjson"
+
+	"github.com/wundergraph/astjson"
+)
+
+// entityIndexMap maps positions in the typed gRPC response back to positions
+// in the original representations array. The slice index is the response
+// position; the value is the representation index. It is built per call by
+// recording the position of every representation whose __typename matches
+// the requested entity type.
+type entityIndexMap []int
+
+// newEntityIndexMap builds the index map for a single entity call by collecting
+// the positions of representations whose __typename matches the requested type.
+// A single pass over representations populates the slice.
+func newEntityIndexMap(requestedEntityType string, representations []gjson.Result) entityIndexMap {
+	indexMap := make(entityIndexMap, 0, len(representations))
+	for i, representation := range representations {
+		if representation.Get(typenameFieldName).String() == requestedEntityType {
+			indexMap = append(indexMap, i)
+		}
+	}
+	return indexMap
+}
+
+// getRepresentations gets the representations from the variables.
+// If no representations are found, it returns nil.
+func getRepresentations(variables gjson.Result) []gjson.Result {
+	r := variables.Get("representations")
+	if !r.Exists() {
+		return nil
+	}
+
+	return r.Array()
+}
+
+// validateEntityResponse verifies that the number of entities returned by the
+// subgraph matches the number of representations of the requested type.
+// Callers should subsequently build an entityIndexMap via newEntityIndexMap to
+// merge the response — mergeEntities relies on the invariant that
+// len(response entities) == len(indexMap), which this function establishes.
+func validateEntityResponse(data *astjson.Value, requestedEntityType string, representations []gjson.Result) error {
+	if data == nil {
+		return errors.New("validateEntityResponse: subgraph response data is nil")
+	}
+
+	if requestedEntityType == "" {
+		return errors.New("validateEntityResponse: requested entity type is empty; the entity RPC plan is missing a RequestedEntityType")
+	}
+
+	if len(representations) == 0 {
+		return errors.New("validateEntityResponse: no entity representations provided in the request variables")
+	}
+
+	expected := 0
+	for _, representation := range representations {
+		if representation.Get(typenameFieldName).String() == requestedEntityType {
+			expected++
+		}
+	}
+
+	entities := data.Get(entityPath).GetArray()
+	if len(entities) != expected {
+		return fmt.Errorf("entity type %s received %d entities in the subgraph response, but %d are expected", requestedEntityType, len(entities), expected)
+	}
+
+	return nil
+}

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

@@ -0,0 +1,186 @@
+package grpcdatasource
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tidwall/gjson"
+
+	"github.com/wundergraph/astjson"
+)
+
+func TestNewEntityIndexMap(t *testing.T) {
+	t.Run("returns empty map when no representations match", func(t *testing.T) {
+		reps := getRepresentations(gjson.Parse(`{"representations":[
+			{"__typename":"Storage","id":"1"}
+		]}`))
+		idx := newEntityIndexMap("Product", reps)
+		assert.Equal(t, entityIndexMap{}, idx)
+	})
+
+	t.Run("returns empty map when representations are nil", func(t *testing.T) {
+		idx := newEntityIndexMap("Product", nil)
+		assert.Equal(t, entityIndexMap{}, idx)
+	})
+
+	t.Run("ordered representations [Product, Product, Storage, Storage]", func(t *testing.T) {
+		reps := getRepresentations(gjson.Parse(`{"representations":[
+			{"__typename":"Product","id":"1"},
+			{"__typename":"Product","id":"2"},
+			{"__typename":"Storage","id":"3"},
+			{"__typename":"Storage","id":"4"}
+		]}`))
+
+		productIdx := newEntityIndexMap("Product", reps)
+		assert.Equal(t, entityIndexMap{0, 1}, productIdx)
+
+		storageIdx := newEntityIndexMap("Storage", reps)
+		assert.Equal(t, entityIndexMap{2, 3}, storageIdx)
+	})
+
+	t.Run("unordered representations [Product, Storage, Product, Storage]", func(t *testing.T) {
+		reps := getRepresentations(gjson.Parse(`{"representations":[
+			{"__typename":"Product","id":"1"},
+			{"__typename":"Storage","id":"2"},
+			{"__typename":"Product","id":"3"},
+			{"__typename":"Storage","id":"4"}
+		]}`))
+
+		productIdx := newEntityIndexMap("Product", reps)
+		assert.Equal(t, entityIndexMap{0, 2}, productIdx)
+
+		storageIdx := newEntityIndexMap("Storage", reps)
+		assert.Equal(t, entityIndexMap{1, 3}, storageIdx)
+	})
+
+	t.Run("interleaved representations across three types", func(t *testing.T) {
+		reps := getRepresentations(gjson.Parse(`{"representations":[
+			{"__typename":"Product","id":"1"},
+			{"__typename":"Storage","id":"2"},
+			{"__typename":"Warehouse","id":"3"},
+			{"__typename":"Product","id":"4"},
+			{"__typename":"Warehouse","id":"5"},
+			{"__typename":"Storage","id":"6"}
+		]}`))
+
+		assert.Equal(t, entityIndexMap{0, 3}, newEntityIndexMap("Product", reps))
+		assert.Equal(t, entityIndexMap{1, 5}, newEntityIndexMap("Storage", reps))
+		assert.Equal(t, entityIndexMap{2, 4}, newEntityIndexMap("Warehouse", reps))
+	})
+
+	t.Run("single matching representation", func(t *testing.T) {
+		reps := getRepresentations(gjson.Parse(`{"representations":[
+			{"__typename":"Storage","id":"1"},
+			{"__typename":"Product","id":"2"},
+			{"__typename":"Storage","id":"3"}
+		]}`))
+
+		assert.Equal(t, entityIndexMap{1}, newEntityIndexMap("Product", reps))
+	})
+
+	t.Run("preserves original positions for fully matching list", func(t *testing.T) {
+		reps := getRepresentations(gjson.Parse(`{"representations":[
+			{"__typename":"Product","id":"1"},
+			{"__typename":"Product","id":"2"},
+			{"__typename":"Product","id":"3"}
+		]}`))
+
+		assert.Equal(t, entityIndexMap{0, 1, 2}, newEntityIndexMap("Product", reps))
+	})
+
+	t.Run("interface entity matches by typename string only", func(t *testing.T) {
+		// Interface-entity representations carry the interface name as __typename
+		// (e.g. "Resource"). The index map cares only about the typename string,
+		// not whether it refers to an interface or a concrete type.
+		reps := getRepresentations(gjson.Parse(`{"representations":[
+			{"__typename":"Resource","id":"1"},
+			{"__typename":"Product","id":"2"},
+			{"__typename":"Resource","id":"3"},
+			{"__typename":"Storage","id":"4"},
+			{"__typename":"Resource","id":"5"}
+		]}`))
+
+		assert.Equal(t, entityIndexMap{0, 2, 4}, newEntityIndexMap("Resource", reps))
+		// Concrete types in the same list are independent.
+		assert.Equal(t, entityIndexMap{1}, newEntityIndexMap("Product", reps))
+		assert.Equal(t, entityIndexMap{3}, newEntityIndexMap("Storage", reps))
+	})
+}
+
+func TestGetRepresentations(t *testing.T) {
+	t.Run("returns nil when representations key missing", func(t *testing.T) {
+		vars := gjson.Parse(`{"other":"value"}`)
+		assert.Nil(t, getRepresentations(vars))
+	})
+
+	t.Run("returns empty slice when representations is empty array", func(t *testing.T) {
+		vars := gjson.Parse(`{"representations":[]}`)
+		reps := getRepresentations(vars)
+		assert.NotNil(t, reps)
+		assert.Empty(t, reps)
+	})
+
+	t.Run("returns representations when present", func(t *testing.T) {
+		vars := gjson.Parse(`{"representations":[{"__typename":"Product","id":"1"},{"__typename":"Storage","id":"2"}]}`)
+		reps := getRepresentations(vars)
+		assert.Len(t, reps, 2)
+		assert.Equal(t, "Product", reps[0].Get("__typename").String())
+		assert.Equal(t, "Storage", reps[1].Get("__typename").String())
+	})
+}
+func TestValidateEntityResponse(t *testing.T) {
+	reps := getRepresentations(gjson.Parse(`{"representations":[
+		{"__typename":"Product","id":"1"},
+		{"__typename":"Product","id":"2"}
+	]}`))
+
+	t.Run("returns error when data is nil", func(t *testing.T) {
+		err := validateEntityResponse(nil, "Product", reps)
+		assert.ErrorContains(t, err, "validateEntityResponse: subgraph response data is nil")
+	})
+
+	t.Run("returns error when requested entity type is empty", func(t *testing.T) {
+		data := astjson.MustParse(`{"_entities":[]}`)
+		err := validateEntityResponse(data, "", reps)
+		assert.ErrorContains(t, err, "validateEntityResponse: requested entity type is empty")
+	})
+
+	t.Run("returns error when representations are empty", func(t *testing.T) {
+		data := astjson.MustParse(`{"_entities":[]}`)
+		err := validateEntityResponse(data, "Product", nil)
+		assert.ErrorContains(t, err, "validateEntityResponse: no entity representations provided")
+	})
+
+	t.Run("returns error when entity count mismatches representation count", func(t *testing.T) {
+		data := astjson.MustParse(`{"_entities":[{"__typename":"Product","id":"1"}]}`)
+		err := validateEntityResponse(data, "Product", reps)
+		assert.ErrorContains(t, err, "entity type Product received 1 entities in the subgraph response, but 2 are expected")
+	})
+
+	t.Run("returns nil when entity count matches representation count", func(t *testing.T) {
+		data := astjson.MustParse(`{"_entities":[{"__typename":"Product","id":"1"},{"__typename":"Product","id":"2"}]}`)
+		assert.NoError(t, validateEntityResponse(data, "Product", reps))
+	})
+
+	t.Run("counts only representations of the requested type", func(t *testing.T) {
+		mixedReps := getRepresentations(gjson.Parse(`{"representations":[
+			{"__typename":"Product","id":"1"},
+			{"__typename":"Storage","id":"2"},
+			{"__typename":"Product","id":"3"}
+		]}`))
+		data := astjson.MustParse(`{"_entities":[{"__typename":"Product","id":"1"},{"__typename":"Product","id":"3"}]}`)
+		assert.NoError(t, validateEntityResponse(data, "Product", mixedReps))
+	})
+
+	t.Run("returns error when _entities key is missing", func(t *testing.T) {
+		data := astjson.MustParse(`{}`)
+		err := validateEntityResponse(data, "Product", reps)
+		assert.ErrorContains(t, err, "entity type Product received 0 entities in the subgraph response, but 2 are expected")
+	})
+
+	t.Run("returns error when _entities path is not an array", func(t *testing.T) {
+		data := astjson.MustParse(`{"_entities":"not an array"}`)
+		err := validateEntityResponse(data, "Product", reps)
+		assert.Error(t, err)
+	})
+}

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

@@ -110,6 +110,9 @@ type RPCCall struct {
 	Response RPCMessage
 	// ResponsePath is the path to the response in the JSON response
 	ResponsePath ast.Path
+	// RequestedEntityType is the type of the entity that is being requested
+	// Empty if the call is not an entity lookup.
+	RequestedEntityType string
 }
 
 // RPCMessage represents a gRPC message structure for requests and responses.
@@ -363,6 +366,11 @@ func NewPlanner(subgraphName string, mapping *GRPCMapping, federationConfigs pla
 
 // formatRPCMessage formats an RPCMessage and adds it to the string builder with the specified indentation
 func formatRPCMessage(sb *strings.Builder, message RPCMessage, indent int) {
+	visited := make(map[*RPCMessage]struct{})
+	formatRPCMessageVisited(sb, message, indent, visited)
+}
+
+func formatRPCMessageVisited(sb *strings.Builder, message RPCMessage, indent int, visited map[*RPCMessage]struct{}) {
 	indentStr := strings.Repeat(" ", indent)
 
 	fmt.Fprintf(sb, "%sName: %s\n", indentStr, message.Name)
@@ -375,25 +383,34 @@ func formatRPCMessage(sb *strings.Builder, message RPCMessage, indent int) {
 		fmt.Fprintf(sb, "%s    JSONPath: %s\n", indentStr, field.JSONPath)
 		fmt.Fprintf(sb, "%s    ResolvePath: %s\n", indentStr, field.ResolvePath.String())
 
-		if field.Message != nil {
-			fmt.Fprintf(sb, "%s    Message:\n", indentStr)
-			formatRPCMessage(sb, *field.Message, indent+6)
+		if field.Message == nil {
+			return
+		}
+
+		fmt.Fprintf(sb, "%s    Message:\n", indentStr)
+		if _, seen := visited[field.Message]; seen {
+			fmt.Fprintf(sb, "%s      <recursive: %s>\n", indentStr, field.Message.Name)
+			continue
 		}
+		visited[field.Message] = struct{}{}
+		formatRPCMessageVisited(sb, *field.Message, indent+6, visited)
 	}
 }
 
 type rpcPlanningContext struct {
-	operation  *ast.Document
-	definition *ast.Document
-	mapping    *GRPCMapping
+	operation         *ast.Document
+	definition        *ast.Document
+	mapping           *GRPCMapping
+	visitedInputTypes map[string]*RPCMessage
 }
 
 // newRPCPlanningContext creates a new RPCPlanningContext.
 func newRPCPlanningContext(operation *ast.Document, definition *ast.Document, mapping *GRPCMapping) *rpcPlanningContext {
 	return &rpcPlanningContext{
-		operation:  operation,
-		definition: definition,
-		mapping:    mapping,
+		operation:         operation,
+		definition:        definition,
+		mapping:           mapping,
+		visitedInputTypes: make(map[string]*RPCMessage, len(definition.InputObjectTypeDefinitions)),
 	}
 }
 
@@ -684,11 +701,26 @@ func (r *rpcPlanningContext) buildMessageFromInputObjectType(node *ast.Node) (*R
 		return nil, fmt.Errorf("unable to build message from input object type definition - incorrect type: %s", node.Kind)
 	}
 
+	typeName := node.NameString(r.definition)
+
+	// If we've already started building this type, return the in-progress message
+	// pointer to break the recursion cycle. The message's fields are populated by
+	// the caller that first entered this type, so the pointer will be complete once
+	// the top-level call returns.
+	if existing, ok := r.visitedInputTypes[typeName]; ok {
+		return existing, nil
+	}
+
 	inputObjectDefinition := r.definition.InputObjectTypeDefinitions[node.Ref]
 	message := &RPCMessage{
-		Name:   node.NameString(r.definition),
+		Name:   typeName,
 		Fields: make(RPCFields, 0, len(inputObjectDefinition.InputFieldsDefinition.Refs)),
 	}
+
+	// Register the message before recursing into fields so that recursive
+	// references resolve to this same pointer.
+	r.visitedInputTypes[typeName] = message
+
 	for _, inputFieldRef := range inputObjectDefinition.InputFieldsDefinition.Refs {
 		field, err := r.buildMessageFieldFromInputValueDefinition(inputFieldRef, node)
 		if err != nil {