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

Reverts #3300 (which reverted #3293), as we're finishing up testing and the original PR should be fine to merge soon.

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 {