From 062572b3253e8640b60a0bf58b83945094b76b6f Mon Sep 17 00:00:00 2001
From: Duckki Oe <duckki.oe@apollographql.com>
Date: Mon, 14 Oct 2024 10:03:49 -0700
Subject: [PATCH] fix(query-planner): fixed missing referenced variables in the
 `variableUsages` 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                  |  9 +++++++++
 internals-js/src/operations.ts                   |  5 +++++
 query-planner-js/src/__tests__/buildPlan.test.ts |  6 ++++--
 query-planner-js/src/buildPlan.ts                | 16 ++++++++++------
 4 files changed, 28 insertions(+), 8 deletions(-)
 create mode 100644 .changeset/stupid-geese-camp.md

diff --git a/.changeset/stupid-geese-camp.md b/.changeset/stupid-geese-camp.md
new file mode 100644
index 000000000..99967382f
--- /dev/null
+++ b/.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.
diff --git a/internals-js/src/operations.ts b/internals-js/src/operations.ts
index 9e2f4be5e..a88bd1865 100644
--- a/internals-js/src/operations.ts
+++ b/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,
diff --git a/query-planner-js/src/__tests__/buildPlan.test.ts b/query-planner-js/src/__tests__/buildPlan.test.ts
index 155c48530..012a229e6 100644
--- a/query-planner-js/src/__tests__/buildPlan.test.ts
+++ b/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', () => {
diff --git a/query-planner-js/src/buildPlan.ts b/query-planner-js/src/buildPlan.ts
index 7dfabd9f2..ca5a2ebdf 100644
--- a/query-planner-js/src/buildPlan.ts
+++ b/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,