feat: cache graphql query in query plan. (#3106)

Co-authored-by: Tushar Mathur <[email protected]>

Author: laststylebender14

src/core/document.rs

@@ -1,10 +1,12 @@
+use std::borrow::Cow;
+use std::fmt::Display;
+
 use async_graphql::parser::types::*;
-use async_graphql::{Pos, Positioned};
-use async_graphql_value::{ConstValue, Name};
+use async_graphql::Positioned;
+use async_graphql_value::ConstValue;
 
-fn pos<A>(a: A) -> Positioned<A> {
-    Positioned::new(a, Pos::default())
-}
+use super::jit::Directive as JitDirective;
+use super::json::JsonLikeOwned;
 
 struct LineBreaker<'a> {
     string: &'a str,
@@ -61,9 +63,12 @@ fn get_formatted_docs(docs: Option<String>, indent: usize) -> String {
     formatted_docs
 }
 
-pub fn print_directives<'a>(directives: impl Iterator<Item = &'a ConstDirective>) -> String {
+pub fn print_directives<'a, T>(directives: impl Iterator<Item = &'a T>) -> String
+where
+    &'a T: Into<Directive<'a>> + 'a,
+{
     directives
-        .map(|d| print_directive(&const_directive_to_sdl(d)))
+        .map(|d| print_directive(d))
         .collect::<Vec<String>>()
         .join(" ")
 }
@@ -102,37 +107,6 @@ fn print_schema(schema: &SchemaDefinition) -> String {
     )
 }
 
-fn const_directive_to_sdl(directive: &ConstDirective) -> DirectiveDefinition {
-    DirectiveDefinition {
-        description: None,
-        name: pos(Name::new(directive.name.node.as_str())),
-        arguments: directive
-            .arguments
-            .iter()
-            .filter_map(|(k, v)| {
-                if v.node != ConstValue::Null {
-                    Some(pos(InputValueDefinition {
-                        description: None,
-                        name: pos(Name::new(k.node.clone())),
-                        ty: pos(Type {
-                            nullable: true,
-                            base: async_graphql::parser::types::BaseType::Named(Name::new(
-                                v.to_string(),
-                            )),
-                        }),
-                        default_value: Some(pos(ConstValue::String(v.to_string()))),
-                        directives: Vec::new(),
-                    }))
-                } else {
-                    None
-                }
-            })
-            .collect(),
-        is_repeatable: true,
-        locations: vec![],
-    }
-}
-
 fn print_type_def(type_def: &TypeDefinition) -> String {
     match &type_def.kind {
         TypeKind::Scalar => {
@@ -320,18 +294,23 @@ fn print_input_value(field: &async_graphql::parser::types::InputValueDefinition)
         print_default_value(field.default_value.as_ref())
     )
 }
-fn print_directive(directive: &DirectiveDefinition) -> String {
+
+pub fn print_directive<'a, T>(directive: &'a T) -> String
+where
+    &'a T: Into<Directive<'a>>,
+{
+    let directive: Directive<'a> = directive.into();
     let args = directive
-        .arguments
+        .args
         .iter()
-        .map(|arg| format!("{}: {}", arg.node.name.node, arg.node.ty.node))
+        .map(|arg| format!("{}: {}", arg.name, arg.value))
         .collect::<Vec<String>>()
         .join(", ");
 
     if args.is_empty() {
-        format!("@{}", directive.name.node)
+        format!("@{}", directive.name)
     } else {
-        format!("@{}({})", directive.name.node, args)
+        format!("@{}({})", directive.name, args)
     }
 }
 
@@ -420,3 +399,60 @@ pub fn print(sd: ServiceDocument) -> String {
 
     sdl_string.trim_end_matches('\n').to_string()
 }
+
+pub struct Directive<'a> {
+    pub name: Cow<'a, str>,
+    pub args: Vec<Arg<'a>>,
+}
+
+pub struct Arg<'a> {
+    pub name: Cow<'a, str>,
+    pub value: Cow<'a, str>,
+}
+
+impl<'a> From<&'a ConstDirective> for Directive<'a> {
+    fn from(value: &'a ConstDirective) -> Self {
+        Self {
+            name: Cow::Borrowed(value.name.node.as_str()),
+            args: value
+                .arguments
+                .iter()
+                .filter_map(|(k, v)| {
+                    if v.node != async_graphql_value::ConstValue::Null {
+                        Some(Arg {
+                            name: Cow::Borrowed(k.node.as_str()),
+                            value: Cow::Owned(v.to_string()),
+                        })
+                    } else {
+                        None
+                    }
+                })
+                .collect(),
+        }
+    }
+}
+
+impl<'a, Input: JsonLikeOwned + Display> From<&'a JitDirective<Input>> for Directive<'a> {
+    fn from(value: &'a JitDirective<Input>) -> Self {
+        let to_mustache = |s: &str| -> String {
+            s.strip_prefix('$')
+                .map(|v| format!("{{{{{}}}}}", v))
+                .unwrap_or_else(|| s.to_string())
+        };
+        Self {
+            name: Cow::Borrowed(value.name.as_str()),
+            args: value
+                .arguments
+                .iter()
+                .filter_map(|(k, v)| {
+                    if !v.is_null() {
+                        let v_str = to_mustache(&v.to_string());
+                        Some(Arg { name: Cow::Borrowed(k), value: Cow::Owned(v_str) })
+                    } else {
+                        None
+                    }
+                })
+                .collect(),
+        }
+    }
+}

