Ensure query variables used in the directives applied at the operatio…

…n level are retained in subgraph queries (#2986)

This fixes another regression created by PR #2967.

Related: #2961

.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)

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);
 }