feat(jit): support introspection in JIT

src/core/jit/builder.rs

@@ -344,7 +344,22 @@
         // skip the fields depending on variables.
         fields.retain(|f| !f.skip(variables));
 
-        let plan = OperationPlan::new(fields, operation.ty, self.index.clone());
+        let is_introspection_query = operation.selection_set.node.items.iter().any(|f| {
+            if let Selection::Field(Positioned { node: gql_field, .. }) = &f.node {
+                let query = gql_field.name.node.as_str();
+                query.contains("__schema") || query.contains("__type")
+            } else {
+                false
+            }
+        });
+
+        let plan = OperationPlan::new(
+            fields,
+            operation.ty,
+            self.index.clone(),
+            is_introspection_query,
+        );
+
         // TODO: operation from [ExecutableDocument] could contain definitions for
         // default values of arguments. That info should be passed to
         // [InputResolver] to resolve defaults properly

src/core/jit/exec_const.rs

@@ -13,7 +13,7 @@ use crate::core::jit::synth::Synth;
 
 /// A specialized executor that executes with async_graphql::Value
 pub struct ConstValueExecutor {
-    plan: OperationPlan<ConstValue>,
+    pub plan: OperationPlan<ConstValue>,
 }
 
 impl ConstValueExecutor {
@@ -24,7 +24,7 @@ impl ConstValueExecutor {
     pub async fn execute(
         self,
         req_ctx: &RequestContext,
-        request: Request<ConstValue>,
+        request: &Request<ConstValue>,
     ) -> Response<ConstValue, Error> {
         let exec = ConstValueExec::new(req_ctx);
         let plan = self.plan;

src/core/jit/graphql_executor.rs

@@ -1,13 +1,15 @@
+use std::collections::BTreeMap;
 use std::future::Future;
 use std::sync::Arc;
 
-use async_graphql::{Data, Executor, Response};
+use async_graphql::{Data, Executor, Response, Value};
 use futures_util::stream::BoxStream;
 
 use crate::core::app_context::AppContext;
 use crate::core::http::RequestContext;
 use crate::core::jit;
 use crate::core::jit::ConstValueExecutor;
+use crate::core::merge_right::MergeRight;
 
 #[derive(Clone)]
 pub struct JITExecutor {
@@ -21,15 +23,47 @@
     }
 }
 
+impl From<jit::Request<Value>> for async_graphql::Request {
+    fn from(value: jit::Request<Value>) -> Self {
+        let mut request = async_graphql::Request::new(value.query);
+        request.variables.extend(
+            value
+                .variables
+                .into_hashmap()
+                .into_iter()
+                .map(|(k, v)| (async_graphql::Name::new(k), v))
+                .collect::<BTreeMap<_, _>>(),
+        );
+        request.extensions = value.extensions;
+        request.operation_name = value.operation_name;
+        request
+    }
+}
+
 impl Executor for JITExecutor {
     fn execute(&self, request: async_graphql::Request) -> impl Future<Output = Response> + Send {
-        let request = jit::Request::from(request);
+        let jit_request = jit::Request::from(request);
 
         async {
-            match ConstValueExecutor::new(&request, self.app_ctx.clone()) {
+            match ConstValueExecutor::new(&jit_request, self.app_ctx.clone()) {
                 Ok(exec) => {
-                    let resp = exec.execute(&self.req_ctx, request).await;
-                    resp.into_async_graphql()
+                    let is_introspection_query =
+                        self.app_ctx.blueprint.server.get_enable_introspection()
+                            && exec.plan.is_introspection_query;
+
+                    let jit_resp = exec
+                        .execute(&self.req_ctx, &jit_request)
+                        .await
+                        .into_async_graphql();
+
+                    if is_introspection_query {
+                        let async_req =
+                            async_graphql::Request::from(jit_request).only_introspection();
+                        let async_resp = self.app_ctx.execute(async_req).await;
+                        jit_resp.merge_right(async_resp)
+                    } else {
+                        jit_resp
+                    }
                 }
                 Err(error) => Response::from_errors(vec![error.into()]),
             }

src/core/jit/input_resolver.rs

@@ -63,6 +63,7 @@ where
             new_fields,
             self.plan.operation_type(),
             self.plan.index.clone(),
+            self.plan.is_introspection_query,
         ))
     }
 }

src/core/jit/model.rs

@@ -29,6 +29,10 @@
     pub fn get(&self, key: &str) -> Option<&Value> {
         self.0.get(key)
     }
+    pub fn into_hashmap(self) -> HashMap<String, Value> {
+        self.0
+    }
+
     pub fn insert(&mut self, key: String, value: Value) {
         self.0.insert(key, value);
     }
@@ -350,6 +354,7 @@
     nested: Vec<Field<Nested<Input>, Input>>,
     // TODO: drop index from here. Embed all the necessary information in each field of the plan.
     pub index: Arc<Index>,
+    pub is_introspection_query: bool,
 }
 
 impl<Input> std::fmt::Debug for OperationPlan<Input> {
@@ -382,6 +387,7 @@
             operation_type: self.operation_type,
             nested,
             index: self.index,
+            is_introspection_query: self.is_introspection_query,
         })
     }
 }
@@ -391,6 +397,7 @@
         fields: Vec<Field<Flat, Input>>,
         operation_type: OperationType,
         index: Arc<Index>,
+        is_introspection_query: bool,
     ) -> Self
     where
         Input: Clone,
@@ -402,7 +409,13 @@
             .map(|f| f.into_nested(&fields))
             .collect::<Vec<_>>();
 
-        Self { flat: fields, nested, operation_type, index }
+        Self {
+            flat: fields,
+            nested,
+            operation_type,
+            index,
+            is_introspection_query,
+        }
     }
 
     /// Returns a graphQL operation type