src/core/graphql/request_template.rs

@@ -6,6 +6,7 @@ use std::hash::{Hash, Hasher};
 use derive_setters::Setters;
 use http::header::{HeaderMap, HeaderValue};
 use tailcall_hasher::TailcallHasher;
+use tracing::info;
 
 use crate::core::config::{GraphQLOperationType, KeyValue};
 use crate::core::has_headers::HasHeaders;
@@ -14,7 +15,35 @@ use crate::core::http::Method::POST;
 use crate::core::ir::model::{CacheKey, IoId};
 use crate::core::ir::{GraphQLOperationContext, RelatedFields};
 use crate::core::mustache::Mustache;
-use crate::core::path::PathGraphql;
+use crate::core::path::{PathGraphql, PathString};
+
+/// Represents a GraphQL selection that can either be resolved or unresolved.
+#[derive(Debug, Clone)]
+pub enum Selection {
+    /// A selection with a resolved string value.
+    Resolved(String),
+    /// A selection that contains a Mustache template to be resolved later.
+    UnResolved(Mustache),
+}
+
+impl Selection {
+    /// Resolves the `Unresolved` variant using the provided `PathString`.
+    pub fn resolve(self, p: &impl PathString) -> Selection {
+        match self {
+            Selection::UnResolved(template) => Selection::Resolved(template.render(p)),
+            resolved => resolved,
+        }
+    }
+}
+
+impl From<Mustache> for Selection {
+    fn from(value: Mustache) -> Self {
+        match value.is_const() {
+            true => Selection::Resolved(value.to_string()),
+            false => Selection::UnResolved(value),
+        }
+    }
+}
 
 /// RequestTemplate for GraphQL requests (See RequestTemplate documentation)
 #[derive(Setters, Debug, Clone)]
@@ -26,6 +55,7 @@ pub struct RequestTemplate {
     pub operation_arguments: Option<Vec<(String, Mustache)>>,
     pub headers: MustacheHeaders,
     pub related_fields: RelatedFields,
+    pub selection: Option<Selection>,
 }
 
 impl RequestTemplate {
@@ -85,7 +115,12 @@ impl RequestTemplate {
         ctx: &C,
     ) -> String {
         let operation_type = &self.operation_type;
-        let selection_set = ctx.selection_set(&self.related_fields).unwrap_or_default();
+
+        let selection_set = match &self.selection {
+            Some(Selection::Resolved(s)) => Cow::Borrowed(s),
+            Some(Selection::UnResolved(u)) => Cow::Owned(u.to_string()),
+            None => Cow::Owned(ctx.selection_set(&self.related_fields).unwrap_or_default()),
+        };
 
         let mut operation = Cow::Borrowed(&self.operation_name);
 
@@ -121,7 +156,10 @@ impl RequestTemplate {
             }
         }
 
-        format!(r#"{{ "query": "{operation_type} {{ {operation} {selection_set} }}" }}"#)
+        let query =
+            format!(r#"{{ "query": "{operation_type} {{ {operation} {selection_set} }}" }}"#);
+        info!("Query {} ", query);
+        query
     }
 
     pub fn new(
@@ -149,6 +187,7 @@ impl RequestTemplate {
             operation_arguments,
             headers,
             related_fields,
+            selection: None,
         })
     }
 }

