Fix graphql union query for nested types (#986)

* Fix graphql entities model union query

* rust fmt

crates/torii/graphql/Cargo.toml

@@ -30,6 +30,7 @@ dojo-types = { path = "../../dojo-types" }
 tracing.workspace = true
 url.workspace = true
 warp.workspace = true
+async-recursion = "1.0.5"
 
 [dev-dependencies]
 camino.workspace = true

crates/torii/graphql/src/object/entity.rs

@@ -2,6 +2,7 @@ use async_graphql::dynamic::{
     Field, FieldFuture, FieldValue, InputValue, SubscriptionField, SubscriptionFieldFuture, TypeRef,
 };
 use async_graphql::{Name, Value};
+use async_recursion::async_recursion;
 use indexmap::IndexMap;
 use sqlx::pool::PoolConnection;
 use sqlx::{Pool, Result, Sqlite};
@@ -13,12 +14,11 @@ use super::connection::{
     connection_arguments, connection_output, decode_cursor, parse_connection_arguments,
     ConnectionArguments,
 };
-use super::model_data::model_data_by_id_query;
+use super::model_data::value_mapping_from_row;
 use super::{ObjectTrait, TypeMapping, ValueMapping};
 use crate::constants::DEFAULT_LIMIT;
 use crate::query::{query_by_id, type_mapping_query};
 use crate::types::{GraphqlType, TypeData};
-use crate::utils::csv_to_vec;
 use crate::utils::extract_value::extract;
 
 pub struct EntityObject {
@@ -161,15 +161,24 @@ fn model_union_field() -> Field {
             match ctx.parent_value.try_to_value()? {
                 Value::Object(indexmap) => {
                     let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
-                    let model_names = csv_to_vec(&extract::<String>(indexmap, "modelNames")?);
-                    let entity_id = extract::<String>(indexmap, "id")?;
+                    let model_names: Vec<String> = extract::<String>(indexmap, "modelNames")?
+                        .split(',')
+                        .map(|s| s.to_string())
+                        .collect();
 
+                    let entity_id = extract::<String>(indexmap, "id")?;
                     let mut results: Vec<FieldValue<'_>> = Vec::new();
                     for name in model_names {
                         let type_mapping = type_mapping_query(&mut conn, &name).await?;
-                        let state =
-                            model_data_by_id_query(&mut conn, &name, &entity_id, &type_mapping)
-                                .await?;
+                        let mut path_array = vec![name.clone()];
+                        let state = model_data_recursive_query(
+                            &mut conn,
+                            &mut path_array,
+                            &entity_id,
+                            &type_mapping,
+                        )
+                        .await?;
+
                         results.push(FieldValue::with_type(FieldValue::owned_any(state), name));
                     }
 
@@ -181,6 +190,33 @@ fn model_union_field() -> Field {
     })
 }
 
+// TODO: flatten query
+#[async_recursion]
+pub async fn model_data_recursive_query(
+    conn: &mut PoolConnection<Sqlite>,
+    path_array: &mut Vec<String>,
+    entity_id: &str,
+    type_mapping: &TypeMapping,
+) -> sqlx::Result<ValueMapping> {
+    let table_name = path_array.join("$");
+    let query = format!("SELECT * FROM {} WHERE entity_id = '{}'", &table_name, entity_id);
+    let row = sqlx::query(&query).fetch_one(conn.as_mut()).await?;
+    let mut value_mapping = value_mapping_from_row(&row, type_mapping)?;
+
+    for (field_name, type_data) in type_mapping {
+        if let TypeData::Nested((nested_type_ref, nested_mapping)) = type_data {
+            path_array.push(nested_type_ref.to_string());
+
+            let nested_values =
+                model_data_recursive_query(conn, path_array, entity_id, nested_mapping).await?;
+
+            value_mapping.insert(Name::new(field_name), Value::Object(nested_values));
+        }
+    }
+
+    Ok(value_mapping)
+}
+
 async fn entities_by_sk(
     conn: &mut PoolConnection<Sqlite>,
     keys: Option<Vec<String>>,

crates/torii/graphql/src/object/model_data.rs

@@ -116,7 +116,8 @@ impl ObjectTrait for ModelDataObject {
 
     fn objects(&self) -> Vec<Object> {
         let mut path_array = vec![self.type_name().to_string()];
-        let mut objects = data_objects(self.type_name(), self.type_mapping(), &mut path_array);
+        let mut objects =
+            data_objects_recursion(self.type_name(), self.type_mapping(), &mut path_array);
 
         // root object requires entity_field association
         let mut root = objects.pop().unwrap();
@@ -127,7 +128,7 @@ impl ObjectTrait for ModelDataObject {
     }
 }
 
-fn data_objects(
+fn data_objects_recursion(
     type_name: &str,
     type_mapping: &TypeMapping,
     path_array: &mut Vec<String>,
@@ -137,7 +138,7 @@ fn data_objects(
     for (_, type_data) in type_mapping {
         if let TypeData::Nested((nested_type, nested_mapping)) = type_data {
             path_array.push(nested_type.to_string());
-            objects.extend(data_objects(
+            objects.extend(data_objects_recursion(
                 &nested_type.to_string(),
                 nested_mapping,
                 &mut path_array.clone(),
@@ -160,33 +161,31 @@ pub fn object(type_name: &str, type_mapping: &TypeMapping, path_array: &[String]
             let type_data = type_data.clone();
             let table_name = table_name.clone();
 
-            // Field resolver for nested types
-            if let TypeData::Nested((_, nested_mapping)) = type_data {
-                return FieldFuture::new(async move {
-                    match ctx.parent_value.try_to_value()? {
-                        Value::Object(indexmap) => {
-                            let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
-                            let entity_id = extract::<String>(indexmap, "entity_id")?;
-
-                            // TODO: remove subqueries and use JOIN in parent query
-                            let result = model_data_by_id_query(
-                                &mut conn,
-                                &table_name,
-                                &entity_id,
-                                &nested_mapping,
-                            )
-                            .await?;
-
-                            Ok(Some(Value::Object(result)))
-                        }
-                        _ => Err("incorrect value, requires Value::Object".into()),
+            return FieldFuture::new(async move {
+                if let Some(value) = ctx.parent_value.as_value() {
+                    // Nested types resolution
+                    if let TypeData::Nested((_, nested_mapping)) = type_data {
+                        return match ctx.parent_value.try_to_value()? {
+                            Value::Object(indexmap) => {
+                                let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
+                                let entity_id = extract::<String>(indexmap, "entity_id")?;
+
+                                // TODO: remove subqueries and use JOIN in parent query
+                                let result = model_data_by_id_query(
+                                    &mut conn,
+                                    &table_name,
+                                    &entity_id,
+                                    &nested_mapping,
+                                )
+                                .await?;
+
+                                Ok(Some(Value::Object(result)))
+                            }
+                            _ => Err("incorrect value, requires Value::Object".into()),
+                        };
                     }
-                });
-            }
 
-            // Field resolver for simple types and model union
-            FieldFuture::new(async move {
-                if let Some(value) = ctx.parent_value.as_value() {
+                    // Simple types resolution
                     return match value {
                         Value::Object(value_mapping) => {
                             Ok(Some(value_mapping.get(&field_name).unwrap().clone()))
@@ -202,7 +201,7 @@ pub fn object(type_name: &str, type_mapping: &TypeMapping, path_array: &[String]
                 }
 
                 Err("Field resolver only accepts Value or IndexMap".into())
-            })
+            });
         });
 
         object = object.field(field);
@@ -338,7 +337,7 @@ pub fn model_connection(
     ]))
 }
 
-fn value_mapping_from_row(row: &SqliteRow, types: &TypeMapping) -> sqlx::Result<ValueMapping> {
+pub fn value_mapping_from_row(row: &SqliteRow, types: &TypeMapping) -> sqlx::Result<ValueMapping> {
     let mut value_mapping = types
         .iter()
         .filter(|(_, type_data)| type_data.is_simple())

crates/torii/graphql/src/utils/mod.rs

@@ -1,7 +1,3 @@
 pub mod extract_value;
 pub mod parse_argument;
 pub mod value_accessor;
-
-pub fn csv_to_vec(csv: &str) -> Vec<String> {
-    csv.split(',').map(|s| s.trim().to_string()).collect()
-}