fix(query-planner): Fix bug where subgraph jump key fields are fetched from the wrong group (#3293)

When a `@requires` edge needs a subgraph jump, this would sometimes cause the locally satisfiable key fields needed by that jump to be fetched from the parent group instead of the current group (which apparently was done as an optimization). This can be a problem, since `canSatisfyConditions()` only really guarantees the current group has those key fields.

This PR updates the code using the parent group to now explicitly check whether the locally satisfiable key can be fetched from that group, and if not, falls back to the current group (adding it as a parent, if needed). This lets us keep the optimization and avoid unnecessarily changing query plans.

Fixes #3113 and #2987 , and derives a simpler example test from the one given in #3246 .

Author: sachindshinde

.changeset/afraid-singers-arrive.md

@@ -0,0 +1,5 @@
+---
+"@apollo/query-planner": patch
+---
+
+Fix bug in query planning where a subgraph jump for `@requires` can sometimes try to fetch `@key` fields from a subgraph that doesn't have them. This bug would previously cause query planning to error with a message that looks like "Cannot add selection of field `T.id` to selection set of parent type `T`".

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

@@ -3162,6 +3162,126 @@ describe('@requires', () => {
       }
     `);
   });
+
+  it('avoids selecting inapplicable @key from parent fetch group', () => {
+    // Previously, there was an issue where a subgraph jump inserted due to a
+    // @requires field would try (as an optimization) to collect its locally
+    // satisfiable key from the parent, but the parent may not have that locally
+    // satisfiable key. We now explicitly check for this, falling back to the
+    // current node if needed (which should be guaranteed to have it).
+    const subgraph1 = {
+      name: 'A',
+      typeDefs: gql`
+        type Query {
+          t: T
+        }
+        type T @key(fields: "id1") {
+          id1: ID!
+        }
+      `,
+    };
+
+    const subgraph2 = {
+      name: 'B',
+      typeDefs: gql`
+        type T @key(fields: "id2") @key(fields: "id1") {
+          id1: ID!
+          id2: ID!
+          x: Int @external
+          req: Int @requires(fields: "x")
+        }
+      `,
+    };
+
+    const subgraph3 = {
+      name: 'C',
+      typeDefs: gql`
+        type T @key(fields: "id2") {
+          id2: ID!
+          x: Int
+        }
+      `,
+    };
+
+    const [api, queryPlanner] = composeAndCreatePlanner(
+      subgraph1,
+      subgraph2,
+      subgraph3,
+    );
+    const operation = operationFromDocument(
+      api,
+      gql`
+        {
+          t {
+            req
+          }
+        }
+      `,
+    );
+
+    const plan = queryPlanner.buildQueryPlan(operation);
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Sequence {
+          Fetch(service: "A") {
+            {
+              t {
+                __typename
+                id1
+              }
+            }
+          },
+          Flatten(path: "t") {
+            Fetch(service: "B") {
+              {
+                ... on T {
+                  __typename
+                  id1
+                }
+              } =>
+              {
+                ... on T {
+                  __typename
+                  id2
+                }
+              }
+            },
+          },
+          Flatten(path: "t") {
+            Fetch(service: "C") {
+              {
+                ... on T {
+                  __typename
+                  id2
+                }
+              } =>
+              {
+                ... on T {
+                  x
+                }
+              }
+            },
+          },
+          Flatten(path: "t") {
+            Fetch(service: "B") {
+              {
+                ... on T {
+                  __typename
+                  x
+                  id2
+                }
+              } =>
+              {
+                ... on T {
+                  req
+                }
+              }
+            },
+          },
+        },
+      }
+    `);
+  });
 });
 
 describe('fetch operation names', () => {

query-planner-js/src/buildPlan.ts

@@ -4814,13 +4814,34 @@ function createPostRequiresGroup(
     // we need the path here, so this will have to do for now, and if this ever breaks in practice, we'll at least have an example to
     // guide us toward improving/fixing.
     assert(parent.path, `Missing path-in-parent for @requires on ${edge} with group ${group} and parent ${parent}`);
+    let requirePath = path.forParentOfGroup(parent.path, parent.group.parentType.schema());
+    let preRequireGroup = parent.group;
+
+    // The `postRequireGroup` needs a key. This can come from `group` (and code
+    // in `canSatisfyConditions()` guarantees such a locally satisfiable key
+    // exists in `group`), but it can also potentially come from `parent.group`,
+    // and previous code had (wrongfully) always assumed it could.
+    //
+    // To keep this previous optimization, we now explicitly check whether the
+    // known locally satisfiable key can be rebased in `parent.group`, and we
+    // fall back to `group` if it doesn't.
+    const keyCondition = getLocallySatisfiableKey(dependencyGraph.federatedQueryGraph, edge.head);
+    assert(keyCondition, () => `Due to @requires, validation should have required a key to be present for ${edge}`);
+    if (!keyCondition.canRebaseOn(typeAtPath(preRequireGroup.selection.parentType, requirePath.inGroup()))) {
+      requirePath = path;
+      preRequireGroup = group;
+      // It's possible we didn't add `group` as a parent previously, so we do so
+      // here similarly to how `handleRequiresTree()` specifies it.
+      postRequireGroup.addParent({ group, path: [] });
+    }
+
     addPostRequireInputs(
       dependencyGraph,
-      path.forParentOfGroup(parent.path, parent.group.parentType.schema()),
+      requirePath,
       entityType,
       edge,
       context,
-      parent.group,
+      preRequireGroup,
       postRequireGroup,
     );
     return {