agent_patch
runs/2026-02-28__13-38-10__gpt-5-4/flux-pr-1262/flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/agent-logs/agent.patch
47786 bytes
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/compiler.go b/app/v2/pkg/engine/datasource/grpc_datasource/compiler.go
index ee49e81..ea13875 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/compiler.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/compiler.go
@@ -3,9 +3,12 @@ package grpcdatasource
import (
"context"
"fmt"
+ "sort"
+ "strings"
"github.com/bufbuild/protocompile"
"github.com/tidwall/gjson"
+ "github.com/tidwall/sjson"
"github.com/wundergraph/graphql-go-tools/v2/pkg/ast"
"github.com/wundergraph/graphql-go-tools/v2/pkg/operationreport"
protoref "google.golang.org/protobuf/reflect/protoreflect"
@@ -330,14 +333,39 @@ type Invocation struct {
Input *dynamicpb.Message
Output *dynamicpb.Message
Call *RPCCall
+ Entity *EntityInvocation
+}
+
+type EntityInvocation struct {
+ Positions []int
}
// Compile processes an RPCExecutionPlan and builds protobuf messages from JSON data
// based on the compiled schema.
func (p *RPCCompiler) Compile(executionPlan *RPCExecutionPlan, inputData gjson.Result) ([]Invocation, error) {
invocations := make([]Invocation, 0, len(executionPlan.Calls))
+ entityAssignments := assignEntityRepresentations(executionPlan.Calls, inputData)
+
+ for callIndex, call := range executionPlan.Calls {
+ callInputData := inputData
+ var entityInvocation *EntityInvocation
+ if call.EntityLookup != nil {
+ positions := entityAssignments[callIndex]
+ if len(positions) == 0 {
+ continue
+ }
+
+ filteredInputData, err := filterRepresentationsInput(inputData, call.EntityLookup.RepresentationsPath, positions)
+ if err != nil {
+ return nil, err
+ }
+
+ callInputData = filteredInputData
+ entityInvocation = &EntityInvocation{
+ Positions: positions,
+ }
+ }
- for _, call := range executionPlan.Calls {
inputMessage, ok := p.doc.MessageByName(call.Request.Name)
if !ok {
return nil, fmt.Errorf("input message %s not found in document", call.Request.Name)
@@ -348,7 +376,7 @@ func (p *RPCCompiler) Compile(executionPlan *RPCExecutionPlan, inputData gjson.R
return nil, fmt.Errorf("output message %s not found in document", call.Response.Name)
}
- request := p.buildProtoMessage(inputMessage, &call.Request, inputData)
+ request := p.buildProtoMessage(inputMessage, &call.Request, callInputData)
response := p.newEmptyMessage(outputMessage)
if p.report.HasErrors() {
@@ -366,12 +394,86 @@ func (p *RPCCompiler) Compile(executionPlan *RPCExecutionPlan, inputData gjson.R
Input: request,
Output: response,
Call: &call,
+ Entity: entityInvocation,
})
}
return invocations, nil
}
+func assignEntityRepresentations(calls []RPCCall, inputData gjson.Result) map[int][]int {
+ assignments := make(map[int][]int)
+ representations := inputData.Get("representations").Array()
+ if len(representations) == 0 {
+ return assignments
+ }
+
+ entityCallIndexes := make([]int, 0, len(calls))
+ for i, call := range calls {
+ if call.EntityLookup != nil {
+ entityCallIndexes = append(entityCallIndexes, i)
+ }
+ }
+
+ sort.SliceStable(entityCallIndexes, func(i, j int) bool {
+ left := entityCallIndexes[i]
+ right := entityCallIndexes[j]
+ leftLen := len(calls[left].EntityLookup.KeyFields)
+ rightLen := len(calls[right].EntityLookup.KeyFields)
+ if leftLen == rightLen {
+ return left < right
+ }
+ return leftLen > rightLen
+ })
+
+ for position, representation := range representations {
+ for _, callIndex := range entityCallIndexes {
+ if entityLookupMatchesRepresentation(calls[callIndex].EntityLookup, representation) {
+ assignments[callIndex] = append(assignments[callIndex], position)
+ break
+ }
+ }
+ }
+
+ return assignments
+}
+
+func entityLookupMatchesRepresentation(entityLookup *EntityLookup, representation gjson.Result) bool {
+ if entityLookup == nil {
+ return false
+ }
+
+ if representation.Get("__typename").String() != entityLookup.TypeName {
+ return false
+ }
+
+ for _, key := range entityLookup.KeyFields {
+ if !representation.Get(key).Exists() {
+ return false
+ }
+ }
+
+ return true
+}
+
+func filterRepresentationsInput(inputData gjson.Result, representationsPath string, positions []int) (gjson.Result, error) {
+ representations := inputData.Get(representationsPath).Array()
+ rawRepresentations := make([]string, 0, len(positions))
+ for _, position := range positions {
+ if position < 0 || position >= len(representations) {
+ continue
+ }
+ rawRepresentations = append(rawRepresentations, representations[position].Raw)
+ }
+
+ filteredInput, err := sjson.SetRaw(inputData.Raw, representationsPath, "["+strings.Join(rawRepresentations, ",")+"]")
+ if err != nil {
+ return gjson.Result{}, fmt.Errorf("failed to filter entity representations: %w", err)
+ }
+
+ return gjson.Parse(filteredInput), nil
+}
+
func (p *RPCCompiler) resolveServiceName(methodName string) (string, bool) {
for _, service := range p.doc.Services {
for _, methodRef := range service.MethodsRefs {
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/compiler_test.go b/app/v2/pkg/engine/datasource/grpc_datasource/compiler_test.go
index 4859896..988fcb5 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/compiler_test.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/compiler_test.go
@@ -293,6 +293,89 @@ func TestBuildProtoMessage(t *testing.T) {
}
+func TestCompileEntityInvocationsFiltersRepresentationsByTypeAndKey(t *testing.T) {
+ compiler, err := NewProtoCompiler(validProto, nil)
+ require.NoError(t, err)
+
+ plan := &RPCExecutionPlan{
+ Calls: []RPCCall{
+ {
+ MethodName: "LookupProductById",
+ Request: RPCMessage{
+ Name: "LookupProductByIdRequest",
+ Fields: []RPCField{{
+ Name: "inputs",
+ TypeName: string(DataTypeMessage),
+ Repeated: true,
+ JSONPath: "representations",
+ Message: &RPCMessage{
+ Name: "LookupProductByIdInput",
+ Fields: []RPCField{{
+ Name: "key",
+ TypeName: string(DataTypeMessage),
+ Message: &RPCMessage{
+ Name: "ProductByIdKey",
+ Fields: []RPCField{{
+ Name: "id",
+ TypeName: string(DataTypeString),
+ JSONPath: "id",
+ }},
+ },
+ }},
+ },
+ }},
+ },
+ Response: RPCMessage{Name: "LookupProductByIdResponse"},
+ EntityLookup: &EntityLookup{
+ TypeName: "Product",
+ KeyFields: []string{"id"},
+ RepresentationsPath: "representations",
+ ResultPath: "_entities",
+ },
+ },
+ {
+ MethodName: "LookupProductByName",
+ Request: RPCMessage{
+ Name: "LookupProductByNameRequest",
+ Fields: []RPCField{{
+ Name: "inputs",
+ TypeName: string(DataTypeMessage),
+ Repeated: true,
+ JSONPath: "representations",
+ Message: &RPCMessage{
+ Name: "LookupProductByNameInput",
+ Fields: []RPCField{{
+ Name: "name",
+ TypeName: string(DataTypeString),
+ JSONPath: "name",
+ }},
+ },
+ }},
+ },
+ Response: RPCMessage{Name: "LookupProductByNameResponse"},
+ EntityLookup: &EntityLookup{
+ TypeName: "Product",
+ KeyFields: []string{"name"},
+ RepresentationsPath: "representations",
+ ResultPath: "_entities",
+ },
+ },
+ },
+ }
+
+ invocations, err := compiler.Compile(plan, gjson.ParseBytes([]byte(`{"representations":[{"__typename":"Product","id":"p1"},{"__typename":"Product","name":"Chair"},{"__typename":"Storage","id":"s1"}]}`)))
+ require.NoError(t, err)
+ require.Len(t, invocations, 2)
+ require.Equal(t, []int{0}, invocations[0].Entity.Positions)
+ require.Equal(t, []int{1}, invocations[1].Entity.Positions)
+
+ inputByID := invocations[0].Input.ProtoReflect()
+ require.Equal(t, 1, inputByID.Get(inputByID.Descriptor().Fields().ByName("inputs")).List().Len())
+
+ inputByName := invocations[1].Input.ProtoReflect()
+ require.Equal(t, 1, inputByName.Get(inputByName.Descriptor().Fields().ByName("inputs")).List().Len())
+}
+
func TestCompileNestedLists(t *testing.T) {
compiler, err := NewProtoCompiler(grpctest.MustProtoSchema(t), testMapping())
require.NoError(t, err)
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/configuration.go b/app/v2/pkg/engine/datasource/grpc_datasource/configuration.go
index 6ee1940..b016fab 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/configuration.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/configuration.go
@@ -16,8 +16,9 @@ type GRPCMapping struct {
MutationRPCs RPCConfigMap
// SubscriptionRPCs maps GraphQL subscription fields to the corresponding gRPC RPC configurations
SubscriptionRPCs RPCConfigMap
- // EntityRPCs defines how GraphQL types are resolved as entities using specific RPCs
- EntityRPCs map[string]EntityRPCConfig
+ // EntityRPCs defines how GraphQL types are resolved as entities using specific RPCs.
+ // A type can expose multiple key directives, so each type may map to multiple RPC configs.
+ EntityRPCs map[string][]EntityRPCConfig
// Fields defines the field mappings between GraphQL types and gRPC messages
Fields map[string]FieldMap
// EnumValues defines the enum values for each enum type
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan.go b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan.go
index ad41e4b..6075d1d 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan.go
@@ -68,6 +68,15 @@ type RPCCall struct {
Request RPCMessage
// Response contains the message structure for the gRPC response
Response RPCMessage
+ // EntityLookup holds federation-specific metadata for _entities lookups.
+ EntityLookup *EntityLookup
+}
+
+type EntityLookup struct {
+ TypeName string
+ KeyFields []string
+ RepresentationsPath string
+ ResultPath string
}
// RPCMessage represents a gRPC message structure for requests and responses.
@@ -252,6 +261,55 @@ func (r RPCFields) Exists(name, alias string) bool {
return false
}
+func cloneRPCMessage(message RPCMessage) RPCMessage {
+ cloned := message
+ cloned.Fields = cloneRPCFields(message.Fields)
+ cloned.FieldSelectionSet = cloneRPCFieldSelectionSet(message.FieldSelectionSet)
+ cloned.MemberTypes = append([]string(nil), message.MemberTypes...)
+ if message.Message != nil {
+ msg := cloneRPCMessage(*message.Message)
+ cloned.Message = &msg
+ }
+ return cloned
+}
+
+func cloneRPCFields(fields RPCFields) RPCFields {
+ if len(fields) == 0 {
+ return nil
+ }
+
+ cloned := make(RPCFields, len(fields))
+ for i, field := range fields {
+ cloned[i] = field
+ if field.Message != nil {
+ msg := cloneRPCMessage(*field.Message)
+ cloned[i].Message = &msg
+ }
+ if field.ListMetadata != nil {
+ levelInfo := append([]LevelInfo(nil), field.ListMetadata.LevelInfo...)
+ cloned[i].ListMetadata = &ListMetadata{
+ NestingLevel: field.ListMetadata.NestingLevel,
+ LevelInfo: levelInfo,
+ }
+ }
+ }
+
+ return cloned
+}
+
+func cloneRPCFieldSelectionSet(fieldSelectionSet RPCFieldSelectionSet) RPCFieldSelectionSet {
+ if len(fieldSelectionSet) == 0 {
+ return nil
+ }
+
+ cloned := make(RPCFieldSelectionSet, len(fieldSelectionSet))
+ for key, fields := range fieldSelectionSet {
+ cloned[key] = cloneRPCFields(fields)
+ }
+
+ return cloned
+}
+
func (r *RPCExecutionPlan) String() string {
var result strings.Builder
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_federation_test.go b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_federation_test.go
index 5756d9d..ed46e82 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_federation_test.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_federation_test.go
@@ -21,15 +21,15 @@ func TestEntityLookup(t *testing.T) {
query: `query EntityLookup($representations: [_Any!]!) { _entities(representations: $representations) { ... on Product { __typename id name price } } }`,
mapping: &GRPCMapping{
Service: "Products",
- EntityRPCs: map[string]EntityRPCConfig{
- "Product": {
+ EntityRPCs: map[string][]EntityRPCConfig{
+ "Product": {{
Key: "id",
RPCConfig: RPCConfig{
RPC: "LookupProductById",
Request: "LookupProductByIdRequest",
Response: "LookupProductByIdResponse",
},
- },
+ }},
},
},
expectedPlan: &RPCExecutionPlan{
@@ -37,6 +37,12 @@ func TestEntityLookup(t *testing.T) {
{
ServiceName: "Products",
MethodName: "LookupProductById",
+ EntityLookup: &EntityLookup{
+ TypeName: "Product",
+ KeyFields: []string{"id"},
+ RepresentationsPath: "representations",
+ ResultPath: "_entities",
+ },
// Define the structure of the request message
Request: RPCMessage{
Name: "LookupProductByIdRequest",
@@ -101,195 +107,106 @@ func TestEntityLookup(t *testing.T) {
},
},
},
-
- // TODO implement multiple entity lookup types
- // {
- // name: "Should create an execution plan for an entity lookup multiple types",
- // query: `
- // query EntityLookup($representations: [_Any!]!) {
- // _entities(representations: $representations) {
- // ... on Product {
- // id
- // name
- // price
- // }
- // ... on Storage {
- // id
- // name
- // location
- // }
- // }
- // }
- // `,
- // expectedPlan: &RPCExecutionPlan{
- // Groups: []RPCCallGroup{
- // {
- // Calls: []RPCCall{
- // {
- // ServiceName: "Products",
- // MethodName: "LookupProductById",
- // // Define the structure of the request message
- // Request: RPCMessage{
- // Name: "LookupProductByIdRequest",
- // Fields: []RPCField{
- // {
- // Name: "inputs",
- // TypeName: string(DataTypeMessage),
- // Repeated: true,
- // JSONPath: "representations", // Path to extract data from GraphQL variables
- //
-
- // Message: &RPCMessage{
- // Name: "LookupProductByIdInput",
- // Fields: []RPCField{
- // {
- // Name: "key",
- // TypeName: string(DataTypeMessage),
- //
-
- // Message: &RPCMessage{
- // Name: "ProductByIdKey",
- // Fields: []RPCField{
- // {
- // Name: "id",
- // TypeName: string(DataTypeString),
- // JSONPath: "id", // Extract 'id' from each representation
- //
-
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // // Define the structure of the response message
- // Response: RPCMessage{
- // Name: "LookupProductByIdResponse",
- // Fields: []RPCField{
- // {
- // Name: "results",
- // TypeName: string(DataTypeMessage),
- // Repeated: true,
- //
-
- // JSONPath: "results",
- // Message: &RPCMessage{
- // Name: "LookupProductByIdResult",
- // Fields: []RPCField{
- // {
- // Name: "product",
- // TypeName: string(DataTypeMessage),
- //
-
- // Message: &RPCMessage{
- // Name: "Product",
- // Fields: []RPCField{
- // {
- // Name: "id",
- // TypeName: string(DataTypeString),
- // JSONPath: "id",
- // },
- // {
- // Name: "name",
- // TypeName: string(DataTypeString),
- // JSONPath: "name",
- // },
- // {
- // Name: "price",
- // TypeName: string(DataTypeFloat),
- // JSONPath: "price",
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // {
- // ServiceName: "Products",
- // MethodName: "LookupStorageById",
- // Request: RPCMessage{
- // Name: "LookupStorageByIdRequest",
- // Fields: []RPCField{
- // {
- // Name: "inputs",
- // TypeName: string(DataTypeMessage),
- // Repeated: true,
- // JSONPath: "representations",
- // Message: &RPCMessage{
- // Name: "LookupStorageByIdInput",
- // Fields: []RPCField{
- // {
- // Name: "key",
- // TypeName: string(DataTypeMessage),
- // Message: &RPCMessage{
- // Name: "StorageByIdKey",
- // Fields: []RPCField{
- // {
- // Name: "id",
- // TypeName: string(DataTypeString),
- // JSONPath: "id",
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // Response: RPCMessage{
- // Name: "LookupStorageByIdResponse",
- // Fields: []RPCField{
- // {
- // Name: "results",
- // TypeName: string(DataTypeMessage),
- // Repeated: true,
- // JSONPath: "results",
- // Message: &RPCMessage{
- // Name: "LookupStorageByIdResult",
- // Fields: []RPCField{
- // {
- // Name: "storage",
- // TypeName: string(DataTypeMessage),
- // Message: &RPCMessage{
- // Name: "Storage",
- // Fields: []RPCField{
- // {
- // Name: "id",
- // TypeName: string(DataTypeString),
- // JSONPath: "id",
- // },
- // {
- // Name: "name",
- // TypeName: string(DataTypeString),
- // JSONPath: "name",
- // },
- // {
- // Name: "location",
- // TypeName: string(DataTypeString),
- // JSONPath: "location",
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
- // },
+ {
+ name: "Should create an execution plan for an entity lookup with multiple types",
+ query: `query EntityLookup($representations: [_Any!]!) { _entities(representations: $representations) { ... on Product { __typename id name price } ... on Storage { __typename id name location } } }`,
+ mapping: testMapping(),
+ expectedPlan: &RPCExecutionPlan{
+ Calls: []RPCCall{
+ {
+ ServiceName: "Products",
+ MethodName: "LookupProductById",
+ EntityLookup: &EntityLookup{
+ TypeName: "Product",
+ KeyFields: []string{"id"},
+ RepresentationsPath: "representations",
+ ResultPath: "_entities",
+ },
+ Request: RPCMessage{
+ Name: "LookupProductByIdRequest",
+ Fields: []RPCField{{
+ Name: "keys",
+ TypeName: string(DataTypeMessage),
+ Repeated: true,
+ JSONPath: "representations",
+ Message: &RPCMessage{
+ Name: "LookupProductByIdKey",
+ Fields: []RPCField{{
+ Name: "id",
+ TypeName: string(DataTypeString),
+ JSONPath: "id",
+ }},
+ },
+ }},
+ },
+ Response: RPCMessage{
+ Name: "LookupProductByIdResponse",
+ Fields: []RPCField{{
+ Name: "result",
+ TypeName: string(DataTypeMessage),
+ Repeated: true,
+ JSONPath: "_entities",
+ Message: &RPCMessage{
+ Name: "Product",
+ Fields: []RPCField{
+ {Name: "__typename", TypeName: string(DataTypeString), JSONPath: "__typename", StaticValue: "Product"},
+ {Name: "id", TypeName: string(DataTypeString), JSONPath: "id"},
+ {Name: "name", TypeName: string(DataTypeString), JSONPath: "name"},
+ {Name: "price", TypeName: string(DataTypeDouble), JSONPath: "price"},
+ },
+ },
+ }},
+ },
+ },
+ {
+ CallID: 1,
+ ServiceName: "Products",
+ MethodName: "LookupStorageById",
+ EntityLookup: &EntityLookup{
+ TypeName: "Storage",
+ KeyFields: []string{"id"},
+ RepresentationsPath: "representations",
+ ResultPath: "_entities",
+ },
+ Request: RPCMessage{
+ Name: "LookupStorageByIdRequest",
+ Fields: []RPCField{{
+ Name: "keys",
+ TypeName: string(DataTypeMessage),
+ Repeated: true,
+ JSONPath: "representations",
+ Message: &RPCMessage{
+ Name: "LookupStorageByIdKey",
+ Fields: []RPCField{{
+ Name: "id",
+ TypeName: string(DataTypeString),
+ JSONPath: "id",
+ }},
+ },
+ }},
+ },
+ Response: RPCMessage{
+ Name: "LookupStorageByIdResponse",
+ Fields: []RPCField{{
+ Name: "result",
+ TypeName: string(DataTypeMessage),
+ Repeated: true,
+ JSONPath: "_entities",
+ Message: &RPCMessage{
+ Name: "Storage",
+ Fields: []RPCField{
+ {Name: "__typename", TypeName: string(DataTypeString), JSONPath: "__typename", StaticValue: "Storage"},
+ {Name: "id", TypeName: string(DataTypeString), JSONPath: "id"},
+ {Name: "name", TypeName: string(DataTypeString), JSONPath: "name"},
+ {Name: "location", TypeName: string(DataTypeString), JSONPath: "location"},
+ },
+ },
+ }},
+ },
+ },
+ },
+ },
+ },
}
for _, tt := range tests {
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_visitor.go b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_visitor.go
index edf2dd8..2cf733b 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_visitor.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_visitor.go
@@ -2,6 +2,7 @@ package grpcdatasource
import (
"fmt"
+ "slices"
"strings"
"github.com/wundergraph/graphql-go-tools/v2/pkg/ast"
@@ -18,12 +19,16 @@ type keyField struct {
type entityInfo struct {
name string
- keyFields []keyField
- keyTypeName string
entityRootFieldRef int
entityInlineFragmentRef int
}
+type entityCallConfig struct {
+ keyFields []keyField
+ keyTypeName string
+ rpcConfig RPCConfig
+}
+
type planningInfo struct {
entityInfo entityInfo
// resolvers []string
@@ -56,6 +61,8 @@ type rpcPlanVisitor struct {
operationFieldRefs []int
currentCall *RPCCall
currentCallID int
+ activeEntityCall *RPCCall
+ activeEntityConfigs []entityCallConfig
}
type rpcPlanVisitorConfig struct {
@@ -73,6 +80,12 @@ func newRPCPlanVisitor(walker *astvisitor.Walker, config rpcPlanVisitorConfig) *
mapping: config.mapping,
operationDefinitionRef: -1,
operationFieldRef: -1,
+ planInfo: planningInfo{
+ entityInfo: entityInfo{
+ entityRootFieldRef: -1,
+ entityInlineFragmentRef: -1,
+ },
+ },
}
walker.RegisterEnterDocumentVisitor(visitor)
@@ -107,8 +120,10 @@ func (r *rpcPlanVisitor) EnterOperationDefinition(ref int) {
selectionSetRef := r.operation.OperationDefinitions[ref].SelectionSet
r.operationFieldRefs = r.operation.SelectionSetFieldSelections(selectionSetRef)
- r.plan.Calls = make([]RPCCall, len(r.operationFieldRefs))
+ r.plan.Calls = make([]RPCCall, 0, len(r.operationFieldRefs))
r.planInfo.operationType = r.operation.OperationDefinitions[ref].OperationType
+ r.planInfo.entityInfo.entityRootFieldRef = -1
+ r.planInfo.entityInfo.entityInlineFragmentRef = -1
}
// EnterArgument implements astvisitor.EnterArgumentVisitor.
@@ -145,6 +160,10 @@ func (r *rpcPlanVisitor) EnterSelectionSet(ref int) {
return
}
+ if r.planInfo.currentResponseMessage == nil {
+ return
+ }
+
if len(r.planInfo.currentResponseMessage.Fields) == 0 || len(r.planInfo.currentResponseMessage.Fields) <= r.planInfo.currentResponseFieldIndex {
return
}
@@ -235,7 +254,22 @@ func (r *rpcPlanVisitor) EnterInlineFragment(ref int) {
entityInfo := &r.planInfo.entityInfo
if entityInfo.entityRootFieldRef != -1 && entityInfo.entityInlineFragmentRef == -1 {
entityInfo.entityInlineFragmentRef = ref
- r.resolveEntityInformation(ref)
+ configs, err := r.resolveEntityInformation(ref)
+ if err != nil {
+ r.walker.StopWithInternalErr(err)
+ return
+ }
+
+ if len(configs) == 0 {
+ return
+ }
+
+ r.activeEntityConfigs = configs
+ r.activeEntityCall = &RPCCall{
+ ServiceName: r.resolveServiceName(),
+ }
+ r.planInfo.currentRequestMessage = &r.activeEntityCall.Request
+ r.planInfo.currentResponseMessage = &r.activeEntityCall.Response
r.scaffoldEntityLookup()
return
@@ -245,6 +279,7 @@ func (r *rpcPlanVisitor) EnterInlineFragment(ref int) {
// LeaveInlineFragment implements astvisitor.InlineFragmentVisitor.
func (r *rpcPlanVisitor) LeaveInlineFragment(ref int) {
if ref == r.planInfo.entityInfo.entityInlineFragmentRef {
+ r.finalizeEntityLookup()
r.planInfo.entityInfo.entityInlineFragmentRef = -1
}
}
@@ -269,6 +304,10 @@ func (r *rpcPlanVisitor) IsInlineFragmentField() (int, bool) {
func (r *rpcPlanVisitor) handleRootField(ref int) error {
r.operationFieldRef = ref
r.planInfo.operationFieldName = r.operation.FieldNameString(ref)
+ r.planInfo.methodName = ""
+ r.planInfo.requestMessageAncestors = nil
+ r.planInfo.responseMessageAncestors = nil
+ r.planInfo.responseFieldIndexAncestors = nil
r.currentCall = &RPCCall{
CallID: r.currentCallID,
@@ -289,13 +328,6 @@ func (r *rpcPlanVisitor) handleRootField(ref int) error {
// EnterField implements astvisitor.EnterFieldVisitor.
func (r *rpcPlanVisitor) EnterField(ref int) {
fieldName := r.operation.FieldNameString(ref)
- if r.IsRootField() {
- if err := r.handleRootField(ref); err != nil {
- r.walker.StopWithInternalErr(err)
- return
- }
- }
-
if fieldName == "_entities" {
// _entities is a special field that is used to look up entities
// Entity lookups are handled differently as we use special types for
@@ -310,8 +342,19 @@ func (r *rpcPlanVisitor) EnterField(ref int) {
return
}
+ if r.IsRootField() {
+ if err := r.handleRootField(ref); err != nil {
+ r.walker.StopWithInternalErr(err)
+ return
+ }
+ }
+
// prevent duplicate fields
fieldAlias := r.operation.FieldAliasString(ref)
+ if r.planInfo.currentResponseMessage == nil {
+ return
+ }
+
if r.planInfo.currentResponseMessage.Fields.Exists(fieldName, fieldAlias) {
return
}
@@ -380,8 +423,17 @@ func (r *rpcPlanVisitor) buildField(fd int, fieldName, fieldAlias string) RPCFie
// LeaveField implements astvisitor.FieldVisitor.
func (r *rpcPlanVisitor) LeaveField(ref int) {
- if ref == r.planInfo.entityInfo.entityRootFieldRef {
+ if r.planInfo.entityInfo.entityRootFieldRef != -1 && ref == r.planInfo.entityInfo.entityRootFieldRef {
r.planInfo.entityInfo.entityRootFieldRef = -1
+ r.planInfo.isEntityLookup = false
+ r.planInfo.currentRequestMessage = nil
+ r.planInfo.currentResponseMessage = nil
+ r.activeEntityCall = nil
+ r.activeEntityConfigs = nil
+ r.planInfo.currentResponseFieldIndex = 0
+ r.planInfo.responseMessageAncestors = nil
+ r.planInfo.responseFieldIndexAncestors = nil
+ return
}
// If we are not in the operation field, we can increment the response field index.
@@ -398,7 +450,7 @@ func (r *rpcPlanVisitor) LeaveField(ref int) {
r.currentCall.Response.Name = methodName + "Response"
}
- r.plan.Calls[r.currentCallID] = *r.currentCall
+ r.plan.Calls = append(r.plan.Calls, *r.currentCall)
r.currentCall = &RPCCall{}
r.currentCallID++
@@ -407,6 +459,8 @@ func (r *rpcPlanVisitor) LeaveField(ref int) {
}
r.planInfo.currentResponseFieldIndex = 0
+ r.planInfo.currentRequestMessage = nil
+ r.planInfo.currentResponseMessage = nil
}
// newMessageFromSelectionSet creates a new message from a selection set.
@@ -584,95 +638,157 @@ func (r *rpcPlanVisitor) createListMetadata(typeRef int) *ListMetadata {
return md
}
-func (r *rpcPlanVisitor) resolveEntityInformation(inlineFragmentRef int) {
- // TODO support multiple entities in a single query
- if !r.planInfo.isEntityLookup || r.planInfo.entityInfo.name != "" {
- return
+func (r *rpcPlanVisitor) resolveEntityInformation(inlineFragmentRef int) ([]entityCallConfig, error) {
+ if !r.planInfo.isEntityLookup {
+ return nil, nil
}
fragmentName := r.operation.InlineFragmentTypeConditionNameString(inlineFragmentRef)
node, found := r.definition.NodeByNameStr(fragmentName)
if !found {
- return
+ return nil, nil
}
- // Only process object type definitions
- // TODO: handle interfaces
if node.Kind != ast.NodeKindObjectTypeDefinition {
- return
+ return nil, nil
}
- // An entity must at least have a key directive
+ r.planInfo.entityInfo.name = fragmentName
+
def := r.definition.ObjectTypeDefinitions[node.Ref]
if !def.HasDirectives {
- return
+ return nil, nil
}
- // TODO: We currently only support one key directive per entity
- // We need to get the used key from the graphql datasource.
+ schemaKeys := make([]entityCallConfig, 0, len(def.Directives.Refs))
for _, directiveRef := range def.Directives.Refs {
if r.definition.DirectiveNameString(directiveRef) != federationKeyDirectiveName {
continue
}
- r.planInfo.entityInfo.name = fragmentName
-
directive := r.definition.Directives[directiveRef]
for _, argRef := range directive.Arguments.Refs {
if r.definition.ArgumentNameString(argRef) != "fields" {
continue
}
- argument := r.definition.Arguments[argRef]
- keyFieldName := r.definition.ValueContentString(argument.Value)
-
- fieldDef, ok := r.definition.NodeFieldDefinitionByName(node, ast.ByteSlice(keyFieldName))
- if !ok {
- r.walker.Report.AddExternalError(operationreport.ExternalError{
- Message: fmt.Sprintf("Field %s not found in definition", keyFieldName),
- })
- return
- }
- fdt := r.definition.FieldDefinitionType(fieldDef)
- ft := r.definition.Types[fdt]
+ keyValue := r.definition.ValueContentString(r.definition.Arguments[argRef].Value)
+ fields, err := r.resolveEntityKeyFields(node, keyValue)
+ if err != nil {
+ return nil, err
+ }
- r.planInfo.entityInfo.keyFields =
- append(r.planInfo.entityInfo.keyFields, keyField{
- fieldName: keyFieldName,
- fieldType: r.toDataType(&ft).String(),
- })
+ schemaKeys = append(schemaKeys, entityCallConfig{
+ keyFields: fields,
+ keyTypeName: r.entityKeyTypeName(fragmentName, fields),
+ })
}
+ }
- break
+ if len(schemaKeys) == 0 {
+ return nil, nil
}
- keyFields := make([]string, 0, len(r.planInfo.entityInfo.keyFields))
- for _, key := range r.planInfo.entityInfo.keyFields {
- keyFields = append(keyFields, key.fieldName)
+ configuredRPCs := r.mapping.EntityRPCs[fragmentName]
+ if len(configuredRPCs) == 0 {
+ return schemaKeys, nil
}
- if ei, exists := r.mapping.EntityRPCs[r.planInfo.entityInfo.name]; exists {
- r.currentCall.Request.Name = ei.RPCConfig.Request
- r.currentCall.Response.Name = ei.RPCConfig.Response
- r.planInfo.methodName = ei.RPCConfig.RPC
+ configs := make([]entityCallConfig, 0, len(configuredRPCs))
+ for _, rpcConfig := range configuredRPCs {
+ fields, err := r.resolveEntityKeyFields(node, rpcConfig.Key)
+ if err != nil {
+ return nil, err
+ }
+
+ configs = append(configs, entityCallConfig{
+ keyFields: fields,
+ keyTypeName: r.entityKeyTypeName(fragmentName, fields),
+ rpcConfig: rpcConfig.RPCConfig,
+ })
}
- r.planInfo.entityInfo.keyTypeName = r.planInfo.entityInfo.name + "By" + strings.Join(titleSlice(keyFields), "And")
+ return configs, nil
}
// scaffoldEntityLookup creates the entity lookup call structure
// by creating the key field message and adding it to the current request message.
// It also adds the results message to the current response message.
func (r *rpcPlanVisitor) scaffoldEntityLookup() {
- if !r.planInfo.isEntityLookup {
+ if !r.planInfo.isEntityLookup || r.planInfo.currentResponseMessage == nil {
return
}
- entityInfo := &r.planInfo.entityInfo
+ // The proto response message has a field `result` which is a list of entities.
+ // As this is a special case we directly map it to _entities.
+ r.planInfo.currentResponseMessage.Fields = []RPCField{
+ {
+ Name: "result",
+ TypeName: DataTypeMessage.String(),
+ JSONPath: "_entities",
+ Repeated: true,
+ },
+ }
+}
+
+func (r *rpcPlanVisitor) finalizeEntityLookup() {
+ if r.activeEntityCall == nil || len(r.activeEntityConfigs) == 0 {
+ r.activeEntityCall = nil
+ r.activeEntityConfigs = nil
+ r.planInfo.currentRequestMessage = nil
+ r.planInfo.currentResponseMessage = nil
+ r.planInfo.currentResponseFieldIndex = 0
+ r.planInfo.responseMessageAncestors = nil
+ r.planInfo.responseFieldIndexAncestors = nil
+ return
+ }
+
+ responseTemplate := cloneRPCMessage(r.activeEntityCall.Response)
+ for _, cfg := range r.activeEntityConfigs {
+ call := RPCCall{
+ CallID: r.currentCallID,
+ ServiceName: r.resolveServiceName(),
+ Response: cloneRPCMessage(responseTemplate),
+ EntityLookup: &EntityLookup{
+ TypeName: r.planInfo.entityInfo.name,
+ KeyFields: keyFieldNames(cfg.keyFields),
+ RepresentationsPath: "representations",
+ ResultPath: "_entities",
+ },
+ }
+
+ requestName := ""
+ if cfg.rpcConfig.RPC != "" {
+ call.MethodName = cfg.rpcConfig.RPC
+ requestName = cfg.rpcConfig.Request
+ call.Response.Name = cfg.rpcConfig.Response
+ } else {
+ methodName := "Lookup" + cfg.keyTypeName
+ call.MethodName = methodName
+ requestName = methodName + "Request"
+ call.Response.Name = methodName + "Response"
+ }
+
+ call.Request = r.entityLookupRequest(requestName, call.MethodName, cfg)
+ r.plan.Calls = append(r.plan.Calls, call)
+ r.currentCallID++
+ }
+
+ r.activeEntityCall = nil
+ r.activeEntityConfigs = nil
+ r.planInfo.currentRequestMessage = nil
+ r.planInfo.currentResponseMessage = nil
+ r.planInfo.currentResponseFieldIndex = 0
+ r.planInfo.responseMessageAncestors = nil
+ r.planInfo.responseFieldIndexAncestors = nil
+}
+
+func (r *rpcPlanVisitor) entityLookupRequest(requestName, methodName string, cfg entityCallConfig) RPCMessage {
keyFieldMessage := &RPCMessage{
- Name: r.rpcMethodName() + "Key",
+ Name: methodName + "Key",
}
- for _, key := range entityInfo.keyFields {
+
+ for _, key := range cfg.keyFields {
keyFieldMessage.Fields = append(keyFieldMessage.Fields, RPCField{
Name: key.fieldName,
TypeName: key.fieldType,
@@ -680,26 +796,54 @@ func (r *rpcPlanVisitor) scaffoldEntityLookup() {
})
}
- r.planInfo.currentRequestMessage.Fields = []RPCField{
- {
- Name: "keys",
- TypeName: DataTypeMessage.String(),
- Repeated: true, // The inputs are always a list of objects
- JSONPath: "representations",
- Message: keyFieldMessage,
+ return RPCMessage{
+ Name: requestName,
+ Fields: []RPCField{
+ {
+ Name: "keys",
+ TypeName: DataTypeMessage.String(),
+ Repeated: true,
+ JSONPath: "representations",
+ Message: keyFieldMessage,
+ },
},
}
+}
- // The proto response message has a field `result` which is a list of entities.
- // As this is a special case we directly map it to _entities.
- r.planInfo.currentResponseMessage.Fields = []RPCField{
- {
- Name: "result",
- TypeName: DataTypeMessage.String(),
- JSONPath: "_entities",
- Repeated: true,
- },
+func (r *rpcPlanVisitor) resolveEntityKeyFields(node ast.Node, key string) ([]keyField, error) {
+ keyNames := strings.Fields(key)
+ fields := make([]keyField, 0, len(keyNames))
+ for _, keyFieldName := range keyNames {
+ fieldDef, ok := r.definition.NodeFieldDefinitionByName(node, ast.ByteSlice(keyFieldName))
+ if !ok {
+ return nil, fmt.Errorf("field %s not found in definition", keyFieldName)
+ }
+
+ fdt := r.definition.FieldDefinitionType(fieldDef)
+ ft := r.definition.Types[fdt]
+ fields = append(fields, keyField{
+ fieldName: keyFieldName,
+ fieldType: r.toDataType(&ft).String(),
+ })
+ }
+
+ slices.SortFunc(fields, func(left, right keyField) int {
+ return strings.Compare(left.fieldName, right.fieldName)
+ })
+
+ return fields, nil
+}
+
+func (r *rpcPlanVisitor) entityKeyTypeName(typeName string, fields []keyField) string {
+ return typeName + "By" + strings.Join(titleSlice(keyFieldNames(fields)), "And")
+}
+
+func keyFieldNames(fields []keyField) []string {
+ names := make([]string, 0, len(fields))
+ for _, field := range fields {
+ names = append(names, field.fieldName)
}
+ return names
}
func (r *rpcPlanVisitor) resolveServiceName() string {
@@ -801,10 +945,6 @@ func (r *rpcPlanVisitor) rpcMethodName() string {
// buildQueryMethodName constructs a method name for query operations.
func (r *rpcPlanVisitor) buildQueryMethodName() string {
- if r.planInfo.isEntityLookup && r.planInfo.entityInfo.name != "" {
- return "Lookup" + r.planInfo.entityInfo.keyTypeName
- }
-
return "Query" + cases.Title(language.Und, cases.NoLower).String(r.planInfo.operationFieldName)
}
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/grpc_datasource.go b/app/v2/pkg/engine/datasource/grpc_datasource/grpc_datasource.go
index a718b71..97938f8 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/grpc_datasource.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/grpc_datasource.go
@@ -12,6 +12,7 @@ import (
"errors"
"fmt"
"strconv"
+ "sync"
"github.com/tidwall/gjson"
"github.com/wundergraph/astjson"
@@ -92,17 +93,46 @@ func (d *DataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer)
a := astjson.Arena{}
root := a.NewObject()
+ entityResults := make([]*astjson.Value, len(variables.Get("representations").Array()))
+
+ invokeErrs := make([]error, len(invocations))
+ var wg sync.WaitGroup
+ for i := range invocations {
+ wg.Add(1)
+ go func(idx int) {
+ defer wg.Done()
+ invocation := invocations[idx]
+ methodName := fmt.Sprintf("/%s/%s", invocation.ServiceName, invocation.MethodName)
+ invokeErrs[idx] = d.cc.Invoke(ctx, methodName, invocation.Input, invocation.Output)
+ }(i)
+ }
+ wg.Wait()
- // make gRPC calls
- for _, invocation := range invocations {
- // Invoke the gRPC method - this will populate invocation.Output
- methodName := fmt.Sprintf("/%s/%s", invocation.ServiceName, invocation.MethodName)
- err := d.cc.Invoke(ctx, methodName, invocation.Input, invocation.Output)
- if err != nil {
- out.Write(writeErrorBytes(err))
+ for i, invocation := range invocations {
+ if invokeErrs[i] != nil {
+ out.Write(writeErrorBytes(invokeErrs[i]))
return nil
}
+ if invocation.Entity != nil {
+ entities, err := d.marshalEntityResponseJSON(&a, &invocation.Call.Response, invocation.Output)
+ if err != nil {
+ return err
+ }
+
+ for idx, position := range invocation.Entity.Positions {
+ if position < 0 || position >= len(entityResults) {
+ continue
+ }
+ if idx < len(entities) {
+ entityResults[position] = entities[idx]
+ continue
+ }
+ entityResults[position] = a.NewNull()
+ }
+ continue
+ }
+
responseJSON, err := d.marshalResponseJSON(&a, &invocation.Call.Response, invocation.Output)
if err != nil {
return err
@@ -115,6 +145,17 @@ func (d *DataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer)
}
}
+ if len(entityResults) > 0 {
+ entities := a.NewArray()
+ for i, entity := range entityResults {
+ if entity == nil {
+ entity = a.NewNull()
+ }
+ entities.SetArrayItem(i, entity)
+ }
+ root.Set("_entities", entities)
+ }
+
data := a.NewObject()
data.Set("data", root)
out.Write(data.MarshalTo(nil))
@@ -122,6 +163,30 @@ func (d *DataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer)
return nil
}
+func (d *DataSource) marshalEntityResponseJSON(arena *astjson.Arena, message *RPCMessage, data protoref.Message) ([]*astjson.Value, error) {
+ if message == nil || len(message.Fields) == 0 {
+ return nil, nil
+ }
+
+ field := message.Fields[0]
+ fd := data.Descriptor().Fields().ByName(protoref.Name(field.Name))
+ if fd == nil || !fd.IsList() {
+ return nil, nil
+ }
+
+ list := data.Get(fd).List()
+ entities := make([]*astjson.Value, 0, list.Len())
+ for i := 0; i < list.Len(); i++ {
+ value, err := d.marshalResponseJSON(arena, field.Message, list.Get(i).Message())
+ if err != nil {
+ return nil, err
+ }
+ entities = append(entities, value)
+ }
+
+ return entities, nil
+}
+
// LoadWithFiles implements resolve.DataSource interface.
// Similar to Load, but handles file uploads if needed.
//
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/grpc_datasource_test.go b/app/v2/pkg/engine/datasource/grpc_datasource/grpc_datasource_test.go
index aae0711..2ec38fc 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/grpc_datasource_test.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/grpc_datasource_test.go
@@ -448,6 +448,47 @@ func TestMarshalResponseJSON(t *testing.T) {
require.Equal(t, `{"_entities":[{"__typename":"Product","id":"123","name_different":"test","price_different":123.45}]}`, responseJSON.String())
}
+func TestDataSourceLoadEntityLookupPreservesRepresentationOrder(t *testing.T) {
+ conn, cleanup := setupTestGRPCServer(t)
+ t.Cleanup(cleanup)
+
+ schemaDoc := grpctest.MustGraphQLSchema(t)
+ query := `query EntityLookup($representations: [_Any!]!) { _entities(representations: $representations) { ... on Product { __typename id name } ... on Storage { __typename id name location } } }`
+ queryDoc, report := astparser.ParseGraphqlDocumentString(query)
+ require.False(t, report.HasErrors(), report.Error())
+
+ compiler, err := NewProtoCompiler(grpctest.MustProtoSchema(t), testMapping())
+ require.NoError(t, err)
+
+ ds, err := NewDataSource(conn, DataSourceConfig{
+ Operation: &queryDoc,
+ Definition: &schemaDoc,
+ SubgraphName: "Products",
+ Compiler: compiler,
+ Mapping: testMapping(),
+ })
+ require.NoError(t, err)
+
+ output := new(bytes.Buffer)
+ input := `{"query":"` + query + `","body":{"variables":{"representations":[{"__typename":"Product","id":"p1"},{"__typename":"Storage","id":"s1"},{"__typename":"Product","id":"p2"}]}}}`
+ err = ds.Load(context.Background(), []byte(input), output)
+ require.NoError(t, err)
+
+ var response struct {
+ Data struct {
+ Entities []map[string]any `json:"_entities"`
+ } `json:"data"`
+ }
+ require.NoError(t, json.Unmarshal(output.Bytes(), &response))
+ require.Len(t, response.Data.Entities, 3)
+ require.Equal(t, "Product", response.Data.Entities[0]["__typename"])
+ require.Equal(t, "p1", response.Data.Entities[0]["id"])
+ require.Equal(t, "Storage", response.Data.Entities[1]["__typename"])
+ require.Equal(t, "s1", response.Data.Entities[1]["id"])
+ require.Equal(t, "Product", response.Data.Entities[2]["__typename"])
+ require.Equal(t, "p2", response.Data.Entities[2]["id"])
+}
+
// Test_DataSource_Load_WithAnimalInterface tests the datasource with Animal interface types (Cat/Dog)
// using a bufconn connection to the mock service
func Test_DataSource_Load_WithAnimalInterface(t *testing.T) {
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/mapping_test_helper.go b/app/v2/pkg/engine/datasource/grpc_datasource/mapping_test_helper.go
index 146f371..77782ab 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/engine/datasource/grpc_datasource/mapping_test_helper.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/mapping_test_helper.go
@@ -218,23 +218,23 @@ func testMapping() *GRPCMapping {
},
},
SubscriptionRPCs: RPCConfigMap{},
- EntityRPCs: map[string]EntityRPCConfig{
- "Product": {
+ EntityRPCs: map[string][]EntityRPCConfig{
+ "Product": {{
Key: "id",
RPCConfig: RPCConfig{
RPC: "LookupProductById",
Request: "LookupProductByIdRequest",
Response: "LookupProductByIdResponse",
},
- },
- "Storage": {
+ }},
+ "Storage": {{
Key: "id",
RPCConfig: RPCConfig{
RPC: "LookupStorageById",
Request: "LookupStorageByIdRequest",
Response: "LookupStorageByIdResponse",
},
- },
+ }},
},
EnumValues: map[string][]EnumValueMapping{
"CategoryKind": {
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/grpctest/mapping/mapping.go b/app/v2/pkg/grpctest/mapping/mapping.go
index 842a029..25e6f43 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-02-28__13-38-10__gpt-5-4/app/v2/pkg/grpctest/mapping/mapping.go
+++ b/app/v2/pkg/grpctest/mapping/mapping.go
@@ -225,23 +225,23 @@ func DefaultGRPCMapping() *grpcdatasource.GRPCMapping {
},
},
SubscriptionRPCs: grpcdatasource.RPCConfigMap{},
- EntityRPCs: map[string]grpcdatasource.EntityRPCConfig{
- "Product": {
+ EntityRPCs: map[string][]grpcdatasource.EntityRPCConfig{
+ "Product": {{
Key: "id",
RPCConfig: grpcdatasource.RPCConfig{
RPC: "LookupProductById",
Request: "LookupProductByIdRequest",
Response: "LookupProductByIdResponse",
},
- },
- "Storage": {
+ }},
+ "Storage": {{
Key: "id",
RPCConfig: grpcdatasource.RPCConfig{
RPC: "LookupStorageById",
Request: "LookupStorageByIdRequest",
Response: "LookupStorageByIdResponse",
},
- },
+ }},
},
EnumValues: map[string][]grpcdatasource.EnumValueMapping{
"CategoryKind": {