fix(federation/query-planning): Fix bug where subgraph jump key fields are fetched from the wrong node

.changesets/fix_sachin_fed_544.md

@@ -0,0 +1,5 @@
+### Fix query planning error where `@requires` subgraph jump fetches `@key` from the wrong subgraph ([PR #8016](https://github.com/apollographql/router/pull/8016))
+
+During query planning, a subgraph jump added due to a `@requires` field may sometimes try to collect the necessary `@key` fields from an upstream subgraph fetch as an optimization, but it wasn't properly checking whether that subgraph had those fields. This is now fixed, and previously could cause query planning errors with messages that look like "Cannot add selection of field `T.id` to selection set of parent type `T`".
+
+By [@sachindshinde](https://github.com/sachindshinde) in https://github.com/apollographql/router/pull/8016

apollo-federation/src/query_plan/fetch_dependency_graph.rs

@@ -4800,14 +4800,56 @@ fn create_post_requires_node(
             parent.parent_node_id,
             parent_path,
         )?;
+
+        // The `post_requires_node_id` node needs a key. This can come from `fetch_node_id` (and
+        // code in `GraphPath::can_satisfy_conditions()` guarantees such a locally satisfiable key
+        // exists in `fetch_node_id`), but it can also potentially come from
+        // `parent.parent_node_id`, 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.parent_node_id`, and we fall back to
+        // `fetch_node_id` if it doesn't.
+        let Some(key_condition) = dependency_graph
+            .federated_query_graph
+            .locally_satisfiable_key(query_graph_edge_id)?
+        else {
+            bail!(
+                "Due to @requires, validation should have required a key to be present for {}",
+                query_graph_edge_id.index()
+            );
+        };
+        let parent_fetch_node = dependency_graph.node_weight(parent.parent_node_id)?;
+        let type_at_path_for_parent = dependency_graph.type_at_path(
+            &parent_fetch_node.selection_set.selection_set.type_position,
+            &parent_fetch_node.selection_set.selection_set.schema,
+            &path_for_parent.path_in_node,
+        )?;
+        let (require_node_path, pre_require_node_id) = if !key_condition.can_rebase_on(
+            &type_at_path_for_parent,
+            &parent_fetch_node.selection_set.selection_set.schema,
+        )? {
+            // It's possible we didn't add `fetch_node_id` as a parent previously, so we do so here
+            // similarly to how `handle_conditions_tree()` specifies it.
+            dependency_graph.add_parent(
+                post_requires_node_id,
+                ParentRelation {
+                    parent_node_id: fetch_node_id,
+                    path_in_parent: Some(Arc::new(Default::default())),
+                },
+            );
+            (fetch_node_path, fetch_node_id)
+        } else {
+            (&path_for_parent, parent.parent_node_id)
+        };
+
         add_post_require_inputs(
             dependency_graph,
-            &path_for_parent,
+            require_node_path,
             &entity_type_schema,
             entity_type_position.clone(),
             query_graph_edge_id,
             context,
-            parent.parent_node_id,
+            pre_require_node_id,
             post_requires_node_id,
         )?;
         created_nodes.insert(post_requires_node_id);

apollo-federation/tests/query_plan/build_query_plan_tests/requires.rs

@@ -1929,3 +1929,108 @@ fn allows_post_requires_input_with_typename_on_interface_object_type() {
     "###
     );
 }
+
+#[test]
+fn avoids_selecting_inapplicable_key_from_parent_node() {
+    // 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).
+    let planner = planner!(
+        A: r#"
+          type Query {
+            t: T
+          }
+
+          type T @key(fields: "id1") {
+            id1: ID!
+          }
+        "#,
+        B: r#"
+          type T @key(fields: "id2") @key(fields: "id1") {
+            id1: ID!
+            id2: ID!
+            x: Int @external
+            req: Int @requires(fields: "x")
+          }
+        "#,
+        C: r#"
+          type T @key(fields: "id2") {
+            id2: ID!
+            x: Int
+          }
+        "#,
+    );
+    assert_plan!(
+        &planner,
+        r#"
+          {
+            t {
+              req
+            }
+          }
+        "#,
+
+        @r###"
+        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
+                  }
+                }
+              },
+            },
+          },
+        }
+      "###
+    );
+}

apollo-federation/tests/query_plan/supergraphs/avoids_selecting_inapplicable_key_from_parent_node.graphql

@@ -0,0 +1,76 @@
+# Composed from subgraphs with hash: 7744f0ec3e7d0ca0b7427d4d80f066ec5650bfcb
+schema
+  @link(url: "https://specs.apollo.dev/link/v1.0")
+  @link(url: "https://specs.apollo.dev/join/v0.5", for: EXECUTION)
+{
+  query: Query
+}
+
+directive @join__directive(graphs: [join__Graph!], name: String!, args: join__DirectiveArguments) repeatable on SCHEMA | OBJECT | INTERFACE | FIELD_DEFINITION
+
+directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
+
+directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String, contextArguments: [join__ContextArgument!]) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
+
+directive @join__graph(name: String!, url: String!) on ENUM_VALUE
+
+directive @join__implements(graph: join__Graph!, interface: String!) repeatable on OBJECT | INTERFACE
+
+directive @join__type(graph: join__Graph!, key: join__FieldSet, extension: Boolean! = false, resolvable: Boolean! = true, isInterfaceObject: Boolean! = false) repeatable on OBJECT | INTERFACE | UNION | ENUM | INPUT_OBJECT | SCALAR
+
+directive @join__unionMember(graph: join__Graph!, member: String!) repeatable on UNION
+
+directive @link(url: String, as: String, for: link__Purpose, import: [link__Import]) repeatable on SCHEMA
+
+input join__ContextArgument {
+  name: String!
+  type: String!
+  context: String!
+  selection: join__FieldValue!
+}
+
+scalar join__DirectiveArguments
+
+scalar join__FieldSet
+
+scalar join__FieldValue
+
+enum join__Graph {
+  A @join__graph(name: "A", url: "none")
+  B @join__graph(name: "B", url: "none")
+  C @join__graph(name: "C", url: "none")
+}
+
+scalar link__Import
+
+enum link__Purpose {
+  """
+  `SECURITY` features provide metadata necessary to securely resolve fields.
+  """
+  SECURITY
+
+  """
+  `EXECUTION` features provide metadata necessary for operation execution.
+  """
+  EXECUTION
+}
+
+type Query
+  @join__type(graph: A)
+  @join__type(graph: B)
+  @join__type(graph: C)
+{
+  t: T @join__field(graph: A)
+}
+
+type T
+  @join__type(graph: A, key: "id1")
+  @join__type(graph: B, key: "id2")
+  @join__type(graph: B, key: "id1")
+  @join__type(graph: C, key: "id2")
+{
+  id1: ID! @join__field(graph: A) @join__field(graph: B)
+  id2: ID! @join__field(graph: B) @join__field(graph: C)
+  x: Int @join__field(graph: B, external: true) @join__field(graph: C)
+  req: Int @join__field(graph: B, requires: "x")
+}