Repository: jensneuse/graphql-go-tools
Pull request #317 (2022-01-14 21:24:29)

Fix federation directive variables
Fixes jensneuse/graphql-go-tools#316.

When an entity query contains a field with a directive that includes a
variable, the variable isn't added to constructed upstream operation.
This PR makes it so variables used with directives are added as needed.


Patch:
diff --git a/pkg/engine/datasource/graphql_datasource/graphql_datasource.go b/pkg/engine/datasource/graphql_datasource/graphql_datasource.go
index 855444ed..d8033c16 100644
--- a/pkg/engine/datasource/graphql_datasource/graphql_datasource.go
+++ b/pkg/engine/datasource/graphql_datasource/graphql_datasource.go
@@ -81,8 +81,83 @@ func (p *Planner) addDirectiveToNode(directiveRef int, node ast.Node) {
 	if p.upstreamDefinition != nil && !p.upstreamDefinition.DirectiveIsAllowedOnNodeKind(upstreamDirectiveName, node.Kind, operationType) {
 		return
 	}
-	directive := p.visitor.Importer.ImportDirectiveWithRename(directiveRef, upstreamDirectiveName, p.visitor.Operation, p.upstreamOperation)
-	p.upstreamOperation.AddDirectiveToNode(directive, node)
+	upstreamDirective := p.visitor.Importer.ImportDirectiveWithRename(directiveRef, upstreamDirectiveName, p.visitor.Operation, p.upstreamOperation)
+	p.upstreamOperation.AddDirectiveToNode(upstreamDirective, node)
+
+	// The directive is allowed on the node, so we know it exists.
+	directive := p.visitor.Operation.Directives[directiveRef]
+
+	var variables []ast.Value
+
+	// Collect all the variable arguments.
+	if directive.HasArguments {
+		for _, argument := range directive.Arguments.Refs {
+			value := p.visitor.Operation.ArgumentValue(argument)
+			// TODO: also handle literal values that CONTAIN variables
+			if value.Kind == ast.ValueKindVariable {
+				variables = append(variables, value)
+			}
+		}
+	}
+
+	// Process each variable, adding it to the upstream operation and
+	// variables, if it hasn't already been added. Note: instead of looking
+	// up the type of the corresponding argument on the directive definition,
+	// this code assumes the type of the variable as defined in the operation
+	// is correct and uses the same (possibly mapped) type for the upstream
+	// operation.
+	for _, value := range variables {
+		variableName := p.visitor.Operation.VariableValueNameBytes(value.Ref)
+
+		for _, i := range p.visitor.Operation.OperationDefinitions[p.visitor.Walker.Ancestors[0].Ref].VariableDefinitions.Refs {
+			// Find the variable declaration in the downstream operation.
+			ref := p.visitor.Operation.VariableDefinitions[i].VariableValue.Ref
+			if !p.visitor.Operation.VariableValueNameBytes(ref).Equals(variableName) {
+				continue
+			}
+
+			// Look up the variable type.
+			variableType := p.visitor.Operation.VariableDefinitions[i].Type
+			typeName := p.visitor.Operation.ResolveTypeNameString(variableType)
+
+			renderer, err := resolve.NewJSONVariableRendererWithValidationFromTypeRef(p.visitor.Operation, p.visitor.Definition, variableType)
+			if err != nil {
+				continue
+			}
+
+			contextVariable := &resolve.ContextVariable{
+				Path:     []string{string(variableName)},
+				Renderer: renderer,
+			}
+
+			// Try to add the variable to the set of upstream variables.
+			contextVariableName, exists := p.variables.AddVariable(contextVariable)
+
+			// If the variable already exists, it also already exists in the
+			// upstream operation; there's nothing to add!
+			if exists {
+				continue
+			}
+
+			// Add the variable to the upstream operation. Be sure to map the
+			// downstream type to the upstream type, if needed.
+			upstreamVariable := p.upstreamOperation.ImportVariableValue(variableName)
+			upstreamTypeName := p.visitor.Config.Types.RenameTypeNameOnMatchStr(typeName)
+			importedType := p.visitor.Importer.ImportTypeWithRename(p.visitor.Operation.VariableDefinitions[i].Type, p.visitor.Operation, p.upstreamOperation, upstreamTypeName)
+			p.upstreamOperation.AddVariableDefinitionToOperationDefinition(p.nodes[0].Ref, upstreamVariable, importedType)
+
+			// Also copy any variable directives in the downstream operation to
+			// the upstream operation.
+			if add, ok := p.addDirectivesToVariableDefinitions[i]; ok {
+				for _, directive := range add {
+					p.addDirectiveToNode(directive, ast.Node{Kind: ast.NodeKindVariableDefinition, Ref: i})
+				}
+			}
+
+			// And finally add the variable to the upstream variables JSON.
+			p.upstreamVariables, _ = sjson.SetRawBytes(p.upstreamVariables, string(variableName), []byte(contextVariableName))
+		}
+	}
 }
 
 func (p *Planner) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) {


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 fa6c4bee..d7af783a 100644
--- a/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go
+++ b/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go
@@ -3292,6 +3292,197 @@ func TestGraphQLDataSource(t *testing.T) {
 			},
 		}))
 
