reconstruct deferred queries with knowledge about fragments

NEXT_CHANGELOG.md

@@ -118,6 +118,12 @@ When we drop Telemetry we spawn a thread to perform the global opentelemetry tra
 
 By [@garypen](https://github.com/garypen) in https://github.com/apollographql/router/pull/2191
 
+### Reconstruct deferred queries with knowledge about fragments ([Issue #2105](https://github.com/apollographql/router/issues/2105))
+
+When we are using `@defer`, response formatting must apply on a subset of the query (primary or deferred), that is reconstructed from information provided by the query planner: a path into the response and a subselection. Previously, that path did not include information on fragment application, which resulted in query reconstruction issues if `@defer` was used under a fragment application on an interface.
+
+By [@Geal](https://github.com/geal) in https://github.com/apollographql/router/pull/2109
+
 ## 🛠 Maintenance
 
 ### improve plugin registration predictability ([PR #2181](https://github.com/apollographql/router/pull/2181))

apollo-router/Cargo.toml

@@ -150,7 +150,7 @@ reqwest = { version = "0.11.13", default-features = false, features = [
     "json",
     "stream",
 ] }
-router-bridge = "0.1.11"
+router-bridge = "0.1.12"
 rust-embed="6.4.2"
 schemars = { version = "0.8.11", features = ["url"] }
 shellexpand = "2.1.2"
@@ -190,6 +190,7 @@ urlencoding = "2.1.2"
 uuid = { version = "1.2.2", features = ["serde", "v4"] }
 yaml-rust = "0.4.5"
 askama = "0.11.1"
+apollo-encoder = "0.4.0"
 
 [target.'cfg(macos)'.dependencies]
 uname = "0.1.1"

apollo-router/src/axum_factory/snapshots/apollo_router__axum_factory__tests__defer_is_not_buffered.snap

@@ -0,0 +1,38 @@
+---
+source: apollo-router/src/axum_factory/tests.rs
+expression: "json!([{\n        \"data\" :\n        {\n            \"topProducts\" :\n            [{ \"upc\" : \"1\", \"name\" : \"Table\", \"reviews\" : null },\n            { \"upc\" : \"2\", \"name\" : \"Couch\", \"reviews\" : null }]\n        }, \"errors\" :\n        [{\n            \"message\" :\n            \"couldn't find mock for query {\\\"query\\\":\\\"query TopProducts__reviews__1($representations:[_Any!]!){_entities(representations:$representations){...on Product{reviews{__typename id product{__typename upc}}}}}\\\",\\\"operationName\\\":\\\"TopProducts__reviews__1\\\",\\\"variables\\\":{\\\"representations\\\":[{\\\"__typename\\\":\\\"Product\\\",\\\"upc\\\":\\\"1\\\"},{\\\"__typename\\\":\\\"Product\\\",\\\"upc\\\":\\\"2\\\"}]}}\"\n        },\n        {\n            \"message\" :\n            \"Subgraph response from 'reviews' was missing key `_entities`\",\n            \"path\" : [\"topProducts\", \"@\"]\n        }], \"hasNext\" : true,\n    }, { \"hasNext\" : false }])"
+---
+[
+  {
+    "data": {
+      "topProducts": [
+        {
+          "upc": "1",
+          "name": "Table",
+          "reviews": null
+        },
+        {
+          "upc": "2",
+          "name": "Couch",
+          "reviews": null
+        }
+      ]
+    },
+    "errors": [
+      {
+        "message": "couldn't find mock for query {\"query\":\"query TopProducts__reviews__1($representations:[_Any!]!){_entities(representations:$representations){...on Product{reviews{__typename id product{__typename upc}}}}}\",\"operationName\":\"TopProducts__reviews__1\",\"variables\":{\"representations\":[{\"__typename\":\"Product\",\"upc\":\"1\"},{\"__typename\":\"Product\",\"upc\":\"2\"}]}}"
+      },
+      {
+        "message": "Subgraph response from 'reviews' was missing key `_entities`",
+        "path": [
+          "topProducts",
+          "@"
+        ]
+      }
+    ],
+    "hasNext": true
+  },
+  {
+    "hasNext": false
+  }
+]

apollo-router/src/axum_factory/tests.rs

@@ -1923,31 +1923,7 @@ async fn test_defer_is_not_buffered() {
 
     let (parts, counts): (Vec<_>, Vec<_>) = parts.map(|part| (part, counter.get())).unzip().await;
     let parts = serde_json::Value::Array(parts);
-    assert_eq!(
-        parts,
-        json!([
-            {
-                "data": {
-                    "topProducts": [
-                        {"upc": "1", "name": "Table", "reviews": null},
-                        {"upc": "2", "name": "Couch", "reviews": null}
-                    ]
-                },
-                "errors": [
-                    {
-                        "message": "couldn't find mock for query {\"query\":\"query TopProducts__reviews__1($representations:[_Any!]!){_entities(representations:$representations){...on Product{reviews{__typename id product{__typename upc}}}}}\",\"operationName\":\"TopProducts__reviews__1\",\"variables\":{\"representations\":[{\"__typename\":\"Product\",\"upc\":\"1\"},{\"__typename\":\"Product\",\"upc\":\"2\"}]}}"
-                    },
-                    {
-                        "message": "Subgraph response from 'reviews' was missing key `_entities`",
-                        "path": [ "topProducts", "@" ]
-                    }],
-                "hasNext": true,
-            },
-            {"hasNext": false}
-        ]),
-        "{}",
-        serde_json::to_string(&parts).unwrap()
-    );
+    insta::assert_json_snapshot!(parts);
 
     // Non-regression test for https://github.com/apollographql/router/issues/1572
     //

apollo-router/src/json_ext.rs

@@ -242,6 +242,7 @@ impl ValueExt for Value {
                         .get_mut(k.as_str())
                         .expect("the value at that key was just inserted");
                 }
+                PathElement::Fragment(_) => {}
             }
         }
 
@@ -319,6 +320,7 @@ impl ValueExt for Value {
                         })
                     }
                 },