src/core/ir/model.rs

@@ -131,6 +131,28 @@ impl Cache {
 }
 
 impl IR {
+    // allows to modify the IO node in the IR tree
+    pub fn modify_io(&mut self, io_modifier: &mut dyn FnMut(&mut IO)) {
+        match self {
+            IR::IO(io) => io_modifier(io),
+            IR::Cache(cache) => io_modifier(&mut cache.io),
+            IR::Discriminate(_, ir) | IR::Protect(_, ir) | IR::Path(ir, _) => {
+                ir.modify_io(io_modifier)
+            }
+            IR::Pipe(ir1, ir2) => {
+                ir1.modify_io(io_modifier);
+                ir2.modify_io(io_modifier);
+            }
+            IR::Entity(hash_map) => {
+                for ir in hash_map.values_mut() {
+                    ir.modify_io(io_modifier);
+                }
+            }
+            IR::Map(map) => map.input.modify_io(io_modifier),
+            _ => {}
+        }
+    }
+
     pub fn pipe(self, next: Self) -> Self {
         IR::Pipe(Box::new(self), Box::new(next))
     }

src/core/jit/model.rs

@@ -1,6 +1,6 @@
 use std::borrow::Cow;
 use std::collections::HashMap;
-use std::fmt::{Debug, Formatter};
+use std::fmt::{Debug, Display, Formatter};
 use std::num::NonZeroU64;
 use std::sync::Arc;
 
@@ -13,12 +13,20 @@ use super::Error;
 use crate::core::blueprint::Index;
 use crate::core::ir::model::IR;
 use crate::core::ir::TypedValue;
-use crate::core::json::JsonLike;
+use crate::core::json::{JsonLike, JsonLikeOwned};
+use crate::core::path::PathString;
 use crate::core::scalar::Scalar;
 
 #[derive(Debug, Deserialize, Clone)]
 pub struct Variables<Value>(HashMap<String, Value>);
 
+impl<V: JsonLikeOwned + Display> PathString for Variables<V> {
+    fn path_string<'a, T: AsRef<str>>(&'a self, path: &'a [T]) -> Option<Cow<'a, str>> {
+        self.get(path[0].as_ref())
+            .map(|v| Cow::Owned(v.to_string()))
+    }
+}
+
 impl<Value> Default for Variables<Value> {
     fn default() -> Self {
         Self::new()
@@ -96,6 +104,22 @@ pub struct Arg<Input> {
     pub default_value: Option<Input>,
 }
 
+impl<Input: Display> Display for Arg<Input> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let v = self
+            .value
+            .as_ref()
+            .map(|v| v.to_string())
+            .unwrap_or_else(|| {
+                self.default_value
+                    .as_ref()
+                    .map(|v| v.to_string())
+                    .unwrap_or_default()
+            });
+        write!(f, "{}: {}", self.name, v)
+    }
+}
+
 impl<Input> Arg<Input> {
     pub fn try_map<Output, Error>(
         self,

src/core/jit/request.rs

@@ -50,6 +50,7 @@ impl Request<ConstValue> {
             .pipe(transform::AuthPlanner::new())
             .pipe(transform::CheckDedupe::new())
             .pipe(transform::CheckCache::new())
+            .pipe(transform::GraphQL::new())
             .transform(plan)
             .to_result()
             // both transformers are infallible right now