fix(jit): add field aliases and args support (#2709)

Co-authored-by: Panagiotis Karatakis <[email protected]>
Co-authored-by: Panagiotis <[email protected]>
Co-authored-by: Kiryl Mialeshka <[email protected]>

src/core/blueprint/from_config.rs

@@ -1,6 +1,7 @@
 use std::collections::{BTreeMap, BTreeSet};
 
 use async_graphql::dynamic::SchemaBuilder;
+use indexmap::IndexMap;
 
 use self::telemetry::to_opentelemetry;
 use super::{Server, TypeLike};
@@ -70,7 +71,7 @@ pub fn apply_batching(mut blueprint: Blueprint) -> Blueprint {
 pub fn to_json_schema_for_field(field: &Field, config: &Config) -> JsonSchema {
     to_json_schema(field, config)
 }
-pub fn to_json_schema_for_args(args: &BTreeMap<String, Arg>, config: &Config) -> JsonSchema {
+pub fn to_json_schema_for_args(args: &IndexMap<String, Arg>, config: &Config) -> JsonSchema {
     let mut schema_fields = BTreeMap::new();
     for (name, arg) in args.iter() {
         schema_fields.insert(name.clone(), to_json_schema(arg, config));

src/core/config/config.rs

@@ -1,11 +1,12 @@
-use std::collections::{BTreeMap, BTreeSet, HashSet};
+use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::fmt::{self, Display};
 use std::num::NonZeroU64;
 
 use anyhow::Result;
 use async_graphql::parser::types::{ConstDirective, ServiceDocument};
 use async_graphql::Positioned;
 use derive_setters::Setters;
+use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use tailcall_macros::{CustomResolver, DirectiveDefinition, InputDefinition};
@@ -254,7 +255,8 @@ pub struct Field {
     ///
     /// Map of argument name and its definition.
     #[serde(default, skip_serializing_if = "is_default")]
-    pub args: BTreeMap<String, Arg>,
+    #[schemars(with = "HashMap::<String, Arg>")]
+    pub args: IndexMap<String, Arg>,
 
     ///
     /// Publicly visible documentation for the field.

src/core/config/from_document.rs

@@ -8,6 +8,7 @@ use async_graphql::parser::types::{
 use async_graphql::parser::Positioned;
 use async_graphql::Name;
 use async_graphql_value::ConstValue;
+use indexmap::IndexMap;
 
 use super::telemetry::Telemetry;
 use super::Alias;
@@ -294,7 +295,7 @@ fn to_field(field_definition: &FieldDefinition) -> Valid<config::Field, String>
 fn to_input_object_field(field_definition: &InputValueDefinition) -> Valid<config::Field, String> {
     to_common_field(
         field_definition,
-        BTreeMap::new(),
+        IndexMap::new(),
         field_definition
             .default_value
             .as_ref()
@@ -303,7 +304,7 @@ fn to_input_object_field(field_definition: &InputValueDefinition) -> Valid<confi
 }
 fn to_common_field<F>(
     field: &F,
-    args: BTreeMap<String, config::Arg>,
+    args: IndexMap<String, config::Arg>,
     default_value: Option<ConstValue>,
 ) -> Valid<config::Field, String>
 where
@@ -356,8 +357,8 @@ fn to_type_of(type_: &Type) -> String {
         BaseType::List(ty) => to_type_of(ty),
     }
 }
-fn to_args(field_definition: &FieldDefinition) -> BTreeMap<String, config::Arg> {
-    let mut args: BTreeMap<String, config::Arg> = BTreeMap::new();
+fn to_args(field_definition: &FieldDefinition) -> IndexMap<String, config::Arg> {
+    let mut args = IndexMap::new();
 
     for arg in field_definition.arguments.iter() {
         let arg_name = pos_name_to_string(&arg.node.name);

src/core/config/transformer/snapshots/tailcall__core__config__transformer__union_input_type__tests__union_in_type.snap

@@ -66,13 +66,13 @@ type NU {
 }
 
 type Query {
-  testVar0Var0(nnu: NNU__nu0, nu: NU__u0!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
-  testVar0Var1(nnu: NNU__nu0, nu: NU__u1!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
-  testVar0Var2(nnu: NNU__nu0, nu: NU__u2!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
-  testVar1Var0(nnu: NNU__nu1, nu: NU__u0!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
-  testVar1Var1(nnu: NNU__nu1, nu: NU__u1!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
-  testVar1Var2(nnu: NNU__nu1, nu: NU__u2!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
-  testVar2Var0(nnu: NNU__nu2, nu: NU__u0!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
-  testVar2Var1(nnu: NNU__nu2, nu: NU__u1!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
-  testVar2Var2(nnu: NNU__nu2, nu: NU__u2!): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar0Var0(nu: NU__u0!, nnu: NNU__nu0): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar0Var1(nu: NU__u0!, nnu: NNU__nu1): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar0Var2(nu: NU__u0!, nnu: NNU__nu2): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar1Var0(nu: NU__u1!, nnu: NNU__nu0): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar1Var1(nu: NU__u1!, nnu: NNU__nu1): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar1Var2(nu: NU__u1!, nnu: NNU__nu2): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar2Var0(nu: NU__u2!, nnu: NNU__nu0): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar2Var1(nu: NU__u2!, nnu: NNU__nu1): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
+  testVar2Var2(nu: NU__u2!, nnu: NNU__nu2): U @http(baseURL: "http://localhost", path: "/users/{{args.nu.u}}")
 }

src/core/generator/tests/snapshots/json_to_config_spec__add_cart.json.snap

@@ -17,7 +17,7 @@ input T4 {
 }
 
 type Mutation {
-  addCart(addCartInput: T4, code: String): T2 @http(baseURL: "https://dummyjson.com", body: "{{.args.addCartInput}}", method: "POST", path: "/carts/add", query: [{key: "code", value: "{{.args.code}}"}])
+  addCart(code: String, addCartInput: T4): T2 @http(baseURL: "https://dummyjson.com", body: "{{.args.addCartInput}}", method: "POST", path: "/carts/add", query: [{key: "code", value: "{{.args.code}}"}])
 }
 
 type T1 {

src/core/ir/resolver_context_like.rs

@@ -95,7 +95,7 @@ impl SelectionField {
     fn from_jit_field(
         field: &crate::core::jit::Field<Nested<ConstValue>, ConstValue>,
     ) -> SelectionField {
-        let name = field.name.clone();
+        let name = field.output_name.to_string();
         let selection_set = field
             .nested_iter(field.type_of.name())
             .map(Self::from_jit_field)

src/core/jit/builder.rs

@@ -206,18 +206,18 @@ impl Builder {
                             Some(Flat::new(id.clone())),
                             fragments,
                         );
-                        let name = gql_field
-                            .alias
-                            .as_ref()
-                            .map(|alias| alias.node.to_string())
-                            .unwrap_or(field_name.to_string());
                         let ir = match field_def {
                             QueryField::Field((field_def, _)) => field_def.resolver.clone(),
                             _ => None,
                         };
                         let flat_field = Field {
                             id,
-                            name,
+                            name: field_name.to_string(),
+                            output_name: gql_field
+                                .alias
+                                .as_ref()
+                                .map(|a| a.node.to_string())
+                                .unwrap_or(field_name.to_owned()),
                             ir,
                             type_of,
                             type_condition: type_condition.to_string(),
@@ -404,6 +404,21 @@ mod tests {
         insta::assert_debug_snapshot!(plan.into_nested());
     }
 
+    #[test]
+    fn test_alias_query() {
+        let plan = plan(
+            r#"
+            query {
+                articles: posts { author: user { identifier: id } }
+            }
+        "#,
+            &Variables::new(),
+        );
+
+        assert!(plan.is_query());
+        insta::assert_debug_snapshot!(plan.into_nested());
+    }
+
     #[test]
     fn test_simple_mutation() {
         let plan = plan(

src/core/jit/context.rs

@@ -2,6 +2,7 @@ use std::sync::{Arc, Mutex, MutexGuard};
 
 use async_graphql::{Name, ServerError};
 use async_graphql_value::ConstValue;
+use indexmap::IndexMap;
 
 use super::error::*;
 use super::{Field, Nested, OperationPlan, Positioned};
@@ -39,28 +40,20 @@ pub struct Context<'a, Input, Output> {
     request: &'a RequestContext<Input>,
 }
 impl<'a, Input: Clone, Output> Context<'a, Input, Output> {
-    pub fn new(field: &'a Field<Nested<Input>, Input>, env: &'a RequestContext<Input>) -> Self {
-        Self { value: None, args: None, field, request: env }
+    pub fn new(field: &'a Field<Nested<Input>, Input>, request: &'a RequestContext<Input>) -> Self {
+        Self { request, value: None, args: Self::build_args(field), field }
     }
 
     pub fn with_value_and_field(
         &self,
         value: &'a Output,
         field: &'a Field<Nested<Input>, Input>,
     ) -> Self {
-        Self { args: None, value: Some(value), field, request: self.request }
-    }
-
-    pub fn with_args(&self, args: indexmap::IndexMap<&str, Input>) -> Self {
-        let mut map = indexmap::IndexMap::new();
-        for (key, value) in args {
-            map.insert(Name::new(key), value);
-        }
         Self {
-            value: self.value,
-            args: Some(map),
-            field: self.field,
             request: self.request,
+            args: Self::build_args(field),
+            value: Some(value),
+            field,
         }
     }
 
@@ -71,6 +64,26 @@ impl<'a, Input: Clone, Output> Context<'a, Input, Output> {
     pub fn field(&self) -> &Field<Nested<Input>, Input> {
         self.field
     }
+
+    fn build_args(field: &Field<Nested<Input>, Input>) -> Option<IndexMap<Name, Input>> {
+        let mut arg_map = IndexMap::new();
+
+        for arg in field.args.iter() {
+            let name = arg.name.as_str();
+            let value = arg
+                .value
+                .clone()
+                // TODO: default value resolution should happen in the InputResolver
+                .or_else(|| arg.default_value.clone());
+            if let Some(value) = value {
+                arg_map.insert(Name::new(name), value);
+            } else if !arg.type_of.is_nullable() {
+                // TODO: throw error here
+                todo!()
+            }
+        }
+        Some(arg_map)
+    }
 }
 
 impl<'a> ResolverContextLike for Context<'a, ConstValue, ConstValue> {

src/core/jit/error.rs

@@ -27,7 +27,7 @@ pub enum ValidationError {
     // TODO: replace with sane error message. Right now, it's defined as is only for compatibility
     // with async_graphql error message for this case
     #[error(r#"internal: invalid value for scalar "{type_of}", expected "FieldValue::Value""#)]
-    ScalarInvalid { type_of: String, path: String },
+    ScalarInvalid { type_of: String },
     #[error("TypeName shape doesn't satisfy the processed object")]
     TypeNameMismatch,
     #[error(r#"internal: invalid item for enum "{type_of}""#)]

src/core/jit/exec.rs

@@ -31,7 +31,7 @@ where
     Exec: IRExecutor<Input = Input, Output = Output, Error = jit::Error>,
 {
     pub fn new(plan: OperationPlan<Input>, exec: Exec) -> Self {
-        Self { exec, ctx: RequestContext::new(plan) }
+        Self { exec, ctx: RequestContext::new(plan.clone()) }
     }
 
     pub async fn store(&self) -> Store<Result<TypedValue<Output>, Positioned<jit::Error>>> {
@@ -73,25 +73,9 @@ where
 
     async fn init(&mut self) {
         join_all(self.request.plan().as_nested().iter().map(|field| async {
-            let mut arg_map = indexmap::IndexMap::new();
-            for arg in field.args.iter() {
-                let name = arg.name.as_str();
-                let value: Option<Input> = arg
-                    .value
-                    .clone()
-                    // TODO: default value resolution should happen in the InputResolver
-                    .or_else(|| arg.default_value.clone());
-
-                if let Some(value) = value {
-                    arg_map.insert(name, value);
-                } else if !arg.type_of.is_nullable() {
-                    // TODO: throw error here
-                    todo!()
-                }
-            }
+            let ctx = Context::new(field, self.request);
             // TODO: with_args should be called on inside iter_field on any level, not only
             // for root fields
-            let ctx = Context::new(field, self.request).with_args(arg_map);
             self.execute(&ctx, DataPath::new()).await
         }))
         .await;
@@ -176,7 +160,7 @@ where
             let default_obj = Output::object(Output::JsonObject::new());
             let value = ctx
                 .value()
-                .and_then(|v| v.get_key(&field.name))
+                .and_then(|v| v.get_key(&field.output_name))
                 // in case there is no value we still put some dumb empty value anyway
                 // to force execution of the nested fields even when parent object is not present.
                 // For async_graphql it's done by `fix_dangling_resolvers` fn that basically creates

src/core/jit/model.rs

@@ -108,7 +108,11 @@ impl FieldId {
 #[derive(Clone)]
 pub struct Field<Extensions, Input> {
     pub id: FieldId,
+    /// Name of key in the value object for this field
     pub name: String,
+    /// Output name (i.e. with alias) that should be used for the result value
+    /// of this field
+    pub output_name: String,
     pub ir: Option<IR>,
     pub type_of: crate::core::blueprint::Type,
     /// Specifies the name of type used in condition to fetch that field
@@ -158,6 +162,7 @@ impl<Input> Field<Nested<Input>, Input> {
         Ok(Field {
             id: self.id,
             name: self.name,
+            output_name: self.output_name,
             ir: self.ir,
             type_of: self.type_of,
             type_condition: self.type_condition,
@@ -187,6 +192,7 @@ impl<Input> Field<Flat, Input> {
         Ok(Field {
             id: self.id,
             name: self.name,
+            output_name: self.output_name,
             ir: self.ir,
             type_of: self.type_of,
             type_condition: self.type_condition,
@@ -262,6 +268,7 @@ impl<Input> Field<Flat, Input> {
         Field {
             id: self.id,
             name: self.name,
+            output_name: self.output_name,
             ir: self.ir,
             type_of: self.type_of,
             type_condition: self.type_condition,
@@ -280,6 +287,7 @@ impl<Extensions: Debug, Input: Debug> Debug for Field<Extensions, Input> {
         let mut debug_struct = f.debug_struct("Field");
         debug_struct.field("id", &self.id);
         debug_struct.field("name", &self.name);
+        debug_struct.field("output_name", &self.output_name);
         if self.ir.is_some() {
             debug_struct.field("ir", &"Some(..)");
         }
@@ -298,6 +306,7 @@ impl<Extensions: Debug, Input: Debug> Debug for Field<Extensions, Input> {
             debug_struct.field("include", &self.include);
         }
         debug_struct.field("directives", &self.directives);
+
         debug_struct.finish()
     }
 }

src/core/jit/snapshots/tailcall__core__jit__builder__tests__alias_query.snap

@@ -0,0 +1,44 @@
+---
+source: src/core/jit/builder.rs
+expression: plan.into_nested()
+---
+[
+    Field {
+        id: 0,
+        name: "posts",
+        output_name: "articles",
+        ir: "Some(..)",
+        type_of: [Post],
+        type_condition: "Query",
+        extensions: Some(
+            Nested(
+                [
+                    Field {
+                        id: 1,
+                        name: "user",
+                        output_name: "author",
+                        ir: "Some(..)",
+                        type_of: User,
+                        type_condition: "Post",
+                        extensions: Some(
+                            Nested(
+                                [
+                                    Field {
+                                        id: 2,
+                                        name: "id",
+                                        output_name: "identifier",
+                                        type_of: ID!,
+                                        type_condition: "User",
+                                        directives: [],
+                                    },
+                                ],
+                            ),
+                        ),
+                        directives: [],
+                    },
+                ],
+            ),
+        ),
+        directives: [],
+    },
+]