Repository: apollographql/router
Pull request #7047: support preserving key fields in connector response mapping
Author: login: lennyburdette
Created: 2025-03-19T17:11:13Z | Merged: 2025-03-20T18:40:48Z

the query planner doesn't ask for key fields when it makes a batch entity request (because it already has them). but we'll need key fields in order to match the response objects with the entity references and ensure that the batch result is consistent with the request

<!-- https://apollographql.atlassian.net/browse/CNN-633 -->

<!-- start metadata -->
---

**Checklist**

Complete the checklist (and note appropriate exceptions) before the PR is marked ready-for-review.

- [ ] Changes are compatible[^1]
- [ ] Documentation[^2] completed
- [ ] Performance impact assessed and acceptable
- Tests added and passing[^3]
    - [ ] Unit Tests
    - [ ] Integration Tests
    - [ ] Manual Tests

**Exceptions**

*Note any exceptions here*

**Notes**

[^1]: It may be appropriate to bring upcoming changes to the attention of other (impacted) groups. Please endeavour to do this before seeking PR approval. The mechanism for doing this will vary considerably, so use your judgement as to how and when to do this.
[^2]: Configuration is an important part of many changes. Where applicable please try to document configuration examples.
[^3]: Tick whichever testing boxes are applicable. If you are adding Manual Tests, please document the manual testing (extensively) in the Exceptions.



Diff:
diff --git a/apollo-federation/src/sources/connect/json_selection/selection_set.rs b/apollo-federation/src/sources/connect/json_selection/selection_set.rs
index f09335d9d9..f46d30d104 100644
--- a/apollo-federation/src/sources/connect/json_selection/selection_set.rs
+++ b/apollo-federation/src/sources/connect/json_selection/selection_set.rs
@@ -19,6 +19,7 @@
 use apollo_compiler::ExecutableDocument;
 use apollo_compiler::Node;
 use apollo_compiler::executable::Field;
+use apollo_compiler::executable::FieldSet;
 use apollo_compiler::executable::Selection;
 use apollo_compiler::executable::SelectionSet;
 use apollo_compiler::name;
@@ -38,14 +39,32 @@ use crate::sources::connect::json_selection::NamedSelection;
 
 impl JSONSelection {
     /// Apply a selection set to create a new [`JSONSelection`]
+    ///
+    /// Operations from the query planner will never contain key fields (because
+    /// it already has them from a previous fetch) but we might need those
+    /// fields for things like sorting entities in a batch. If the optional
+    /// `required_keys` is provided, we'll merge those fields into the selection
+    /// set before applying it to the JSONSelection.
     pub fn apply_selection_set(
         &self,
         document: &ExecutableDocument,
         selection_set: &SelectionSet,
+        required_keys: Option<&FieldSet>,
     ) -> Self {
+        let selection_set = match required_keys {
+            Some(keys) => {
+                let mut selection_set = selection_set.clone();
+                for selection in keys.selection_set.selections.iter() {
+                    selection_set.push(selection.clone());
+                }
+                selection_set
+            }
+            None => selection_set.clone(),
+        };
+
         match self {
-            Self::Named(sub) => Self::Named(sub.apply_selection_set(document, selection_set)),
-            Self::Path(path) => Self::Path(path.apply_selection_set(document, selection_set)),
+            Self::Named(sub) => Self::Named(sub.apply_selection_set(document, &selection_set)),
+            Self::Path(path) => Self::Path(path.apply_selection_set(document, &selection_set)),
         }
     }
 }
