fix(query-planner): fixed missing referenced variables in the `variab…

…leUsages` field (#3166)

Fixed missing referenced variables in the `variableUsages` field of
fetch operations

Query variables used in fetch operation should be listed in the
`variableUsages` field. However, there was a bug where variables
referenced by query-level directives could be missing in the field.

The previous PR (#2986) fixed a similar issue in fetch queries, but it
didn't fully fix the issue by failing to update the `variableUsages`
field. This PR completes the remaining issue.

<!-- [FED-387] -->


[FED-387]:
https://apollographql.atlassian.net/browse/FED-387?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

.changeset/stupid-geese-camp.md

@@ -0,0 +1,9 @@
+---
+"@apollo/query-planner": patch
+"@apollo/query-graphs": patch
+"@apollo/federation-internals": patch
+---
+
+Fixed missing referenced variables in the `variableUsages` field of fetch operations
+
+Query variables used in fetch operation should be listed in the `variableUsages` field. However, there was a bug where variables referenced by query-level directives could be missing in the field.

internals-js/src/operations.ts

@@ -1193,6 +1193,11 @@ export class NamedFragmentDefinition extends DirectiveTargetElement<NamedFragmen
     }
   }
 
+  collectVariables(collector: VariableCollector) {
+    this.selectionSet.collectVariables(collector);
+    this.collectVariablesInAppliedDirectives(collector);
+  }
+
   toFragmentDefinitionNode() : FragmentDefinitionNode {
     return {
       kind: Kind.FRAGMENT_DEFINITION,

query-planner-js/src/__tests__/buildPlan.test.ts

@@ -8977,8 +8977,10 @@ describe('handles operations with directives', () => {
       query testQuery__Subgraph1__0($some_var: String!) @withArgs(arg1: $some_var) {
         test
       }
-    `); // end of test
-  });
+    `);
+    // Make sure the `variableUsage` also captures the variable as well.
+    expect(fetch_nodes[0].variableUsages?.includes('some_var')).toBe(true);
+  }); // end of test
 }); // end of `describe`
 
 describe('@fromContext impacts on query planning', () => {

query-planner-js/src/buildPlan.ts

@@ -1600,22 +1600,26 @@ class FetchGroup {
       );
     }
 
-    // collect all used variables in the selection and in used Fragments
-    const usedVariables = new Set(selection.usedVariables().map(v => v.name));
+    // collect all used variables in the selection and the used directives and fragments.
+    const collector = new VariableCollector();
+    // Note: The operation's selectionSet includes `representations` variable,
+    //       thus we use `selection` here instead.
+    selection.collectVariables(collector);
+    operation.collectVariablesInAppliedDirectives(collector);
     if (operation.fragments) {
       for (const namedFragment of operation.fragments.definitions()) {
-        namedFragment.selectionSet.usedVariables().forEach(v => {
-          usedVariables.add(v.name);  
-        });
+        namedFragment.collectVariables(collector);
       }
     }
+    const usedVariables = collector.variables();
+
     const operationDocument = operationToDocument(operation);
     const fetchNode: FetchNode = {
       kind: 'Fetch',
       id: this.id,
       serviceName: this.subgraphName,
       requires: inputNodes ? trimSelectionNodes(inputNodes.selections) : undefined,
-      variableUsages: Array.from(usedVariables),
+      variableUsages: usedVariables.map((v) => v.name),
       operation: stripIgnoredCharacters(print(operationDocument)),
       operationKind: schemaRootKindToOperationKind(operation.rootKind),
       operationName: operation.name,