fix(query-planner): @requires fetch group optimization selects invalid key

.changeset/pink-pans-stare.md

@@ -0,0 +1,5 @@
+---
+"@apollo/query-planner": patch
+---
+
+Fixed a bug where, when composing fetch groups, the query planner would select a key field on a parent group, without checking if the field is defined in the respective schema. This happened particularly when optimizing fetch groups that satisfy @requires conditions. With these changes, the query planner will properly resolve a selectable key from the parent group's schema.

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

@@ -3162,6 +3162,157 @@ describe('@requires', () => {
       }
     `);
   });
+
+  it('selects a valid key from parent fetch group (for another subgraph) when handling @requires', () => {
+    /**
+     * Previously, there was an issue where, when recognizing that we can select
+     * a key on `T` at `t.y.t2` in the fetch group C-y to be used as input for
+     * fetch group A-t::y::t2, we would search A's schema for a key field on
+     * type `T` to select, assuming that `T`'s definition in C's schema has the
+     * same key(s).
+     */
+    const subgraph1 = {
+      name: 'A',
+      typeDefs: gql`
+        type Query {
+          t: T
+        }
+
+        type T @key(fields: "id1") @key(fields: "id2") {
+          id1: ID!
+          id2: ID!
+          x: Int @external
+          req: Int @requires(fields: "x")
+        }
+      `,
+    };
+
+    const subgraph2 = {
+      name: 'B',
+      typeDefs: gql`
+        type T @key(fields: "id1") {
+          id1: ID!
+          x: Int
+        }
+      `,
+    };
+
+    const subgraph3 = {
+      name: 'C',
+      typeDefs: gql`
+        type Y {
+          t2: T
+        }
+
+        type T @key(fields: "id2") {
+          id2: ID!
+          y: Y
+        }
+      `,
+    };
+
+    const [api, queryPlanner] = composeAndCreatePlanner(
+      subgraph1,
+      subgraph2,
+      subgraph3,
+    );
+    const operation = operationFromDocument(
+      api,
+      gql`
+        {
+          t {
+            y {
+              t2 {
+                req
+              }
+            }
+          }
+        }
+      `,
+    );
+
+    const plan = queryPlanner.buildQueryPlan(operation);
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Sequence {
+          Fetch(service: "A") {
+            {
+              t {
+                __typename
+                id2
+              }
+            }
+          },
+          Flatten(path: "t") {
+            Fetch(service: "C") {
+              {
+                ... on T {
+                  __typename
+                  id2
+                }
+              } =>
+              {
+                ... on T {
+                  y {
+                    t2 {
+                      __typename
+                      id2
+                    }
+                  }
+                }
+              }
+            },
+          },
+          Flatten(path: "t.y.t2") {
+            Fetch(service: "A") {
+              {
+                ... on T {
+                  __typename
+                  id2
+                }
+              } =>
+              {
+                ... on T {
+                  id1
+                }
+              }
+            },
+          },
+          Flatten(path: "t.y.t2") {
+            Fetch(service: "B") {
+              {
+                ... on T {
+                  __typename
+                  id1
+                }
+              } =>
+              {
+                ... on T {
+                  x
+                }
+              }
+            },
+          },
+          Flatten(path: "t.y.t2") {
+            Fetch(service: "A") {
+              {
+                ... on T {
+                  __typename
+                  x
+                  id2
+                }
+              } =>
+              {
+                ... on T {
+                  req
+                }
+              }
+            },
+          },
+        },
+      }
+    `);
+  });
 });
 
 describe('fetch operation names', () => {

query-planner-js/src/buildPlan.ts

@@ -65,6 +65,7 @@
   NamedType,
   VariableCollector,
   DEFAULT_MIN_USAGES_TO_OPTIMIZE,
+  parseFieldSetArgument,
 } from "@apollo/federation-internals";
 import {
   advanceSimultaneousPathsWithOperation,
@@ -5063,7 +5064,26 @@
   preRequireGroup: FetchGroup,
   postRequireGroup: FetchGroup,
 ) {
-  const { inputs, keyInputs } = inputsForRequire(dependencyGraph, entityType, edge, context);
+  const keyConditionIsCrossSubgraph =
+    edge.head.source !== preRequireGroup.subgraphName;
+
+  const crossSubgraphKeyCondition = keyConditionIsCrossSubgraph
+    ? resolveKeyConditionFromParentGroup(
+      dependencyGraph,
+      edge,
+      preRequireGroup.subgraphName,
+    )
+    : undefined;
+
+  const { inputs, keyInputs } = inputsForRequire(
+    dependencyGraph,
+    entityType,
+    edge,
+    context,
+    true,
+    crossSubgraphKeyCondition,
+  );
+
   // Note that `computeInputRewritesOnKeyFetch` will return `undefined` in general, but if `entityType` is an interface/interface object,
   // then we need those rewrites to ensure the underlying fetch is valid.
   postRequireGroup.addInputs(
@@ -5077,6 +5097,75 @@
   }
 }
 
+function resolveKeyConditionFromParentGroup(
+  dependencyGraph: FetchDependencyGraph,
+  edge: Edge,
+  parentSource: string,
+): SelectionSet {
+  const parentSchema = dependencyGraph.federatedQueryGraph.sources
+    .get(parentSource);
+  const parentSchemaMetadata = parentSchema && federationMetadata(parentSchema);
+  assert(
+    parentSchemaMetadata,
+    () => `Could not find federation metadata for source ${parentSource}`,
+  );
+
+  const inputTypeSchema = dependencyGraph.federatedQueryGraph.sources
+    .get(edge.head.source);
+  const inputTypeSchemaMetadata = inputTypeSchema
+    && federationMetadata(inputTypeSchema);
+  assert(
+    inputTypeSchemaMetadata,
+    () => `Could not find federation metadata for source ${edge.head.source}`,
+  );
+
+  const isInterfaceObjectDownCast =
+    edge.transition.kind === 'InterfaceObjectFakeDownCast';
+  const inputType = edge.head.type as CompositeType;
+
+  let typeToRebaseOn: CompositeType;
+  if (isInterfaceObjectDownCast || isInterfaceObjectType(inputType)) {
+    const supergraphItfType = dependencyGraph.supergraphSchema
+      .type(inputType.name);
+    assert(
+      supergraphItfType && isInterfaceType(supergraphItfType),
+      () => `Type ${inputType.name} should be an interface in the supergraph`,
+    );
+    typeToRebaseOn = supergraphItfType;
+  } else {
+    const parentType = parentSchema.type(inputType.name);
+    assert(
+      parentType && isObjectType(parentType),
+      () => `Type ${inputType.name} should be in schema for source ${parentSource}`,
+    );
+    typeToRebaseOn = parentType;
+  }
+
+  const inputTypeSchemaKeyDirective = inputTypeSchemaMetadata.keyDirective();
+  let keyFieldSelectionSet: SelectionSet | undefined;
+  for (const key of inputType.appliedDirectivesOf(inputTypeSchemaKeyDirective)) {
+    if (!(key.arguments().resolvable ?? true)) continue;
+
+    const selectionSet = parseFieldSetArgument({
+      parentType: inputType,
+      directive: key,
+    });
+    if (!inputTypeSchemaMetadata.selectionSelectsAnyExternalField(selectionSet)
+      && selectionSet.canRebaseOn(typeToRebaseOn)
+    ) {
+      keyFieldSelectionSet = selectionSet;
+      break;
+    }
+  }
+
+  assert(
+    keyFieldSelectionSet,
+    () => `Due to @requires, validation should have required a key to be present for ${edge}`,
+  );
+
+  return keyFieldSelectionSet;
+}
+
 function newCompositeTypeSelectionSet(type: CompositeType): MutableSelectionSet {
   const selectionSet = MutableSelectionSet.empty(type);
   selectionSet.updates().add(new FieldSelection(new Field(type.typenameField()!)));
@@ -5088,7 +5177,8 @@
   entityType: ObjectType,
   edge: Edge,
   context: PathContext,
-  includeKeyInputs: boolean = true
+  includeKeyInputs: boolean = true,
+  keyConditionOverride?: SelectionSet,
 ): {
   inputs: SelectionSet,
   keyInputs: SelectionSet | undefined,
@@ -5109,7 +5199,10 @@
   }
   let keyInputs: MutableSelectionSet | undefined = undefined;
   if (includeKeyInputs) {
-    const keyCondition = getLocallySatisfiableKey(dependencyGraph.federatedQueryGraph, edge.head);
+    const keyCondition = keyConditionOverride ?? getLocallySatisfiableKey(
+      dependencyGraph.federatedQueryGraph,
+      edge.head,
+    );
     assert(keyCondition, () => `Due to @require, validation should have required a key to be present for ${edge}`);
     let keyConditionAsInput = keyCondition;
     if (isInterfaceObjectDownCast) {