+                PathElement::Fragment(_) => {}
             }
         }
 
@@ -402,8 +404,17 @@ where
                     iterate_path(parent, &path[1..], value, f);
                     parent.pop();
                 }
+            } else if let Value::Array(array) = data {
+                for (i, value) in array.iter().enumerate() {
+                    parent.push(PathElement::Index(i));
+                    iterate_path(parent, path, value, f);
+                    parent.pop();
+                }
             }
         }
+        Some(PathElement::Fragment(_)) => {
+            iterate_path(parent, &path[1..], data, f);
+        }
     }
 }
 
@@ -422,6 +433,10 @@ pub enum PathElement {
     /// An index path element.
     Index(usize),
 
+    /// A fragment application
+    #[serde(deserialize_with = "deserialize_fragment")]
+    Fragment(String),
+
     /// A key path element.
     Key(String),
 }
@@ -464,6 +479,37 @@ where
     serializer.serialize_str("@")
 }
 
+fn deserialize_fragment<'de, D>(deserializer: D) -> Result<String, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    deserializer.deserialize_str(FragmentVisitor)
+}
+
+struct FragmentVisitor;
+
+impl<'de> serde::de::Visitor<'de> for FragmentVisitor {
+    type Value = String;
+
+    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+        write!(formatter, "a string that begins with '... '")
+    }
+
+    fn visit_str<E>(self, s: &str) -> Result<Self::Value, E>
+    where
+        E: serde::de::Error,
+    {
+        if s.starts_with("... ") {
+            Ok(s.to_string())
+        } else {
+            Err(serde::de::Error::invalid_value(
+                serde::de::Unexpected::Str(s),
+                &self,
+            ))
+        }
+    }
+}
+
 /// A path into the result document.
 ///
 /// This can be composed of strings and numbers
@@ -587,6 +633,7 @@ impl fmt::Display for Path {
                 PathElement::Index(index) => write!(f, "{}", index)?,
                 PathElement::Key(key) => write!(f, "{}", key)?,
                 PathElement::Flatten => write!(f, "@")?,
+                PathElement::Fragment(fragment) => write!(f, "{fragment}")?,
             }
         }
         Ok(())

