fix: resolve union type in response (#2019)

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

Cargo.toml

@@ -162,6 +162,8 @@ Inflector = "0.11.4"
 path-clean = "=1.0.1"
 pathdiff = "0.2.1"
 num = "0.4.3"
+indenter = "0.3.3"
+derive_more = "0.99.18"
 
 [dev-dependencies]
 tailcall-prettier = { path = "tailcall-prettier" }
@@ -175,6 +177,7 @@ tempfile = "3.10.1"
 temp-env = "0.3.6"
 maplit = "1.0.2"
 tailcall-fixtures = { path = "./tailcall-fixtures" }
+test-log = { version = "0.2.16", default-features = false, features = ["color", "trace"] }
 
 [features]
 

examples/jsonplaceholder.graphql

@@ -1,5 +1,5 @@
 schema
-  @server(port: 8000, headers: {cors: {allowOrigins: ["*"], allowHeaders: ["*"], allowMethods: [POST, GET, OPTIONS]}})
+  @server(port: 8000)
   @upstream(baseURL: "http://jsonplaceholder.typicode.com", httpCache: 42, batch: {delay: 100}) {
   query: Query
 }

src/core/blueprint/blueprint.rs

@@ -69,6 +69,7 @@ impl Type {
             Type::ListType { of_type, .. } => of_type.name(),
         }
     }
