Repository: jensneuse/graphql-go-tools
Pull request #307 (2021-11-05 15:54:02)

add supporting multiple key directives
Add supporting multiple key directives.
According to [federation-spec](https://www.apollographql.com/docs/federation/federation-spec/#key) multiple keys can be defined on a single object type:
```
type Product @key(fields: "upc") @key(fields: "sku") {
  upc: UPC!
  sku: SKU!
  name: String
}

```


Patch:
diff --git a/pkg/engine/datasource/graphql_datasource/graphql_datasource.go b/pkg/engine/datasource/graphql_datasource/graphql_datasource.go
index 983f016b..8e25678d 100644
--- a/pkg/engine/datasource/graphql_datasource/graphql_datasource.go
+++ b/pkg/engine/datasource/graphql_datasource/graphql_datasource.go
@@ -370,13 +370,42 @@ func (p *Planner) addRepresentationsVariable() {
 		p.visitor.Walker.StopWithInternalErr(fmt.Errorf("GraphQL Planner: failed parsing Federation SDL"))
 		return
 	}
-	directive := -1
+
+	fields := p.selectFederationKeyFields(doc)
+	if len(fields) == 0 {
+		return
+	}
+
+	representationsJson, _ := sjson.SetRawBytes(nil, "__typename", []byte("\""+p.lastFieldEnclosingTypeName+"\""))
+	for i := range fields {
+		objectVariable := &resolve.ObjectVariable{
+			Path:                 []string{fields[i]},
+			RenderAsGraphQLValue: true,
+		}
+		fieldDef := p.fieldDefinition(fields[i], p.lastFieldEnclosingTypeName)
+		if fieldDef == nil {
+			continue
+		}
+		objectVariable.SetJsonValueType(p.visitor.Definition, fieldDef.Type)
+		variable, exists := p.variables.AddVariable(objectVariable)
+		if exists {
+			continue
+		}
+		representationsJson, _ = sjson.SetRawBytes(representationsJson, fields[i], []byte(variable))
+	}
+	representationsJson = append([]byte("["), append(representationsJson, []byte("]")...)...)
+	p.upstreamVariables, _ = sjson.SetRawBytes(p.upstreamVariables, "representations", representationsJson)
+	p.extractEntities = true
+}
+
+func (p *Planner) selectFederationKeyFields(doc *ast.Document) []string {
+	var directiveRefs []int
+
 	for i := range doc.ObjectTypeExtensions {
 		if p.lastFieldEnclosingTypeName == doc.ObjectTypeExtensionNameString(i) {
 			for _, j := range doc.ObjectTypeExtensions[i].Directives.Refs {
 				if doc.DirectiveNameString(j) == "key" {
-					directive = j
-					break
+					directiveRefs = append(directiveRefs, j)
 				}
 			}
 			break
@@ -386,47 +415,34 @@ func (p *Planner) addRepresentationsVariable() {
 		if p.lastFieldEnclosingTypeName == doc.ObjectTypeDefinitionNameString(i) {
 			for _, j := range doc.ObjectTypeDefinitions[i].Directives.Refs {
 				if doc.DirectiveNameString(j) == "key" {
-					directive = j
-					break
+					directiveRefs = append(directiveRefs, j)
 				}
 			}
 			break
 		}
 	}
-	if directive == -1 {
-		return
-	}
-	value, exists := doc.DirectiveArgumentValueByName(directive, []byte("fields"))
-	if !exists {
-		return
-	}
-	if value.Kind != ast.ValueKindString {
-		return
+
+	if len(directiveRefs) == 0 {
+		return nil
 	}
-	fieldsStr := doc.StringValueContentString(value.Ref)
-	fields := strings.Split(fieldsStr, " ")
-	representationsJson, _ := sjson.SetRawBytes(nil, "__typename", []byte("\""+p.lastFieldEnclosingTypeName+"\""))
-	for i := range fields {
-		objectVariable := &resolve.ObjectVariable{
-			Path:                 []string{fields[i]},
-			RenderAsGraphQLValue: true,
-		}
-		fieldDef := p.fieldDefinition(fields[i], p.lastFieldEnclosingTypeName)
-		if fieldDef == nil {
+
+	var fields []string
+
+	for _, directiveRef := range directiveRefs {
+		value, exists := doc.DirectiveArgumentValueByName(directiveRef, []byte("fields"))
+		if !exists {
 			continue
 		}
-		objectVariable.SetJsonValueType(p.visitor.Definition, fieldDef.Type)
-		variable, exists := p.variables.AddVariable(objectVariable)
-		if exists {
+		if value.Kind != ast.ValueKindString {
 			continue
 		}
-		representationsJson, _ = sjson.SetRawBytes(representationsJson, fields[i], []byte(variable))
+
+		fieldsStr := doc.StringValueContentString(value.Ref)
+		fields = append(fields, strings.Split(fieldsStr, " ")...)
 	}
-	representationsJson = append([]byte("["), append(representationsJson, []byte("]")...)...)
-	p.upstreamVariables, _ = sjson.SetRawBytes(p.upstreamVariables, "representations", representationsJson)
-	p.extractEntities = true
-}
 
+	return fields
+}
 func (p *Planner) fieldDefinition(fieldName, typeName string) *ast.FieldDefinition {
 	node, ok := p.visitor.Definition.Index.FirstNodeByNameStr(typeName)
 	if !ok {


Test patch:
diff --git a/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go b/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go
index f6fde194..8ba6b8e0 100644
--- a/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go
+++ b/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go
@@ -2365,7 +2365,7 @@ func TestGraphQLDataSource(t *testing.T) {
 																	&resolve.BatchFetch{
 																		Fetch: &resolve.SingleFetch{
 																			BufferId:   2,
-																			Input:      `{"method":"POST","url":"http://product.service","body":{"query":"query($representations: [_Any!]!){_entities(representations: $representations){... on Product {name price}}}","variables":{"representations":[{"upc":$$0$$,"__typename":"Product"}]}}}`,
+																			Input:      `{"method":"POST","url":"http://product.service","body":{"query":"query($representations: [_Any!]!){_entities(representations: $representations){... on Product {name price}}}","variables":{"representations":[{"name":$$1$$,"upc":$$0$$,"__typename":"Product"}]}}}`,
 																			DataSource: &Source{},
 																			Variables: resolve.NewVariables(
 																				&resolve.ObjectVariable{
@@ -2373,6 +2373,11 @@ func TestGraphQLDataSource(t *testing.T) {
 																					JsonValueType:        jsonparser.String,
 																					RenderAsGraphQLValue: true,
 																				},
+																				&resolve.ObjectVariable{
+																					Path:                 []string{"name"},
+																					JsonValueType:        jsonparser.String,
+																					RenderAsGraphQLValue: true,
+																				},
 																			),
 																			DataSourceIdentifier: []byte("graphql_datasource.Source"),
 																			ProcessResponseConfig: resolve.ProcessResponseConfig{
@@ -2385,13 +2390,18 @@ func TestGraphQLDataSource(t *testing.T) {
 																	&resolve.BatchFetch{
 																		Fetch: &resolve.SingleFetch{
 																			BufferId: 3,
-																			Input:    `{"method":"POST","url":"http://review.service","body":{"query":"query($representations: [_Any!]!){_entities(representations: $representations){... on Product {reviews {body author {id username}}}}}","variables":{"representations":[{"upc":$$0$$,"__typename":"Product"}]}}}`,
+																			Input:    `{"method":"POST","url":"http://review.service","body":{"query":"query($representations: [_Any!]!){_entities(representations: $representations){... on Product {reviews {body author {id username}}}}}","variables":{"representations":[{"name":$$1$$,"upc":$$0$$,"__typename":"Product"}]}}}`,
 																			Variables: resolve.NewVariables(
 																				&resolve.ObjectVariable{
 																					Path:                 []string{"upc"},
 																					JsonValueType:        jsonparser.String,
 																					RenderAsGraphQLValue: true,
 																				},
+																				&resolve.ObjectVariable{
+																					Path:                 []string{"name"},
+																					JsonValueType:        jsonparser.String,
+																					RenderAsGraphQLValue: true,
+																				},
 																			),
 																			DataSource:           &Source{},
 																			DataSourceIdentifier: []byte("graphql_datasource.Source"),
@@ -2559,7 +2569,7 @@ func TestGraphQLDataSource(t *testing.T) {
 						},
 						Federation: FederationConfiguration{
 							Enabled:    true,
-							ServiceSDL: "extend type Query {topProducts(first: Int = 5): [Product]} type Product @key(fields: \"upc\") {upc: String! name: String! price: Int!}",
+							ServiceSDL: "extend type Query {topProducts(first: Int = 5): [Product]} type Product @key(fields: \"upc\") @key(fields: \"name\"){upc: String! name: String! price: Int!}",
 						},
 					}),
 					Factory: federationFactory,
@@ -2596,7 +2606,7 @@ func TestGraphQLDataSource(t *testing.T) {
 						},
 						Federation: FederationConfiguration{
 							Enabled:    true,
-							ServiceSDL: "type Review { body: String! author: User! @provides(fields: \"username\") product: Product! } extend type User @key(fields: \"id\") { id: ID! @external reviews: [Review] } extend type Product @key(fields: \"upc\") { upc: String! @external reviews: [Review] }",
+							ServiceSDL: "type Review { body: String! author: User! @provides(fields: \"username\") product: Product! } extend type User @key(fields: \"id\") { id: ID! @external reviews: [Review] } extend type Product @key(fields: \"upc\") @key(fields: \"name\") { upc: String! @external name: String! reviews: [Review] }",
 						},
 					}),
 				},