Fixed query planner to pass the directives from original query to subgraph operations (#2961)

.changeset/six-shoes-fly.md

@@ -0,0 +1,6 @@
+---
+"@apollo/query-planner": patch
+"@apollo/federation-internals": patch
+---
+
+Fixed query planner to pass the directives from original query to subgraph operations (#2961)

internals-js/src/definitions.ts

@@ -392,26 +392,11 @@ export class DirectiveTargetElement<T extends DirectiveTargetElement<T>> {
   }
 
   appliedDirectivesToDirectiveNodes() : ConstDirectiveNode[] | undefined {
-    if (this.appliedDirectives.length == 0) {
-      return undefined;
-    }
-
-    return this.appliedDirectives.map(directive => {
-      return {
-        kind: Kind.DIRECTIVE,
-        name: {
-          kind: Kind.NAME,
-          value: directive.name,
-        },
-        arguments: directive.argumentsToAST()
-      };
-    });
+    return directivesToDirectiveNodes(this.appliedDirectives);
   }
 
   appliedDirectivesToString(): string {
-    return this.appliedDirectives.length == 0
-      ? ''
-      : ' ' + this.appliedDirectives.join(' ');
+    return directivesToString(this.appliedDirectives);
   }
 
   collectVariablesInAppliedDirectives(collector: VariableCollector) {
@@ -3257,6 +3242,37 @@ export class Directive<
   }
 }
 
+/**
+ * Formats a Directive array as a string (with a leading space, if present).
+ */
+export function directivesToString(directives?: readonly Directive<any>[])
+  : string
+{
+  return (!directives || directives.length == 0)
+        ? ''
+        : ' ' + directives.join(' ');
+}
+
+/**
+ * Converts a Directive array into DirectiveNode array.
+ */
+export function directivesToDirectiveNodes(directives?: readonly Directive<any>[])
+  : ConstDirectiveNode[] | undefined
+{
+  return (!directives || directives.length === 0)
+    ? undefined
+    : directives.map(directive => {
+      return {
+        kind: Kind.DIRECTIVE,
+        name: {
+          kind: Kind.NAME,
+          value: directive.name,
+        },
+        arguments: directive.argumentsToAST()
+      };
+    });
+}
+
 /**
  * Checks if 2 directive applications should be considered equal.
  *

internals-js/src/operations.ts

@@ -48,6 +48,8 @@ import {
   isObjectType,
   NamedType,
   isUnionType,
+  directivesToString,
+  directivesToDirectiveNodes,
 } from "./definitions";
 import { isInterfaceObjectType } from "./federation";
 import { ERRORS } from "./error";
@@ -877,15 +879,15 @@ function computeFragmentsToKeep(
   return toExpand.size === 0 ? fragments : fragments.filter((f) => !toExpand.has(f.name));
 }
 
-// TODO Operations can also have directives
 export class Operation {
   constructor(
     readonly schema: Schema,
     readonly rootKind: SchemaRootKind,
     readonly selectionSet: SelectionSet,
     readonly variableDefinitions: VariableDefinitions,
     readonly fragments?: NamedFragments,
-    readonly name?: string) {
+    readonly name?: string,
+    readonly directives?: readonly Directive<any>[]) {
   }
 
   // Returns a copy of this operation with the provided updated selection set.
@@ -901,7 +903,8 @@ export class Operation {
       newSelectionSet,
       this.variableDefinitions,
       this.fragments,
-      this.name
+      this.name,
+      this.directives
     );
   }
 
@@ -917,7 +920,8 @@ export class Operation {
       newSelectionSet,
       this.variableDefinitions,
       newFragments,
-      this.name
+      this.name,
+      this.directives
     );
   }
 
@@ -982,6 +986,7 @@ export class Operation {
       this.variableDefinitions,
       fragments,
       this.name,
+      this.directives
     );
   }
 
@@ -1053,7 +1058,7 @@ export class Operation {
   }
 
   toString(expandFragments: boolean = false, prettyPrint: boolean = true): string {
-    return this.selectionSet.toOperationString(this.rootKind, this.variableDefinitions, this.fragments, this.name, expandFragments, prettyPrint);
+    return this.selectionSet.toOperationString(this.rootKind, this.variableDefinitions, this.fragments, this.name, this.directives, expandFragments, prettyPrint);
   }
 }
 
@@ -2077,6 +2082,7 @@ export class SelectionSet {
     variableDefinitions: VariableDefinitions,
     fragments: NamedFragments | undefined,
     operationName?: string,
+    directives?: readonly Directive<any>[],
     expandFragments: boolean = false,
     prettyPrint: boolean = true
   ): string {
@@ -2090,7 +2096,8 @@ export class SelectionSet {
     const nameAndVariables = operationName
       ? " " + (operationName + (variableDefinitions.isEmpty() ? "" : variableDefinitions.toString()))
       : (variableDefinitions.isEmpty() ? "" : " " + variableDefinitions.toString());
-    return fragmentsDefinitions + rootKind + nameAndVariables + " " + this.toString(expandFragments, true, indent);
+    const directives_str = directivesToString(directives);
+    return fragmentsDefinitions + rootKind + nameAndVariables + directives_str + " " + this.toString(expandFragments, true, indent);
   }
 
   /**
@@ -3562,7 +3569,7 @@ class FragmentSpreadSelection extends FragmentSelection {
 
   key(): string {
     if (!this.computedKey) {
-      this.computedKey = '...' + this.namedFragment.name + (this.spreadDirectives.length === 0 ? '' : ' ' + this.spreadDirectives.join(' '));
+      this.computedKey = '...' + this.namedFragment.name + directivesToString(this.spreadDirectives);
     }
     return this.computedKey;
   }
@@ -3588,18 +3595,7 @@ class FragmentSpreadSelection extends FragmentSelection {
   }
 
   toSelectionNode(): FragmentSpreadNode {
-    const directiveNodes = this.spreadDirectives.length === 0
-      ? undefined
-      : this.spreadDirectives.map(directive => {
-        return {
-          kind: Kind.DIRECTIVE,
-          name: {
-            kind: Kind.NAME,
-            value: directive.name,
-          },
-          arguments: directive.argumentsToAST()
-        } as DirectiveNode;
-      });
+    const directiveNodes = directivesToDirectiveNodes(this.spreadDirectives);
     return {
       kind: Kind.FRAGMENT_SPREAD,
       name: { kind: Kind.NAME, value: this.namedFragment.name },
@@ -3744,9 +3740,7 @@ class FragmentSpreadSelection extends FragmentSelection {
     if (expandFragments) {
       return (indent ?? '') + this.element + ' ' + this.selectionSet.toString(true, true, indent);
     } else {
-      const directives = this.spreadDirectives;
-      const directiveString = directives.length == 0 ? '' : ' ' + directives.join(' ');
-      return (indent ?? '') + '...' + this.namedFragment.name + directiveString;
+      return (indent ?? '') + '...' + this.namedFragment.name + directivesToString(this.spreadDirectives);
     }
   }
 }
@@ -3832,6 +3826,7 @@ export function operationFromDocument(
   }
 ) : Operation {
   let operation: OperationDefinitionNode | undefined;
+  let operation_directives: Directive<any>[] | undefined; // the directives on `operation`
   const operationName = options?.operationName;
   const fragments = new NamedFragments();
   // We do a first pass to collect the operation, and create all named fragment, but without their selection set yet.
@@ -3842,6 +3837,7 @@ export function operationFromDocument(
         validate(!operation || operationName, () => 'Must provide operation name if query contains multiple operations.');
         if (!operationName || (definition.name && definition.name.value === operationName)) {
           operation = definition;
+          operation_directives = directivesOfNodes(schema, definition.directives);
         }
         break;
       case Kind.FRAGMENT_DEFINITION:
@@ -3875,18 +3871,20 @@ export function operationFromDocument(
     }
   });
   fragments.validate(variableDefinitions);
-  return operationFromAST({schema, operation, variableDefinitions, fragments, validateInput: options?.validate});
+  return operationFromAST({schema, operation, operation_directives, variableDefinitions, fragments, validateInput: options?.validate});
 }
 
 function operationFromAST({
   schema,
   operation,
+  operation_directives,
   variableDefinitions,
   fragments,
   validateInput,
 }:{
   schema: Schema,
   operation: OperationDefinitionNode,
+  operation_directives?: Directive<any>[],
   variableDefinitions: VariableDefinitions,
   fragments: NamedFragments,
   validateInput?: boolean,
@@ -3906,7 +3904,8 @@ function operationFromAST({
     }),
     variableDefinitions,
     fragmentsIfAny,
-    operation.name?.value
+    operation.name?.value,
+    operation_directives
   );
 }
 
@@ -3961,6 +3960,7 @@ export function operationToDocument(operation: Operation): DocumentNode {
     name: operation.name ? { kind: Kind.NAME, value: operation.name } : undefined,
     selectionSet: operation.selectionSet.toSelectionSetNode(),
     variableDefinitions: operation.variableDefinitions.toVariableDefinitionNodes(),
+    directives: directivesToDirectiveNodes(operation.directives),
   };
   const fragmentASTs: DefinitionNode[] = operation.fragments
     ? operation.fragments?.toFragmentDefinitionNodes()

query-planner-js/src/QueryPlan.ts

@@ -46,7 +46,7 @@ export interface FetchNode {
   operationKind: OperationTypeNode;
   operationDocumentNode?: DocumentNode;
   // Optionally describe a number of "rewrites" that query plan executors should apply to the data that is sent as input of this fetch.
-  // Note that such rewrites should only impact the inputs of the fetch they are applied to (meaning that, as those inputs are collected 
+  // Note that such rewrites should only impact the inputs of the fetch they are applied to (meaning that, as those inputs are collected
   // from the current in-memory result, the rewrite should _not_ impact said in-memory results, only what is sent in the fetch).
   inputRewrites?: FetchDataRewrite[];
   // Similar, but for optional "rewrites" to apply to the data that received from a fetch (and before it is applied to the current in-memory results).