fix: fixed another regression created by PR apollographql#2967

- Ensure query variables used in the directives applied at the operation level are retained in subgraph queries

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`

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<any>[]) {
+ 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);
 }