@@ -274,8 +293,11 @@ fn map_fields_by_name_impl<'a>(
 mod tests {
     use apollo_compiler::ExecutableDocument;
     use apollo_compiler::Schema;
+    use apollo_compiler::executable::FieldSet;
     use apollo_compiler::executable::SelectionSet;
+    use apollo_compiler::name;
     use apollo_compiler::validation::Valid;
+    use insta::assert_snapshot;
     use pretty_assertions::assert_eq;
 
     fn selection_set(schema: &Valid<Schema>, s: &str) -> (ExecutableDocument, SelectionSet) {
@@ -339,7 +361,7 @@ mod tests {
             "{ t { z: a, y: b, x: d, w: h v: k { u: l t: m } } }",
         );
 
-        let transformed = json.apply_selection_set(&document, &selection_set);
+        let transformed = json.apply_selection_set(&document, &selection_set, None);
         assert_eq!(
             transformed.to_string(),
             r###"$.result {
@@ -413,7 +435,7 @@ mod tests {
             "{ t { a b_alias c { e: e_alias h group { j } } path_to_f } }",
         );
 
-        let transformed = json_selection.apply_selection_set(&document, &selection_set);
+        let transformed = json_selection.apply_selection_set(&document, &selection_set, None);
         assert_eq!(
             transformed.to_string(),
             r###"$.result {
@@ -507,7 +529,7 @@ mod tests {
 
         let (document, selection_set) = selection_set(&schema, "{ t { a { b { renamed } } } }");
 
-        let transformed = json.apply_selection_set(&document, &selection_set);
+        let transformed = json.apply_selection_set(&document, &selection_set, None);
         assert_eq!(
             transformed.to_string(),
             r###"$.result {
@@ -575,7 +597,7 @@ mod tests {
         let (document, selection_set) =
             selection_set(&schema, "{ t { id __typename author { __typename id } } }");
 
-        let transformed = json.apply_selection_set(&document, &selection_set);
+        let transformed = json.apply_selection_set(&document, &selection_set, None);
         assert_eq!(
             transformed.to_string(),
             r###"$.result {
@@ -653,7 +675,7 @@ mod tests {
             }",
         );
 
-        let transformed = json.apply_selection_set(&document, &selection_set);
+        let transformed = json.apply_selection_set(&document, &selection_set, None);
         assert_eq!(
             transformed.to_string(),
             r###"reviews: result {
@@ -669,4 +691,53 @@ mod tests {
 }"###
         );
     }
+
+    #[test]
+    fn test_ensuring_key_fields() {
+        let json = super::JSONSelection::parse(
+            r###"
+            id
+            store { id }
+            name
+            price
+            "###,
+        )
+        .unwrap();
+
+        let schema = Schema::parse_and_validate(
+            r###"
+        type Query {
+            product: Product
+        }
+
+        type Product {
+            id: ID!
+            store: Store!
+            name: String!
+            price: String
+        }
+
+        type Store {
+          id: ID!
+        }
+        "###,
+            "./",
+        )
+        .unwrap();
+
+        let (document, selection_set) = selection_set(&schema, "{ product { name price } }");
+
+        let keys =
+            FieldSet::parse_and_validate(&schema, name!(Product), "id store { id }", "").unwrap();
+
+        let transformed = json.apply_selection_set(&document, &selection_set, Some(&keys));
+        assert_snapshot!(transformed.to_string(), @r"
+        id
+        store {
+          id
+        }
+        name
+        price
+        ");
+    }
 }
diff --git a/apollo-router/src/plugins/connectors/make_requests.rs b/apollo-router/src/plugins/connectors/make_requests.rs
index 481c5f76ed..ac74a15985 100644
--- a/apollo-router/src/plugins/connectors/make_requests.rs
+++ b/apollo-router/src/plugins/connectors/make_requests.rs
@@ -314,11 +314,11 @@ fn root_fields(
 
                 let response_key = ResponseKey::RootField {
                     name: response_name,
-                    selection: Arc::new(
-                        connector
-                            .selection
-                            .apply_selection_set(&request.operation, &field.selection_set),
-                    ),
+                    selection: Arc::new(connector.selection.apply_selection_set(
+                        &request.operation,
+                        &field.selection_set,
+                        None,
+                    )),
                     inputs: request_inputs,
                 };
 
@@ -373,11 +373,11 @@ fn entities_from_request(
 
     let (entities_field, _) = graphql_utils::get_entity_fields(&request.operation, op)?;
 
-    let selection = Arc::new(
-        connector
-            .selection
-            .apply_selection_set(&request.operation, &entities_field.selection_set),
-    );
+    let selection = Arc::new(connector.selection.apply_selection_set(
+        &request.operation,
+        &entities_field.selection_set,
+        None,
+    ));
 
     representations
         .as_array()
@@ -528,11 +528,11 @@ fn entities_with_fields_from_request(
         .into_iter()
         .flatten()
         .flat_map(|(typename, field)| {
-            let selection = Arc::new(
-                connector
-                    .selection
-                    .apply_selection_set(&request.operation, &field.selection_set),
-            );
+            let selection = Arc::new(connector.selection.apply_selection_set(
+                &request.operation,
+                &field.selection_set,
+                None,
+            ));
 
             representations.iter().map(move |(i, representation)| {
                 let args = graphql_utils::field_arguments_map(field, &request.variables.variables)


Reviews:
user: login: benjamn
body: Looks good. Is there any way to use this outside of the tests, yet?
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-19T18:17:11Z
user: login: lennyburdette
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T18:21:28Z

Review comments:
diff_hunk: @@ -669,4 +691,53 @@ mod tests {
 }"###
         );
     }
+
+    #[test]
+    fn test_ensuring_key_fields() {
+        let json = super::JSONSelection::parse(
+            r###"
+            id
+            store { id }
+            name
+            price
+            "###,
+        )
+        .unwrap();
+
+        let schema = Schema::parse_and_validate(
+            r###"
+        type Query {
+            product: Product
+        }
+
+        type Product {
+            id: ID!
+            store: Store!
+            name: String!
+            price: String
+        }
+
+        type Store {
+          id: ID!
+        }
+        "###,
+            "./",
+        )
+        .unwrap();
+
+        let (document, selection_set) = selection_set(&schema, "{ product { name price } }");
+
+        let keys =
+            FieldSet::parse_and_validate(&schema, name!(Product), "id store { id }", "").unwrap();
path: apollo-federation/src/sources/connect/json_selection/selection_set.rs
user: login: benjamn
body: We still need a general way of extracting these key `FieldSet`s from a `JSONSelection`, right?
created_at: 2025-03-19T18:16:10Z
author_association: MEMBER
diff_hunk: @@ -669,4 +691,53 @@ mod tests {
 }"###
         );
     }
+
+    #[test]
+    fn test_ensuring_key_fields() {
+        let json = super::JSONSelection::parse(
+            r###"
+            id
+            store { id }
+            name
+            price
+            "###,
+        )
+        .unwrap();
+
+        let schema = Schema::parse_and_validate(
+            r###"
+        type Query {
+            product: Product
+        }
+
+        type Product {
+            id: ID!
+            store: Store!
+            name: String!
+            price: String
+        }
+
+        type Store {
+          id: ID!
+        }
+        "###,
+            "./",
+        )
+        .unwrap();
+
+        let (document, selection_set) = selection_set(&schema, "{ product { name price } }");
+
+        let keys =
+            FieldSet::parse_and_validate(&schema, name!(Product), "id store { id }", "").unwrap();
path: apollo-federation/src/sources/connect/json_selection/selection_set.rs
user: login: lennyburdette
body: yeah, this currently only works with something like `{$this.id}/{$this.store.id}` so we still need a way to support `$this { id store { id } }`
created_at: 2025-03-19T18:21:28Z
author_association: CONTRIBUTOR

PR comments:
user: login: svc-apollo-docs
created_at: 2025-03-19T17:11:17Z
body: ### ⚠️ Docs preview not attached to branch

The preview was not built because the PR's base branch `lb/expand-batch` is not in the list of sources.

An Apollo team member can comment one of the following commands to dictate which branch to attach the preview to:

- `!docs set-base-branch dev`
- `!docs set-base-branch 1.x`

**Build ID**: 6b27aeed82a241abc73aa792
author_association: NONE
user: login: lennyburdette
created_at: 2025-03-19T18:21:48Z
body: i'll use it in the next PR!
author_association: CONTRIBUTOR