diff --git a/.changeset/two-dryers-deny.md b/.changeset/two-dryers-deny.md new file mode 100644 index 000000000..10d74e1f9 --- /dev/null +++ b/.changeset/two-dryers-deny.md @@ -0,0 +1,5 @@ +--- +"@apollo/query-planner": patch +--- + +Ensure query variables used in the directives applied at the operation level are retained in subgraph queries (#2986) diff --git a/query-planner-js/src/__tests__/buildPlan.test.ts b/query-planner-js/src/__tests__/buildPlan.test.ts index 932e827d4..ae23c0b2d 100644 --- a/query-planner-js/src/__tests__/buildPlan.test.ts +++ b/query-planner-js/src/__tests__/buildPlan.test.ts @@ -8719,4 +8719,42 @@ describe('handles operations with directives', () => { } `); }); + + test('subgraph query retains the query variables used in the directives applied to the query', () => { + const subgraph1 = { + name: 'Subgraph1', + typeDefs: gql` + directive @withArgs(arg1: String) on QUERY + + type Query { + test: String! + } + `, + }; + + const subgraph2 = { + name: 'Subgraph2', + typeDefs: gql` + directive @withArgs(arg1: String) on QUERY + `, + }; + + const query = gql` + query testQuery($some_var: String!) @withArgs(arg1: $some_var) { + test + } + `; + + const [api, qp] = composeAndCreatePlanner(subgraph1, subgraph2); + const op = operationFromDocument(api, query); + const queryPlan = qp.buildQueryPlan(op); + const fetch_nodes = findFetchNodes(subgraph1.name, queryPlan.node); + expect(fetch_nodes).toHaveLength(1); + // Note: `($some_var: String!)` used to be missing. + expect(parse(fetch_nodes[0].operation)).toMatchInlineSnapshot(` + query testQuery__Subgraph1__0($some_var: String!) @withArgs(arg1: $some_var) { + test + } + `); // end of test + }); }); // end of `describe` diff --git a/query-planner-js/src/buildPlan.ts b/query-planner-js/src/buildPlan.ts index cdf337135..03ce32394 100644 --- a/query-planner-js/src/buildPlan.ts +++ b/query-planner-js/src/buildPlan.ts @@ -25,6 +25,7 @@ import { Variable, VariableDefinition, VariableDefinitions, + VariableCollector, newDebugLogger, selectionOfElement, selectionSetOfElement, @@ -4753,6 +4754,20 @@ function representationsVariableDefinition(schema: Schema): VariableDefinition { return new VariableDefinition(schema, representationsVariable, representationsType); } +// Collects all variables used in the operation to be created. +// - It's computed based on its selection set and directives. +function collectUsedVariables(selectionSet: SelectionSet, operationDirectives?: readonly Directive[]) { + const collector = new VariableCollector(); + selectionSet.collectVariables(collector); + + if (operationDirectives) { + for (const applied of operationDirectives) { + collector.collectInArguments(applied.arguments()); + } + } + return collector.variables(); +} + function operationForEntitiesFetch( subgraphSchema: Schema, selectionSet: SelectionSet, @@ -4763,7 +4778,7 @@ function operationForEntitiesFetch( const variableDefinitions = new VariableDefinitions(); variableDefinitions.add(representationsVariableDefinition(subgraphSchema)); variableDefinitions.addAll( - allVariableDefinitions.filter(selectionSet.usedVariables()), + allVariableDefinitions.filter(collectUsedVariables(selectionSet, directives)), ); const queryType = subgraphSchema.schemaDefinition.rootType('query'); @@ -4799,6 +4814,6 @@ function operationForQueryFetch( // Note that this is called _before_ named fragments reuse is attempted, so there is not spread in // the selection, hence the `undefined` for fragments. return new Operation(subgraphSchema, rootKind, selectionSet, - allVariableDefinitions.filter(selectionSet.usedVariables()), + allVariableDefinitions.filter(collectUsedVariables(selectionSet, directives)), /*fragments*/undefined, operationName, directives); }