Fragment variable definitions erased in subgraph queries (#3119)

Fixes #3112

Before creating an Operation, we call `collectUsedVariables`, which will
only pull values from the selection set and not from fragments. This
isn't great, because a variable used in a fragment won't be collected,
and it doesn't make sense to collect variables from fragments because
it's before they are optimized and many will be unused.

The inelegant solution I came up with is to pass in available variables
in calls to `optimize` or `generateQueryFragments` for an operation
where we can add back in the unused variables. This should be ok,
because we are guaranteed that exactly one of them will get called by
`toPlanNode`. Pretty sure there won't be too much overhead added because
we'll only call this once per subgraph fetch.

---------

Co-authored-by: Sachin D. Shinde <sachin@apollographql.com>

Author: clenfest

.changeset/poor-seahorses-whisper.md

@@ -0,0 +1,6 @@
+---
+"@apollo/query-planner": patch
+"@apollo/federation-internals": patch
+---
+
+Fix issue where variable was not passed into subgraph when embedded in a fragment

internals-js/src/operations.ts

@@ -56,6 +56,8 @@ import { assert, mapKeys, mapValues, MapWithCachedArrays, MultiMap, SetMultiMap
 import { argumentsEquals, argumentsFromAST, isValidValue, valueToAST, valueToString } from "./values";
 import { v1 as uuidv1 } from 'uuid';
 
+export const DEFAULT_MIN_USAGES_TO_OPTIMIZE = 2;
+
 function validate(condition: any, message: () => string, sourceAST?: ASTNode): asserts condition {
   if (!condition) {
     throw ERRORS.INVALID_GRAPHQL.err(message(), { nodes: sourceAST });
@@ -934,25 +936,55 @@ export class Operation extends DirectiveTargetElement<Operation> {
       this.appliedDirectives,
     );
   }
+  
+  private collectUndefinedVariablesFromFragments(fragments: NamedFragments): Variable[] {
+    const collector = new VariableCollector();
+    for (const namedFragment of fragments.definitions()) {
+      namedFragment.selectionSet.usedVariables().forEach(v => {
+        if (!this.variableDefinitions.definition(v)) {
+          collector.add(v);
+        }
+      });
+    }
+    return collector.variables();
+  }
 
   // Returns a copy of this operation with the provided updated selection set and fragments.
-  private withUpdatedSelectionSetAndFragments(newSelectionSet: SelectionSet, newFragments: NamedFragments | undefined): Operation {
+  private withUpdatedSelectionSetAndFragments(
+    newSelectionSet: SelectionSet,
+    newFragments: NamedFragments | undefined,
+    allAvailableVariables?: VariableDefinitions,
+  ): Operation {
     if (this.selectionSet === newSelectionSet && newFragments === this.fragments) {
       return this;
     }
+    
+    let newVariableDefinitions = this.variableDefinitions;
+    if (allAvailableVariables && newFragments) {
+      const undefinedVariables = this.collectUndefinedVariablesFromFragments(newFragments);
+      if (undefinedVariables.length > 0) {
+        newVariableDefinitions = new VariableDefinitions();
+        newVariableDefinitions.addAll(this.variableDefinitions);
+        newVariableDefinitions.addAll(allAvailableVariables.filter(undefinedVariables));
+      }
+    }
 
     return new Operation(
       this.schema(),
       this.rootKind,
       newSelectionSet,
-      this.variableDefinitions,
+      newVariableDefinitions,
       newFragments,
       this.name,
       this.appliedDirectives,
     );
   }
 
-  optimize(fragments?: NamedFragments, minUsagesToOptimize: number = 2): Operation {
+  optimize(
+    fragments?: NamedFragments,
+    minUsagesToOptimize: number = DEFAULT_MIN_USAGES_TO_OPTIMIZE,
+    allAvailableVariables?: VariableDefinitions,
+  ): Operation {
     assert(minUsagesToOptimize >= 1, `Expected 'minUsagesToOptimize' to be at least 1, but got ${minUsagesToOptimize}`)
     if (!fragments || fragments.isEmpty()) {
       return this;
@@ -1001,11 +1033,16 @@ export class Operation extends DirectiveTargetElement<Operation> {
       }
     }
 
-    return this.withUpdatedSelectionSetAndFragments(optimizedSelection, finalFragments ?? undefined);
+    return this.withUpdatedSelectionSetAndFragments(
+      optimizedSelection,
+      finalFragments ?? undefined,
+      allAvailableVariables,
+    );
   }
 
   generateQueryFragments(): Operation {
     const [minimizedSelectionSet, fragments] = this.selectionSet.minimizeSelectionSet();
+    
     return new Operation(
       this.schema(),
       this.rootKind,

query-planner-js/src/buildPlan.ts

@@ -25,7 +25,6 @@ import {
   Variable,
   VariableDefinition,
   VariableDefinitions,
-  VariableCollector,
   newDebugLogger,
   selectionOfElement,
   selectionSetOfElement,
@@ -64,6 +63,8 @@ import {
   isInputType,
   possibleRuntimeTypes,
   NamedType,
+  VariableCollector,
+  DEFAULT_MIN_USAGES_TO_OPTIMIZE,
 } from "@apollo/federation-internals";
 import {
   advanceSimultaneousPathsWithOperation,
@@ -1591,16 +1592,29 @@ class FetchGroup {
     if (this.generateQueryFragments) {
       operation = operation.generateQueryFragments();
     } else {
-      operation = operation.optimize(fragments?.forSubgraph(this.subgraphName, subgraphSchema));
+      operation = operation.optimize(
+        fragments?.forSubgraph(this.subgraphName, subgraphSchema),
+        DEFAULT_MIN_USAGES_TO_OPTIMIZE,
+        variableDefinitions,
+      );
     }
 
+    // collect all used variables in the selection and in used Fragments
+    const usedVariables = new Set(selection.usedVariables().map(v => v.name));
+    if (operation.fragments) {
+      for (const namedFragment of operation.fragments.definitions()) {
+        namedFragment.selectionSet.usedVariables().forEach(v => {
+          usedVariables.add(v.name);  
+        });
+      }
+    }
     const operationDocument = operationToDocument(operation);
     const fetchNode: FetchNode = {
       kind: 'Fetch',
       id: this.id,
       serviceName: this.subgraphName,
       requires: inputNodes ? trimSelectionNodes(inputNodes.selections) : undefined,
-      variableUsages: selection.usedVariables().map(v => v.name),
+      variableUsages: Array.from(usedVariables),
       operation: stripIgnoredCharacters(print(operationDocument)),
       operationKind: schemaRootKindToOperationKind(operation.rootKind),
       operationName: operation.name,