fix: query interface with fragments

tailcallhq · tusharmath · Jan 17, 2025 · Jan 4, 2025 · Jan 4, 2025 · Jan 5, 2025
commit d23720edc17b336e51a64371692c420b74ea2791
diff --git a/src/core/blueprint/index.rs b/src/core/blueprint/index.rs
@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use indexmap::IndexMap;
 
 use super::InputObjectTypeDefinition;
@@ -38,6 +40,18 @@ impl Index {
         matches!(def, Some(Definition::Scalar(_))) || scalar::Scalar::is_predefined(type_name)
     }
 
+    pub fn get_interfaces(&self) -> Box<HashSet<String>> {
+        Box::new(
+            self.map
+                .iter()
+                .filter_map(|(_, (def, _))| match def {
+                    Definition::Interface(i) => Some(i.name.to_owned()),
+                    _ => None,
+                })
+                .collect(),
+        )
+    }
+
     pub fn type_is_enum(&self, type_name: &str) -> bool {
         let def = self.map.get(type_name).map(|(def, _)| def);
 

diff --git a/src/core/jit/builder.rs b/src/core/jit/builder.rs
@@ -128,6 +128,7 @@ impl<'a> Builder<'a> {
     #[inline(always)]
     fn iter(
         &self,
+        parent_fragment: Option<&str>,
         selection: &SelectionSet,
         type_condition: &str,
         fragments: &HashMap<&str, &FragmentDefinition>,
@@ -168,6 +169,7 @@ impl<'a> Builder<'a> {
                         .map(|(k, v)| (k.node.as_str().to_string(), v.node.to_owned()))
                         .collect::<HashMap<_, _>>();
 
+                    let parent_fragment = parent_fragment.map(|s| s.to_owned());
                     // Check if the field is present in the schema index
                     if let Some(field_def) = self.index.get_field(type_condition, field_name) {
                         let mut args = Vec::with_capacity(request_args.len());
@@ -198,8 +200,12 @@ impl<'a> Builder<'a> {
                         let id = FieldId::new(self.field_id.next());
 
                         // Recursively gather child fields for the selection set
-                        let child_fields =
-                            self.iter(&gql_field.selection_set.node, type_of.name(), fragments);
+                        let child_fields = self.iter(
+                            None,
+                            &gql_field.selection_set.node,
+                            type_of.name(),
+                            fragments,
+                        );
 
                         let ir = match field_def {
                             QueryField::Field((field_def, _)) => field_def.resolver.clone(),
@@ -220,6 +226,7 @@ impl<'a> Builder<'a> {
                         let field = Field {
                             id,
                             selection: child_fields,
+                            parent_fragment,
                             name: field_name.to_string(),
                             output_name: gql_field
                                 .alias
@@ -252,6 +259,7 @@ impl<'a> Builder<'a> {
                             args: Vec::new(),
                             pos: selection.pos.into(),
                             selection: vec![], // __typename has no child selection
+                            parent_fragment,
                             directives,
                             is_enum: false,
                             scalar: Some(scalar::Scalar::Empty),
@@ -265,6 +273,7 @@ impl<'a> Builder<'a> {
                         fragments.get(fragment_spread.fragment_name.node.as_str())
                     {
                         fields.extend(self.iter(
+                            Some(fragment.type_condition.node.on.node.as_str()),
                             &fragment.selection_set.node,
                             fragment.type_condition.node.on.node.as_str(),
                             fragments,
@@ -277,8 +286,12 @@ impl<'a> Builder<'a> {
                         .as_ref()
                         .map(|cond| cond.node.on.node.as_str())
                         .unwrap_or(type_condition);
-
-                    fields.extend(self.iter(&fragment.selection_set.node, type_of, fragments));
+                    fields.extend(self.iter(
+                        Some(type_of),
+                        &fragment.selection_set.node,
+                        type_of,
+                        fragments,
+                    ));
                 }
             }
         }
@@ -334,7 +347,7 @@ impl<'a> Builder<'a> {
         let name = self
             .get_type(operation.ty)
             .ok_or(BuildError::RootOperationTypeNotDefined { operation: operation.ty })?;
-        let fields = self.iter(&operation.selection_set.node, name, &fragments);
+        let fields = self.iter(None, &operation.selection_set.node, name, &fragments);
 
         let is_introspection_query = operation.selection_set.node.items.iter().any(|f| {
             if let Selection::Field(Positioned { node: gql_field, .. }) = &f.node {
@@ -351,6 +364,7 @@ impl<'a> Builder<'a> {
             operation.ty,
             self.index.clone(),
             is_introspection_query,
+            Some(*self.index.get_interfaces()),
         );
         Ok(plan)
     }

diff --git a/src/core/jit/model.rs b/src/core/jit/model.rs
@@ -1,5 +1,5 @@
 use std::borrow::Cow;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Debug, Display, Formatter};
 use std::num::NonZeroU64;
 use std::sync::Arc;
@@ -187,6 +187,7 @@ pub struct Field<Input> {
     pub include: Option<Variable>,
     pub args: Vec<Arg<Input>>,
     pub selection: Vec<Field<Input>>,
+    pub parent_fragment: Option<String>,
     pub pos: Pos,
     pub directives: Vec<Directive<Input>>,
     pub is_enum: bool,
@@ -245,6 +246,7 @@ impl<Input> Field<Input> {
                 .into_iter()
                 .map(|f| f.try_map(map))
                 .collect::<Result<Vec<Field<Output>>, Error>>()?,
+            parent_fragment: None,
             skip: self.skip,
             include: self.include,
             pos: self.pos,
@@ -315,6 +317,7 @@ pub struct OperationPlan<Input> {
     pub min_cache_ttl: Option<NonZeroU64>,
     pub selection: Vec<Field<Input>>,
     pub before: Option<IR>,
+    pub interfaces: Option<HashSet<String>>,
 }
 
 impl<Input> OperationPlan<Input> {
@@ -339,6 +342,7 @@ impl<Input> OperationPlan<Input> {
             is_protected: self.is_protected,
             min_cache_ttl: self.min_cache_ttl,
             before: self.before,
+            interfaces: None,
         })
     }
 }
@@ -351,6 +355,7 @@ impl<Input> OperationPlan<Input> {
         operation_type: OperationType,
         index: Arc<Index>,
         is_introspection_query: bool,
+        interfaces: Option<HashSet<String>>,
     ) -> Self
     where
         Input: Clone,
@@ -366,6 +371,7 @@ impl<Input> OperationPlan<Input> {
             is_protected: false,
             min_cache_ttl: None,
             before: Default::default(),
+            interfaces,
         }
     }
 

diff --git a/src/core/jit/transform/graphql.rs b/src/core/jit/transform/graphql.rs
@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::collections::{HashMap, HashSet};
 use std::convert::Infallible;
 use std::fmt::{Debug, Display};
 use std::marker::PhantomData;
@@ -20,18 +21,25 @@ impl<A> GraphQL<A> {
     }
 }
 
-fn compute_selection_set<A: Display + Debug + JsonLikeOwned>(base_field: &mut [Field<A>]) {
+fn compute_selection_set<A: Display + Debug + JsonLikeOwned>(
+    base_field: &mut [Field<A>],
+    interfaces: &HashSet<String>,
+) {
     for field in base_field.iter_mut() {
         if let Some(ir) = field.ir.as_mut() {
             ir.modify_io(&mut |io| {
                 if let IO::GraphQL { req_template, .. } = io {
-                    if let Some(v) = format_selection_set(field.selection.iter()) {
+                    if let Some(v) = format_selection_set(
+                        field.selection.iter(),
+                        interfaces,
+                        interfaces.contains(field.type_of.name()),
+                    ) {
                         req_template.selection = Some(Mustache::parse(&v).into());
                     }
                 }
             });
         }
-        compute_selection_set(field.selection.as_mut());
+        compute_selection_set(field.selection.as_mut(), interfaces);
     }
 }
 
@@ -40,15 +48,23 @@ impl<A: Display + Debug + JsonLikeOwned + Clone> Transform for GraphQL<A> {
     type Error = Infallible;
 
     fn transform(&self, mut plan: Self::Value) -> Valid<Self::Value, Self::Error> {
-        compute_selection_set(&mut plan.selection);
+        let interfaces = match plan.interfaces {
+            Some(ref interfaces) => interfaces,
+            None => &HashSet::new(),
+        };
+        compute_selection_set(&mut plan.selection, interfaces);
 
         Valid::succeed(plan)
     }
 }
 
 fn format_selection_set<'a, A: 'a + Display + JsonLikeOwned>(
     selection_set: impl Iterator<Item = &'a Field<A>>,
+    interfaces: &HashSet<String>,
+    is_parent_interface: bool,
 ) -> Option<String> {
+    let mut fragments_fields = HashMap::new();
+    let mut normal_fields = vec![];
     let set = selection_set
         .filter(|field| !matches!(&field.ir, Some(IR::IO(_)) | Some(IR::Dynamic(_))))
         .map(|field| {
@@ -58,20 +74,51 @@ fn format_selection_set<'a, A: 'a + Display + JsonLikeOwned>(
             } else {
                 field.name.to_string()
             };
-            format_selection_field(field, &field_name)
+            let is_this_field_interface = interfaces.contains(field.type_of.name());
+            let formatted_selection_fields =
+                format_selection_field(field, &field_name, interfaces, is_this_field_interface);
+            match &field.parent_fragment {
+                Some(fragment) if is_parent_interface => {
+                    fragments_fields
+                        .entry(fragment.to_owned())
+                        .or_insert_with(Vec::new)
+                        .push(formatted_selection_fields);
+                }
+                _ => {
+                    normal_fields.push(formatted_selection_fields);
+                }
+            }
         })
         .collect::<Vec<_>>();
 
     if set.is_empty() {
         return None;
     }
 
-    Some(format!("{{ {} }}", set.join(" ")))
+    let string_set: Vec<String> = fragments_fields
+        .into_iter()
+        .map(|(fragment_name, fields)| {
+            format!("... on {} {{ {} }}", fragment_name, fields.join(" "))
+        })
+        .collect();
+
+    //Don't force user to query the type and get it automatically
+    if is_parent_interface {
+        normal_fields.push("__typename".to_owned());
+    }
+    normal_fields.push(string_set.join(" "));
+    Some(format!("{{ {} }}", normal_fields.join(" ")))
 }
 
-fn format_selection_field<A: Display + JsonLikeOwned>(field: &Field<A>, name: &str) -> String {
+fn format_selection_field<A: Display + JsonLikeOwned>(
+    field: &Field<A>,
+    name: &str,
+    interfaces: &HashSet<String>,
+    is_parent_interface: bool,
+) -> String {
     let arguments = format_selection_field_arguments(field);
-    let selection_set = format_selection_set(field.selection.iter());
+    let selection_set =
+        format_selection_set(field.selection.iter(), interfaces, is_parent_interface);
 
     let mut output = format!("{}{}", name, arguments);
 

diff --git a/src/core/jit/transform/input_resolver.rs b/src/core/jit/transform/input_resolver.rs
@@ -84,6 +84,7 @@ where
             is_const: self.plan.is_const,
             is_protected: self.plan.is_protected,
             min_cache_ttl: self.plan.min_cache_ttl,
+            interfaces: None,
             selection,
             before: self.plan.before,
         })

diff --git a/tests/core/snapshots/graphql-with-fragments.md_0.snap b/tests/core/snapshots/graphql-with-fragments.md_0.snap
@@ -0,0 +1,27 @@
+---
+source: tests/core/spec.rs
+expression: response
+snapshot_kind: text
+---
+{
+  "status": 200,
+  "headers": {
+    "content-type": "application/json"
+  },
+  "body": {
+    "data": {
+      "queryNodeC": [
+        {
+          "name": "nodeA",
+          "__typename": "NodeA",
+          "nodeA_id": "nodeA_id"
+        },
+        {
+          "name": "nodeB",
+          "__typename": "NodeB",
+          "nodeB_id": "nodeB_id"
+        }
+      ]
+    }
+  }
+}
diff --git a/tests/core/snapshots/graphql-with-fragments.md_client.snap b/tests/core/snapshots/graphql-with-fragments.md_client.snap
@@ -0,0 +1,28 @@
+---
+source: tests/core/spec.rs
+expression: formatted
+snapshot_kind: text
+---
+type NodeA implements NodeC {
+  name: String
+  nodeA_id: String
+}
+
+type NodeB implements NodeC {
+  name: String
+  nodeB_id: String
+}
+
+interface NodeC {
+  name: String
+}
+
+type Query {
+  queryNodeA: [NodeA!]
+  queryNodeB: [NodeB!]
+  queryNodeC: [NodeC!]
+}
+
+schema {
+  query: Query
+}
diff --git a/tests/core/snapshots/graphql-with-fragments.md_merged.snap b/tests/core/snapshots/graphql-with-fragments.md_merged.snap
@@ -0,0 +1,28 @@
+---
+source: tests/core/spec.rs
+expression: formatter
+snapshot_kind: text
+---
+schema @server @upstream @link(src: "schema_0.graphql", type: Config) {
+  query: Query
+}
+
+interface NodeC {
+  name: String
+}
+
+type NodeA implements NodeC {
+  name: String
+  nodeA_id: String
+}
+
+type NodeB implements NodeC {
+  name: String
+  nodeB_id: String
+}
+
+type Query {
+  queryNodeA: [NodeA!] @graphQL(url: "http://upstream/graphql", name: "nodeA")
+  queryNodeB: [NodeB!] @graphQL(url: "http://upstream/graphql", name: "nodeB")
+  queryNodeC: [NodeC!] @graphQL(url: "http://upstream/graphql", name: "nodeC")
+}