feat: add support for multiple key directives

execution/engine/execution_engine_grpc_test.go

@@ -234,7 +234,7 @@ func TestGRPCSubgraphExecution(t *testing.T) {
 			Query:         "query UserQuery { users { id name } }",
 		}
 
-		response, err := executeOperation(t, conn, operation)
+		response, err := executeOperation(t, conn, operation, withGRPCMapping(mapping.MustDefaultGRPCMapping(t)))
 		require.NoError(t, err)
 		require.Equal(t, `{"data":{"users":[{"id":"user-1","name":"User 1"},{"id":"user-2","name":"User 2"},{"id":"user-3","name":"User 3"}]}}`, response)
 	})
@@ -255,7 +255,7 @@ func TestGRPCSubgraphExecution(t *testing.T) {
 			`,
 		}
 
-		response, err := executeOperation(t, conn, operation)
+		response, err := executeOperation(t, conn, operation, withGRPCMapping(mapping.MustDefaultGRPCMapping(t)))
 		require.NoError(t, err)
 		require.Equal(t, `{"data":{"user":{"id":"1","name":"User 1"}}}`, response)
 	})

v2/pkg/astvisitor/visitor.go

@@ -3998,3 +3998,27 @@ func (w *Walker) FieldDefinitionDirectiveArgumentValueByName(field int, directiv
 
 	return w.definition.DirectiveArgumentValueByName(directive, argumentName)
 }
+
+// InRootField returns true if the current field is a root field.
+func (w *Walker) InRootField() bool {
+	return w.CurrentKind == ast.NodeKindField &&
+		len(w.Ancestors) == 2 &&
+		w.Ancestors[0].Kind == ast.NodeKindOperationDefinition
+}
+
+// ResolveInlineFragment returns the inline fragment ref if the current field is inside of
+// an inline fragment.
+// It returns the inline fragment ref and true if the current field is inside of an inline fragment.
+// It returns -1 and false if the current field is not inside of an inline fragment.
+func (w *Walker) ResolveInlineFragment() (int, bool) {
+	if len(w.Ancestors) < 2 {
+		return ast.InvalidRef, false
+	}
+
+	node := w.Ancestors[len(w.Ancestors)-2]
+	if node.Kind != ast.NodeKindInlineFragment {
+		return ast.InvalidRef, false
+	}
+
+	return node.Ref, true
+}

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource.go

@@ -368,11 +368,12 @@ func (p *Planner[T]) ConfigureFetch() resolve.FetchConfiguration {
 		}
 
 		dataSource, err = grpcdatasource.NewDataSource(p.grpcClient, grpcdatasource.DataSourceConfig{
-			Operation:  &opDocument,
-			Definition: p.config.schemaConfiguration.upstreamSchemaAst,
-			Mapping:    p.config.grpc.Mapping,
-			Compiler:   p.config.grpc.Compiler,
-			Disabled:   p.config.grpc.Disabled,
+			Operation:         &opDocument,
+			Definition:        p.config.schemaConfiguration.upstreamSchemaAst,
+			Mapping:           p.config.grpc.Mapping,
+			Compiler:          p.config.grpc.Compiler,
+			Disabled:          p.config.grpc.Disabled,
+			FederationConfigs: p.dataSourcePlannerConfig.RequiredFields,
 			// TODO: remove fallback logic in visitor for subgraph name and
 			// add proper error handling if the subgraph name is not set in the mapping
 			SubgraphName: p.dataSourceConfig.Name(),

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

@@ -3,6 +3,7 @@ package grpcdatasource
 import (
 	"context"
 	"fmt"
+	"slices"
 
 	"github.com/bufbuild/protocompile"
 	"github.com/tidwall/gjson"
@@ -433,6 +434,12 @@ func (p *RPCCompiler) buildProtoMessage(inputMessage Message, rpcMessage *RPCMes
 			for _, element := range elements {
 				switch field.Type {
 				case DataTypeMessage:
+					// If we handle entity lookups, we get a list of representation variables that need to
+					// be applied to the correct type names.
+					if !isAllowedForTypename(rpcField.Message, element) {
+						continue
+					}
+
 					fieldMsg := p.buildProtoMessage(p.doc.Messages[field.MessageRef], rpcField.Message, element)
 					list.Append(protoref.ValueOfMessage(fieldMsg))
 				default:
@@ -785,3 +792,26 @@ func (p *RPCCompiler) parseField(f protoref.FieldDescriptor) Field {
 		MessageRef: -1,
 	}
 }
+
+func isAllowedForTypename(message *RPCMessage, element gjson.Result) bool {
+	if message == nil {
+		// We assume that having a nil message expects a null value.
+		return true
+	}
+
+	// If we don't have a member types, we assume that the message is allowed for all types.
+	if message.MemberTypes == nil {
+		return true
+	}
+
+	typeName := element.Get("__typename")
+	if !typeName.Exists() {
+		// If we don't have a type name, we assume that the message is allowed for all types.
+		return true
+	}
+
+	typeString := typeName.String()
+
+	// If we have a type name, we need to check if the message is restricted to a specific type.
+	return slices.Contains(message.MemberTypes, typeString)
+}

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

@@ -1,9 +1,19 @@
 package grpcdatasource
 
+import (
+	"strings"
+
+	"github.com/wundergraph/graphql-go-tools/v2/pkg/lexer/runes"
+)
+
 type (
 	// RPCConfigMap is a map of RPC names to RPC configurations
+	// The key is the field name in the GraphQL operation type (query, mutation, subscription).
+	// The value is the RPC configuration for that field.
 	RPCConfigMap map[string]RPCConfig
 	// FieldMap defines the mapping between a GraphQL field and a gRPC field
+	// The key is the field name in the GraphQL type.
+	// The value is the FieldMapData for that field which contains the target name and the argument mappings.
 	FieldMap map[string]FieldMapData
 )
 
@@ -17,22 +27,27 @@ type GRPCMapping struct {
 	// 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
+	// The key is the type name and the value is a list of EntityRPCConfig for that type
+	EntityRPCs map[string][]EntityRPCConfig
 	// Fields defines the field mappings between GraphQL types and gRPC messages
+	// The key is the type name and the value is a map of field name to FieldMapData for that type
 	Fields map[string]FieldMap
 	// EnumValues defines the enum values for each enum type
+	// The key is the enum type name and the value is a list of EnumValueMapping for that enum type
 	EnumValues map[string][]EnumValueMapping
 }
 
+// EnumValueMapping defines the mapping between a GraphQL enum value and a gRPC enum value
 type EnumValueMapping struct {
-	Value       string
-	TargetValue string
+	Value       string // The GraphQL enum value
+	TargetValue string // The gRPC enum value
 }
 
+// GRPCConfiguration defines the configuration for a gRPC datasource
 type GRPCConfiguration struct {
-	Disabled bool
-	Mapping  *GRPCMapping
-	Compiler *RPCCompiler
+	Disabled bool         // Whether the RPC is disabled
+	Mapping  *GRPCMapping // The mapping between GraphQL types and gRPC messages
+	Compiler *RPCCompiler // The compiler for the RPC
 }
 
 // RPCConfig defines the configuration for a specific RPC operation
@@ -53,9 +68,10 @@ type EntityRPCConfig struct {
 	RPCConfig
 }
 
+// FieldMapData defines the mapping between a GraphQL field and a gRPC field
 type FieldMapData struct {
-	TargetName       string
-	ArgumentMappings FieldArgumentMap
+	TargetName       string           // The name of the gRPC field
+	ArgumentMappings FieldArgumentMap // The mapping between GraphQL field arguments and gRPC request arguments
 }
 
 // FieldArgumentMap defines the mapping between a GraphQL field argument and a gRPC field
@@ -121,3 +137,97 @@ func (g *GRPCMapping) ResolveEnumValue(enumName, enumValue string) (string, bool
 
 	return "", false
 }
+
+func (g *GRPCMapping) ResolveEntityRPCConfig(typeName, key string) (RPCConfig, bool) {
+	rpcConfig, ok := g.EntityRPCs[typeName]
+	if !ok {
+		return RPCConfig{}, false
+	}
+
+	for _, ei := range rpcConfig {
+		if compareKeyFields(ei.Key, key) {
+			return ei.RPCConfig, true
+		}
+
+	}
+
+	return RPCConfig{}, false
+}
+
+type keySet map[string]struct{}
+
+func (k keySet) add(keys ...string) {
+	for _, key := range keys {
+		trimmedKey := strings.TrimSpace(key)
+		if trimmedKey != "" {
+			k[trimmedKey] = struct{}{}
+		}
+	}
+}
+
+func (k keySet) equals(other keySet) bool {
+	if len(k) != len(other) {
+		return false
+	}
+
+	for key := range k {
+		if _, ok := other[key]; !ok {
+			return false
+		}
+	}
+
+	return true
+}
+
+// We compare only top level key
+func compareKeyFields(left, right string) bool {
+	if left == right {
+		return true
+	}
+
+	left = stripSelectionSets(left)
+	right = stripSelectionSets(right)
+
+	leftKeys := strings.Split(left, " ")
+	rightKeys := strings.Split(right, " ")
+
+	leftSet := make(keySet)
+	leftSet.add(leftKeys...)
+
+	rightSet := make(keySet)
+	rightSet.add(rightKeys...)
+
+	return leftSet.equals(rightSet)
+}
+
+func stripSelectionSets(keyString string) string {
+	depth := 0
+
+	var prev rune
+
+	keyString = strings.ReplaceAll(keyString, ",", " ")
+
+	var sb strings.Builder
+
+	for _, r := range keyString {
+		switch r {
+		case runes.LBRACE:
+			depth++
+		case runes.RBRACE:
+			if depth == 0 {
+				continue
+			}
+
+			depth--
+		default:
+			if depth != 0 || (r == runes.SPACE && prev == runes.SPACE) {
+				break
+			}
+
+			sb.WriteRune(r)
+			prev = r
+		}
+	}
+
+	return strings.TrimSpace(sb.String())
+}