+	t.Run("federation with variables", RunTest(federationTestSchema,
+		`	query MyReviews($publicOnly: Boolean!, $someSkipCondition: Boolean!) {
+						me {
+							reviews {
+								body
+								notes @skip(if: $someSkipCondition)
+								likes(filterToPublicOnly: $publicOnly)
+							}
+						}
+					}`,
+		"MyReviews",
+		&plan.SynchronousResponsePlan{
+			Response: &resolve.GraphQLResponse{
+				Data: &resolve.Object{
+					Fetch: &resolve.SingleFetch{
+						BufferId:              0,
+						Input:                 `{"method":"POST","url":"http://user.service","body":{"query":"{me {id}}"}}`,
+						DataSource:            &Source{},
+						DataSourceIdentifier:  []byte("graphql_datasource.Source"),
+						ProcessResponseConfig: resolve.ProcessResponseConfig{ExtractGraphqlResponse: true},
+					},
+					Fields: []*resolve.Field{
+						{
+							HasBuffer: true,
+							BufferID:  0,
+							Name:      []byte("me"),
+							Position: resolve.Position{
+								Line:   2,
+								Column: 7,
+							},
+							Value: &resolve.Object{
+								Fetch: &resolve.BatchFetch{
+									Fetch: &resolve.SingleFetch{
+										BufferId: 1,
+										Input:    `{"method":"POST","url":"http://review.service","body":{"query":"query($representations: [_Any!]!, $someSkipCondition: Boolean!, $publicOnly: Boolean!){_entities(representations: $representations){... on User {reviews {body notes @skip(if: $someSkipCondition) likes(filterToPublicOnly: $publicOnly)}}}}","variables":{"publicOnly":$$2$$,"someSkipCondition":$$1$$,"representations":[{"id":$$0$$,"__typename":"User"}]}}}`,
+										Variables: resolve.NewVariables(
+											&resolve.ObjectVariable{
+												Path:     []string{"id"},
+												Renderer: resolve.NewJSONVariableRendererWithValidation(`{"type":["string","integer"]}`),
+											},
+											&resolve.ContextVariable{
+												Path:     []string{"someSkipCondition"},
+												Renderer: resolve.NewJSONVariableRendererWithValidation(`{"type":"boolean"}`),
+											},
+											&resolve.ContextVariable{
+												Path:     []string{"publicOnly"},
+												Renderer: resolve.NewJSONVariableRendererWithValidation(`{"type":"boolean"}`),
+											},
+										),
+										DataSource:           &Source{},
+										DataSourceIdentifier: []byte("graphql_datasource.Source"),
+										ProcessResponseConfig: resolve.ProcessResponseConfig{
+											ExtractGraphqlResponse:    true,
+											ExtractFederationEntities: true,
+										},
+									},
+									BatchFactory: batchFactory,
+								},
+								Path:     []string{"me"},
+								Nullable: true,
+								Fields: []*resolve.Field{
+									{
+										HasBuffer: true,
+										BufferID:  1,
+										Name:      []byte("reviews"),
+										Position: resolve.Position{
+											Line:   3,
+											Column: 8,
+										},
+										Value: &resolve.Array{
+											Path:     []string{"reviews"},
+											Nullable: true,
+											Item: &resolve.Object{
+												Nullable: true,
+												Fields: []*resolve.Field{
+													{
+														Name: []byte("body"),
+														Value: &resolve.String{
+															Path: []string{"body"},
+														},
+														Position: resolve.Position{
+															Line:   4,
+															Column: 9,
+														},
+													},
+													{
+														Name: []byte("notes"),
+														Value: &resolve.String{
+															Path:     []string{"notes"},
+															Nullable: true,
+														},
+														Position: resolve.Position{
+															Line:   5,
+															Column: 9,
+														},
+													},
+													{
+														Name: []byte("likes"),
+														Value: &resolve.String{
+															Path: []string{"likes"},
+														},
+														Position: resolve.Position{
+															Line:   6,
+															Column: 9,
+														},
+													},
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		plan.Configuration{
+			DataSources: []plan.DataSourceConfiguration{
+				{
+					RootNodes: []plan.TypeField{
+						{
+							TypeName:   "Query",
+							FieldNames: []string{"me"},
+						},
+					},
+					ChildNodes: []plan.TypeField{
+						{
+							TypeName:   "User",
+							FieldNames: []string{"id"},
+						},
+					},
+					Custom: ConfigJson(Configuration{
+						Fetch: FetchConfiguration{
+							URL: "http://user.service",
+						},
+						Federation: FederationConfiguration{
+							Enabled:    true,
+							ServiceSDL: "extend type Query {me: User} type User @key(fields: \"id\"){ id: ID! }",
+						},
+					}),
+					Factory: federationFactory,
+				},
+				{
+					RootNodes: []plan.TypeField{
+						{
+							TypeName:   "User",
+							FieldNames: []string{"reviews"},
+						},
+					},
+					ChildNodes: []plan.TypeField{
+						{
+							TypeName:   "Review",
+							FieldNames: []string{"body", "notes", "likes"},
+						},
+						{
+							TypeName:   "User",
+							FieldNames: []string{"id"},
+						},
+					},
+					Factory: federationFactory,
+					Custom: ConfigJson(Configuration{
+						Fetch: FetchConfiguration{
+							URL: "http://review.service",
+						},
+						Federation: FederationConfiguration{
+							Enabled:    true,
+							ServiceSDL: "type Review { body: String! notes: String likes(filterToPublicOnly: Boolean): Int! } extend type User @key(fields: \"id\") { id: ID! @external reviews: [Review] }",
+						},
+					}),
+				},
+			},
+			Fields: []plan.FieldConfiguration{
+				{
+					TypeName:       "User",
+					FieldName:      "reviews",
+					RequiresFields: []string{"id"},
+				},
+				{
+					TypeName:  "Review",
+					FieldName: "likes",
+					Arguments: []plan.ArgumentConfiguration{
+						{
+							Name:       "filterToPublicOnly",
+							SourceType: plan.FieldArgumentSource,
+						},
+					},
+				},
+			},
+		}))
+
 	t.Run("federation with renamed schema", RunTest(renamedFederationTestSchema,
 		`	query MyReviews {
 						api_me {
@@ -4021,7 +4212,6 @@ func runTestOnTestDefinition(operation, operationName string, expectedPlan plan.
 }
 
 func BenchmarkFederationBatching(b *testing.B) {
-
 	userService := FakeDataSource(`{"data":{"me": {"id": "1234","username": "Me","__typename": "User"}}}`)
 	reviewsService := FakeDataSource(`{"data":{"_entities":[{"reviews": [{"body": "A highly effective form of birth control.","product": {"upc": "top-1","__typename": "Product"}},{"body": "Fedoras are one of the most fashionable hats around and can look great with a variety of outfits.","product": {"upc": "top-2","__typename": "Product"}}]}]}}`)
 	productsService := FakeDataSource(`{"data":{"_entities":[{"name": "Trilby"},{"name": "Fedora"}]}}`)
@@ -5012,6 +5202,8 @@ type Review {
   body: String!
   author: User!
   product: Product!
+  notes: String
+  likes(filterToPublicOnly: Boolean): Int!
 }
 
 type User {
diff --git a/pkg/engine/datasourcetesting/datasourcetesting.go b/pkg/engine/datasourcetesting/datasourcetesting.go
index f7af3a67..8967a70d 100644
--- a/pkg/engine/datasourcetesting/datasourcetesting.go
+++ b/pkg/engine/datasourcetesting/datasourcetesting.go
@@ -60,6 +60,5 @@ func RunTest(definition, operation, operationName string, expectedPlan plan.Plan
 		for _, extraCheck := range extraChecks {
 			extraCheck(t, op, actualPlan)
 		}
-
 	}
 }