src/core/jit/request.rs

@@ -19,6 +19,7 @@ pub struct Request<V> {
     pub extensions: HashMap<String, V>,
 }
 
+// NOTE: This is hot code and should allocate minimal memory
 impl From<async_graphql::Request> for Request<ConstValue> {
     fn from(mut value: async_graphql::Request) -> Self {
         let variables = std::mem::take(value.variables.deref_mut());

src/core/jit/response.rs

@@ -1,8 +1,11 @@
+use std::borrow::BorrowMut;
+
 use derive_setters::Setters;
 use serde::Serialize;
 
 use super::Positioned;
 use crate::core::jit;
+use crate::core::merge_right::MergeRight;
 
 #[derive(Setters, Serialize)]
 pub struct Response<Value, Error> {
@@ -31,6 +34,22 @@ impl<Value, Error> Response<Value, Error> {
     }
 }
 
+impl MergeRight for async_graphql::Response {
+    fn merge_right(mut self, other: Self) -> Self {
+        if let async_graphql::Value::Object(mut other_obj) = other.data {
+            if let async_graphql::Value::Object(self_obj) = std::mem::take(self.data.borrow_mut()) {
+                other_obj.extend(self_obj);
+                self.data = async_graphql::Value::Object(other_obj);
+            }
+        }
+
+        self.errors.extend(other.errors);
+        self.extensions.extend(other.extensions);
+
+        self
+    }
+}
+
 impl Response<async_graphql::Value, jit::Error> {
     pub fn into_async_graphql(self) -> async_graphql::Response {
         let mut resp = async_graphql::Response::new(self.data.unwrap_or_default());
@@ -50,6 +69,7 @@ mod test {
 
     use super::Response;
     use crate::core::jit::{self, Pos, Positioned};
+    use crate::core::merge_right::MergeRight;
 
     #[test]
     fn test_with_response() {
@@ -117,4 +137,66 @@ mod test {
         assert_eq!(async_response.errors.len(), 2);
         insta::assert_debug_snapshot!(async_response);
     }
+
+    #[test]
+    pub fn test_merging_of_responses() {
+        let introspection_response = r#"
+        {
+            "__type": {
+                "name": "User",
+                "fields": [
+                    {
+                        "name": "birthday",
+                        "type": {
+                            "name": "Date"
+                        }
+                    },
+                    {
+                        "name": "id",
+                        "type": {
+                            "name": "String"
+                        }
+                    }
+                ]
+            }
+        }
+        "#;
+        let introspection_data =
+            ConstValue::from_json(serde_json::from_str(introspection_response).unwrap()).unwrap();
+        let introspection_response = async_graphql::Response::new(introspection_data);
+
+        let user_response = r#"
+        {
+            "me": {
+                "id": 1,
+                "name": "John Smith",
+                "birthday": "2023-03-08T12:45:26-05:00"
+            }
+        }
+        "#;
+        let user_data =
+            ConstValue::from_json(serde_json::from_str(user_response).unwrap()).unwrap();
+        let query_response = async_graphql::Response::new(user_data);
+
+        let merged_response = introspection_response.merge_right(query_response);
+
+        insta::assert_json_snapshot!(merged_response);
+    }
+
+    #[test]
+    pub fn test_merging_of_errors() {
+        let mut resp1 = async_graphql::Response::new(ConstValue::default());
+        let mut err1 = vec![async_graphql::ServerError::new("Error-1", None)];
+        resp1.errors.append(&mut err1);
+
+        let mut resp2 = async_graphql::Response::new(ConstValue::default());
+        let mut err2 = vec![async_graphql::ServerError::new(
+            "Error-2",
+            Some(async_graphql::Pos::default()),
+        )];
+        resp2.errors.append(&mut err2);
+
+        let merged_resp = resp1.merge_right(resp2);
+        insta::assert_json_snapshot!(merged_resp);
+    }
 }

src/core/jit/snapshots/tailcall__core__jit__response__test__merging.snap

@@ -0,0 +1,30 @@
+---
+source: src/core/jit/response.rs
+expression: merged_response
+---
+{
+  "data": {
+    "me": {
+      "id": 1,
+      "name": "John Smith",
+      "birthday": "2023-03-08T12:45:26-05:00"
+    },
+    "__type": {
+      "name": "User",
+      "fields": [
+        {
+          "name": "birthday",
+          "type": {
+            "name": "Date"
+          }
+        },
+        {
+          "name": "id",
+          "type": {
+            "name": "String"
+          }
+        }
+      ]
+    }
+  }
+}

src/core/jit/snapshots/tailcall__core__jit__response__test__merging_of_errors.snap

@@ -0,0 +1,21 @@
+---
+source: src/core/jit/response.rs
+expression: merged_resp
+---
+{
+  "data": null,
+  "errors": [
+    {
+      "message": "Error-1"
+    },
+    {
+      "message": "Error-2",
+      "locations": [
+        {
+          "line": 0,
+          "column": 0
+        }
+      ]
+    }
+  ]
+}

src/core/jit/snapshots/tailcall__core__jit__response__test__merging_of_responses.snap

@@ -0,0 +1,30 @@
+---
+source: src/core/jit/response.rs
+expression: merged_response
+---
+{
+  "data": {
+    "me": {
+      "id": 1,
+      "name": "John Smith",
+      "birthday": "2023-03-08T12:45:26-05:00"
+    },
+    "__type": {
+      "name": "User",
+      "fields": [
+        {
+          "name": "birthday",
+          "type": {
+            "name": "Date"
+          }
+        },
+        {
+          "name": "id",
+          "type": {
+            "name": "String"
+          }
+        }
+      ]
+    }
+  }
+}