apollo-router/src/query_planner/execution.rs

@@ -389,7 +389,7 @@ impl DeferredNode {
         let mut stream: stream::FuturesUnordered<_> = deferred_receivers.into_iter().collect();
         //FIXME/ is there a solution without cloning the entire node? Maybe it could be moved instead?
         let deferred_inner = self.node.clone();
-        let deferred_path = self.path.clone();
+        let deferred_path = self.query_path.clone();
         let subselection = self.subselection();
         let label = self.label.clone();
         let mut tx = sender;

apollo-router/src/query_planner/plan.rs

@@ -183,6 +183,7 @@ impl PlanNode {
         // re-create full query with the right path
         // parse the subselection
         let mut subselections = HashMap::new();
+
         let operation_kind = if self.contains_mutations() {
             OperationKind::Mutation
         } else {
@@ -225,6 +226,7 @@ impl PlanNode {
                 let primary_path = initial_path.join(&primary.path.clone().unwrap_or_default());
                 if let Some(primary_subselection) = &primary.subselection {
                     let query = reconstruct_full_query(&primary_path, kind, primary_subselection);
+
                     // ----------------------- Parse ---------------------------------
                     let sub_selection = Query::parse(&query, schema, &Default::default())?;
                     // ----------------------- END Parse ---------------------------------
@@ -240,14 +242,15 @@ impl PlanNode {
 
                 deferred.iter().try_fold(subselections, |subs, current| {
                     if let Some(subselection) = &current.subselection {
-                        let query = reconstruct_full_query(&current.path, kind, subselection);
+                        let query = reconstruct_full_query(&current.query_path, kind, subselection);
+
                         // ----------------------- Parse ---------------------------------
                         let sub_selection = Query::parse(&query, schema, &Default::default())?;
                         // ----------------------- END Parse ---------------------------------
 
                         subs.insert(
                             SubSelection {
-                                path: current.path.clone(),
+                                path: current.query_path.clone(),
                                 subselection: subselection.clone(),
                             },
                             sub_selection,
@@ -256,7 +259,7 @@ impl PlanNode {
                     if let Some(current_node) = &current.node {
                         current_node.collect_subselections(
                             schema,
-                            &initial_path.join(&current.path),
+                            &initial_path.join(&current.query_path),
                             kind,
                             subs,
                         )?;
@@ -342,6 +345,11 @@ fn reconstruct_full_query(path: &Path, kind: &OperationKind, subselection: &str)
                     .expect("writing to a String should not fail because it can reallocate");
                 len += 1;
             }
+            json_ext::PathElement::Fragment(fragment) => {
+                write!(&mut query, "{{ {fragment}")
+                    .expect("writing to a String should not fail because it can reallocate");
+                len += 1;
+            }
         }
     }
 
@@ -392,7 +400,7 @@ pub(crate) struct DeferredNode {
     /// The optional defer label.
     pub(crate) label: Option<String>,
     /// Path to the @defer this correspond to. `subselection` start at that `path`.
-    pub(crate) path: Path,
+    pub(crate) query_path: Path,
     /// The part of the original query that "selects" the data to send
     /// in that deferred response (once the plan in `node` completes).
     /// Will be set _unless_ `node` is a `DeferNode` itself.

apollo-router/src/query_planner/testdata/defer_clause_plan.json

@@ -25,7 +25,7 @@
           }
         ],
         "label": null,
-        "path": ["me"],
+        "queryPath": ["me"],
         "subselection": "{ ... on User { name username } }",
         "node": {
           "kind": "Flatten",

apollo-router/src/query_planner/tests.rs

@@ -242,7 +242,7 @@ async fn defer() {
                         defer_label: None,
                     }],
                     label: None,
-                    path: Path(vec![PathElement::Key("t".to_string())]),
+                    query_path: Path(vec![PathElement::Key("t".to_string())]), 
                     subselection: Some("{ y }".to_string()),
                     node: Some(Arc::new(PlanNode::Flatten(FlattenNode {
                         path: Path(vec![PathElement::Key("t".to_string())]),