+
     /// checks if the type is nullable
     pub fn is_nullable(&self) -> bool {
         !match self {
@@ -171,7 +172,7 @@ pub struct FieldDefinition {
 impl FieldDefinition {
     ///
     /// Transforms the current expression if it exists on the provided field.
-    pub fn map_expr<F: FnMut(IR) -> IR>(&mut self, mut wrapper: F) {
+    pub fn map_expr<F: FnOnce(IR) -> IR>(&mut self, wrapper: F) {
         if let Some(resolver) = self.resolver.take() {
             self.resolver = Some(wrapper(resolver))
         }

src/core/blueprint/definitions.rs

@@ -2,6 +2,7 @@ use std::collections::HashSet;
 
 use async_graphql_value::ConstValue;
 use regex::Regex;
+use union_resolver::update_union_resolver;
 
 use crate::core::blueprint::Type::ListType;
 use crate::core::blueprint::*;
@@ -512,6 +513,7 @@ pub fn to_field_definition(
         .and(update_cache_resolvers())
         .and(update_protected(object_name).trace(Protected::trace_name().as_str()))
         .and(update_enum_alias())
+        .and(update_union_resolver())
         .try_fold(
             &(config_module, field, type_of, name),
             FieldDefinition::default(),

src/core/blueprint/from_config.rs

@@ -102,14 +102,7 @@ where
                 .collect::<BTreeSet<String>>(),
         )
     } else {
-        match type_of {
-            "String" => JsonSchema::Str,
-            "Int" => JsonSchema::Num,
-            "Boolean" => JsonSchema::Bool,
-            "Empty" => JsonSchema::Empty,
-            "JSON" => JsonSchema::Any,
-            _ => JsonSchema::Any,
-        }
+        JsonSchema::from_scalar_type(type_of)
     };
 
     if !required {

src/core/blueprint/into_schema.rs

@@ -1,6 +1,7 @@
 use std::borrow::Cow;
 use std::sync::Arc;
 
+use anyhow::{bail, Result};
 use async_graphql::dynamic::{self, FieldFuture, FieldValue, SchemaBuilder};
 use async_graphql::ErrorExtensions;
 use async_graphql_value::ConstValue;
@@ -9,7 +10,7 @@ use tracing::Instrument;
 
 use crate::core::blueprint::{Blueprint, Definition, Type};
 use crate::core::http::RequestContext;
-use crate::core::ir::{EvalContext, ResolverContext};
+use crate::core::ir::{EvalContext, ResolverContext, TypeName};
 use crate::core::scalar::CUSTOM_SCALARS;
 
 fn to_type_ref(type_of: &Type) -> dynamic::TypeRef {
@@ -57,6 +58,29 @@ fn set_default_value(
     }
 }
 
+fn to_field_value<'a>(
+    ctx: &mut EvalContext<'a, ResolverContext<'a>>,
+    value: async_graphql::Value,
+) -> Result<FieldValue<'static>> {
+    let type_name = ctx.type_name.take();
+
+    Ok(match (value, type_name) {
+        // NOTE: Mostly type_name is going to be None so we should keep that as the first check.
+        (value, None) => FieldValue::from(value),
+        (ConstValue::List(values), Some(TypeName::Vec(names))) => FieldValue::list(
+            values
+                .into_iter()
+                .zip(names)
+                .map(|(value, type_name)| FieldValue::from(value).with_type(type_name)),
+        ),
+        (value @ ConstValue::Object(_), Some(TypeName::Single(type_name))) => {
+            FieldValue::from(value).with_type(type_name)
+        }
+        (ConstValue::Null, _) => FieldValue::NULL,
+        (_, Some(_)) => bail!("Failed to match type_name"),
+    })
+}
+
 fn to_type(def: &Definition) -> dynamic::Type {
     match def {
         Definition::Object(def) => {
@@ -65,6 +89,7 @@ fn to_type(def: &Definition) -> dynamic::Type {
                 let field = field.clone();
                 let type_ref = to_type_ref(&field.of_type);
                 let field_name = &field.name.clone();
+
                 let mut dyn_schema_field = dynamic::Field::new(
                     field_name,
                     type_ref.clone(),
@@ -81,6 +106,7 @@ fn to_type(def: &Definition) -> dynamic::Type {
                             None => {
                                 let ctx: ResolverContext = ctx.into();
                                 let ctx = EvalContext::new(req_ctx, &ctx);
+
                                 FieldFuture::from_value(
                                     ctx.path_value(&[field_name]).map(|a| a.into_owned()),
                                 )
@@ -90,22 +116,21 @@ fn to_type(def: &Definition) -> dynamic::Type {
                                     "field_resolver",
                                     otel.name = ctx.path_node.map(|p| p.to_string()).unwrap_or(field_name.clone()), graphql.returnType = %type_ref
                                 );
+
                                 let expr = expr.to_owned();
                                 FieldFuture::new(
                                     async move {
                                         let ctx: ResolverContext = ctx.into();
-                                        let mut ctx = EvalContext::new(req_ctx, &ctx);
-
-                                        let const_value = expr
-                                            .eval(&mut ctx)
-                                            .await
-                                            .map_err(|err| err.extend())?;
-                                        let p = match const_value {
-                                            ConstValue::List(a) => Some(FieldValue::list(a)),
-                                            ConstValue::Null => FieldValue::NONE,
-                                            a => Some(FieldValue::from(a)),
-                                        };
-                                        Ok(p)
+                                        let ctx = &mut EvalContext::new(req_ctx, &ctx);
+
+                                        let value =
+                                            expr.eval(ctx).await.map_err(|err| err.extend())?;
+
+                                        if let ConstValue::Null = value {
+                                            Ok(FieldValue::NONE)
+                                        } else {
+                                            Ok(Some(to_field_value(ctx, value)?))
+                                        }
                                     }
                                     .instrument(span)
                                     .inspect_err(|err| tracing::error!(?err)),

src/core/blueprint/mod.rs

@@ -14,6 +14,7 @@ mod schema;
 mod server;
 pub mod telemetry;
 mod timeout;
+mod union_resolver;
 mod upstream;
 
 pub use auth::*;

src/core/blueprint/union_resolver.rs

@@ -0,0 +1,49 @@
+use crate::core::blueprint::FieldDefinition;
+use crate::core::config;
+use crate::core::config::{ConfigModule, Field};
+use crate::core::ir::model::IR;
+use crate::core::ir::Discriminator;
+use crate::core::try_fold::TryFold;
+use crate::core::valid::{Valid, Validator};
+
+fn compile_union_resolver(
+    config: &ConfigModule,
+    union_name: &str,
+    union_: &config::Union,
+) -> Valid<Discriminator, String> {
+    Valid::from_iter(&union_.types, |type_name| {
+        Valid::from_option(
+            config
+                .find_type(type_name)
+                .map(|type_| (type_name.as_str(), type_)),
+            "Can't find a type that is member of union type".to_string(),
+        )
+    })
+    .and_then(|types| {
+        let types: Vec<_> = types.into_iter().collect();
+
+        Discriminator::new(union_name, &types)
+    })
+}
+
+pub fn update_union_resolver<'a>(
+) -> TryFold<'a, (&'a ConfigModule, &'a Field, &'a config::Type, &'a str), FieldDefinition, String>
+{
+    TryFold::<(&ConfigModule, &Field, &config::Type, &str), FieldDefinition, String>::new(
+        |(config, field, _, _), mut b_field| {
+            let Some(union_) = config.find_union(&field.type_of) else {
+                return Valid::succeed(b_field);
+            };
+
+            compile_union_resolver(config, &field.type_of, union_).map(|discriminator| {
+                b_field.resolver = Some(
+                    b_field
+                        .resolver
+                        .unwrap_or(IR::ContextPath(vec![b_field.name.clone()])),
+                );
+                b_field.map_expr(move |expr| IR::Discriminate(discriminator, expr.into()));
+                b_field
+            })
+        },
+    )
+}