various setContext bug fixes (#3017)

Some setContext bug fixes
- needsJoinDirective() logic is incorrect. We need to make sure that we
add a join field when @fromDirective exists on the arguments, not the
definition
- Query plans were incorrect if type was entirely in one subgraph.
selectionIsFullyLocalFromAllVertices was calling
SelectionSet.canRebaseOn so we fixed to return false if the selection
contained a field with a contextual argument
- For top level queries, we don't want to have "... on Query" in the
rewrite path.
- Fixed up  selectionSetAsKeyRenamers() logic

.changeset/popular-chairs-happen.md

@@ -0,0 +1,11 @@
+---
+"apollo-federation-integration-testsuite": patch
+"@apollo/query-planner": patch
+"@apollo/query-graphs": patch
+"@apollo/composition": patch
+"@apollo/federation-internals": patch
+"@apollo/subgraph": patch
+"@apollo/gateway": patch
+---
+
+Various set context bugfixes

composition-js/src/merging/merge.ts

@@ -1661,11 +1661,11 @@ class Merger {
       return true;
     }
 
-    // if there is a @fromContext directive on one of the sources, we need a join__field
+    // if there is a @fromContext directive on one of the source's arguments, we need a join__field
     if (sources.some((s, idx) => {
       const fromContextDirective = this.subgraphs.values()[idx].metadata().fromContextDirective();
-      if (isFederationDirectiveDefinedInSchema(fromContextDirective)) {
-        return (s?.appliedDirectivesOf(fromContextDirective).length ?? 0) > 0;
+      if (s && isFederationDirectiveDefinedInSchema(fromContextDirective)) {
+        return s.kind === 'FieldDefinition' && s.arguments().some(arg => arg.appliedDirectivesOf(fromContextDirective).length > 0);
       }
       return false;
     })) {

internals-js/src/federation.ts

@@ -413,11 +413,12 @@ const validateFieldValueType = ({
     }
     const { element, selectionSet: childSelectionSet } = selection;
     assert(element.definition.type, 'Element type definition should exist');
-    const type = element.definition.type;
+    let type = element.definition.type;
+
     if (childSelectionSet) {
-      assert(isCompositeType(type), 'Child selection sets should only exist on composite types');
+      assert(isCompositeType(baseType(type)), 'Child selection sets should only exist on composite types');
       const { resolvedType } = validateFieldValueType({
-        currentType: type,
+        currentType: baseType(type) as CompositeType,
         selectionSet: childSelectionSet,
         errorCollector,
         metadata,

internals-js/src/operations.ts

@@ -368,10 +368,23 @@ export class Field<TArgs extends {[key: string]: any} = {[key: string]: any}> ex
     if (this.name === typenameFieldName) {
       return parentType.typenameField()?.type;
     }
-
-    return this.canRebaseOn(parentType)
+    
+    const returnType = this.canRebaseOn(parentType)
       ? parentType.field(this.name)?.type
       : undefined;
+
+    // If the field has an argument with fromContextDirective on it. We should not rebase it.
+    const fromContextDirective = federationMetadata(parentType.schema())?.fromContextDirective();
+    if (fromContextDirective && isFederationDirectiveDefinedInSchema(fromContextDirective)) {
+      const fieldInParent = parentType.field(this.name);
+      if (fieldInParent && fieldInParent.arguments()
+          .some(arg => arg.appliedDirectivesOf(fromContextDirective).length > 0 && (!this.args || this.args[arg.name] === undefined))
+        ) {
+        return undefined;
+      }  
+    }
+
+    return returnType;
   }
 
   hasDefer(): boolean {

query-graphs-js/src/graphPath.ts

@@ -653,6 +653,12 @@ export class GraphPath<TTrigger, RV extends Vertex = Vertex, TNullEdge extends n
     if (!enteringEdge) {
       return undefined;
     }
+
+    // TODO: Temporary fix to avoid optimization if context exists.
+    // permanent fix is described here: https://github.com/apollographql/federation/pull/3017#pullrequestreview-2083949094
+    if (this.graph.subgraphToArgs.size > 0) {
+      return undefined;
+    }
 
     // Usually, the starting subgraph in which we want to look for a direct path is the head of
     // `subgraphEnteringEdge`, that is, where we were just before coming to the current subgraph.