Add support for multi-object graphql queries (#436)

crates/torii/src/graphql/constants.rs

@@ -1 +1 @@
-pub const _DEFAULT_LIMIT: usize = 10;
+pub const DEFAULT_LIMIT: u64 = 10;

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

@@ -3,11 +3,12 @@ use async_graphql::{Name, Value};
 use chrono::{DateTime, Utc};
 use indexmap::IndexMap;
 use serde::Deserialize;
-use sqlx::pool::PoolConnection;
-use sqlx::{FromRow, Pool, Result, Sqlite};
+use sqlx::{FromRow, Pool, Sqlite};
 
+use super::query::{query_all, query_by_id, ID};
 use super::storage::{storage_by_name, type_mapping_from};
 use super::{ObjectTrait, TypeMapping, ValueMapping};
+use crate::graphql::constants::DEFAULT_LIMIT;
 use crate::graphql::types::ScalarType;
 use crate::graphql::utils::extract_value::extract;
 use crate::graphql::utils::{format_name, remove_quotes};
@@ -52,6 +53,21 @@ impl ComponentObject {
             storage_names,
         }
     }
+
+    pub fn value_mapping(component: Component) -> ValueMapping {
+        IndexMap::from([
+            (Name::new("id"), Value::from(component.id)),
+            (Name::new("name"), Value::from(component.name)),
+            (Name::new("classHash"), Value::from(component.class_hash)),
+            (Name::new("transactionHash"), Value::from(component.transaction_hash)),
+            (
+                Name::new("createdAt"),
+                Value::from(
+                    component.created_at.to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
+                ),
+            ),
+        ])
+    }
 }
 
 impl ObjectTrait for ComponentObject {
@@ -75,7 +91,7 @@ impl ObjectTrait for ComponentObject {
     }
 
     fn nested_fields(&self) -> Option<Vec<Field>> {
-        Some(vec![Field::new("storage", TypeRef::named("Storage"), |ctx| {
+        Some(vec![Field::new("storage", TypeRef::named("Storage"), move |ctx| {
             FieldFuture::new(async move {
                 let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
                 let component_values = ctx.parent_value.try_downcast_ref::<ValueMapping>()?;
@@ -97,39 +113,33 @@ impl ObjectTrait for ComponentObject {
                 FieldFuture::new(async move {
                     let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
                     let id = remove_quotes(ctx.args.try_get("id")?.string()?);
-                    let component_values = component_by_id(&mut conn, &id).await?;
-                    Ok(Some(FieldValue::owned_any(component_values)))
+                    let component = query_by_id(&mut conn, "components", ID::Str(id)).await?;
+                    let result = ComponentObject::value_mapping(component);
+                    Ok(Some(FieldValue::owned_any(result)))
                 })
             })
             .argument(InputValue::new("id", TypeRef::named_nn(TypeRef::ID))),
+            Field::new("components", TypeRef::named_list(self.type_name()), |ctx| {
+                FieldFuture::new(async move {
+                    let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
+                    let limit = ctx
+                        .args
+                        .try_get("limit")
+                        .and_then(|limit| limit.u64())
+                        .unwrap_or(DEFAULT_LIMIT);
+
+                    let components: Vec<Component> =
+                        query_all(&mut conn, "components", limit).await?;
+                    let result: Vec<FieldValue<'_>> = components
+                        .into_iter()
+                        .map(ComponentObject::value_mapping)
+                        .map(FieldValue::owned_any)
+                        .collect();
+
+                    Ok(Some(FieldValue::list(result)))
+                })
+            })
+            .argument(InputValue::new("limit", TypeRef::named(TypeRef::INT))),
         ]
     }
 }
-
-async fn component_by_id(conn: &mut PoolConnection<Sqlite>, id: &str) -> Result<ValueMapping> {
-    let component: Component =
-        sqlx::query_as("SELECT * FROM components WHERE id = $1").bind(id).fetch_one(conn).await?;
-
-    Ok(value_mapping(component))
-}
-
-#[allow(dead_code)]
-pub async fn components(conn: &mut PoolConnection<Sqlite>) -> Result<Vec<ValueMapping>> {
-    let components: Vec<Component> =
-        sqlx::query_as("SELECT * FROM components").fetch_all(conn).await?;
-
-    Ok(components.into_iter().map(value_mapping).collect())
-}
-
-fn value_mapping(component: Component) -> ValueMapping {
-    IndexMap::from([
-        (Name::new("id"), Value::from(component.id)),
-        (Name::new("name"), Value::from(component.name)),
-        (Name::new("classHash"), Value::from(component.class_hash)),
-        (Name::new("transactionHash"), Value::from(component.transaction_hash)),
-        (
-            Name::new("createdAt"),
-            Value::from(component.created_at.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)),
-        ),
-    ])
-}

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

@@ -4,9 +4,11 @@ use chrono::{DateTime, Utc};
 use indexmap::IndexMap;
 use serde::Deserialize;
 use sqlx::pool::PoolConnection;
-use sqlx::{FromRow, Pool, Result, Sqlite};
+use sqlx::{FromRow, Pool, QueryBuilder, Result, Sqlite};
 
+use super::query::{query_by_id, ID};
 use super::{ObjectTrait, TypeMapping, ValueMapping};
+use crate::graphql::constants::DEFAULT_LIMIT;
 use crate::graphql::types::ScalarType;
 use crate::graphql::utils::remove_quotes;
 
@@ -38,6 +40,20 @@ impl EntityObject {
             ]),
         }
     }
+
+    pub fn value_mapping(entity: Entity) -> ValueMapping {
+        IndexMap::from([
+            (Name::new("id"), Value::from(entity.id)),
+            (Name::new("name"), Value::from(entity.name)),
+            (Name::new("partitionId"), Value::from(entity.partition_id)),
+            (Name::new("keys"), Value::from(entity.keys.unwrap_or_default())),
+            (Name::new("transactionHash"), Value::from(entity.transaction_hash)),
+            (
+                Name::new("createdAt"),
+                Value::from(entity.created_at.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)),
+            ),
+        ])
+    }
 }
 
 impl ObjectTrait for EntityObject {
@@ -59,32 +75,62 @@ impl ObjectTrait for EntityObject {
                 FieldFuture::new(async move {
                     let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
                     let id = remove_quotes(ctx.args.try_get("id")?.string()?);
-                    let entity_values = entity_by_id(&mut conn, &id).await?;
-                    Ok(Some(FieldValue::owned_any(entity_values)))
+                    let entity = query_by_id(&mut conn, "entities", ID::Str(id)).await?;
+                    let result = EntityObject::value_mapping(entity);
+                    Ok(Some(FieldValue::owned_any(result)))
                 })
             })
             .argument(InputValue::new("id", TypeRef::named_nn(TypeRef::ID))),
+            Field::new("entities", TypeRef::named_nn_list_nn(self.type_name()), |ctx| {
+                FieldFuture::new(async move {
+                    let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
+                    let parition_id = remove_quotes(ctx.args.try_get("partitionId")?.string()?);
+
+                    // handle optional keys argument
+                    let maybe_keys = ctx.args.try_get("keys").ok();
+                    let keys_arr = if let Some(keys_val) = maybe_keys {
+                        keys_val
+                            .list()?
+                            .iter()
+                            .map(|val| val.string().ok().map(remove_quotes))
+                            .collect()
+                    } else {
+                        None
+                    };
+
+                    let limit = ctx
+                        .args
+                        .try_get("limit")
+                        .and_then(|limit| limit.u64())
+                        .unwrap_or(DEFAULT_LIMIT);
+
+                    let entities = entities_by_sk(&mut conn, &parition_id, keys_arr, limit).await?;
+                    Ok(Some(FieldValue::list(entities.into_iter().map(FieldValue::owned_any))))
+                })
+            })
+            .argument(InputValue::new("partitionId", TypeRef::named_nn(ScalarType::FELT)))
+            .argument(InputValue::new("keys", TypeRef::named_list(TypeRef::STRING)))
+            .argument(InputValue::new("limit", TypeRef::named(TypeRef::INT))),
         ]
     }
 }
 
-async fn entity_by_id(conn: &mut PoolConnection<Sqlite>, id: &str) -> Result<ValueMapping> {
-    let entity: Entity =
-        sqlx::query_as("SELECT * FROM entities WHERE id = $1").bind(id).fetch_one(conn).await?;
+async fn entities_by_sk(
+    conn: &mut PoolConnection<Sqlite>,
+    partition_id: &str,
+    keys: Option<Vec<String>>,
+    limit: u64,
+) -> Result<Vec<ValueMapping>> {
+    let mut builder: QueryBuilder<'_, Sqlite> = QueryBuilder::new("SELECT * FROM entities");
+    builder.push(" WHERE partition_id = ").push_bind(partition_id);
 
-    Ok(value_mapping(entity))
-}
+    if let Some(keys) = keys {
+        let keys_str = format!("{}%", keys.join("/"));
+        builder.push(" AND keys LIKE ").push_bind(keys_str);
+    }
+
+    builder.push(" ORDER BY created_at DESC LIMIT ").push(limit);
 
-fn value_mapping(entity: Entity) -> ValueMapping {
-    IndexMap::from([
-        (Name::new("id"), Value::from(entity.id)),
-        (Name::new("name"), Value::from(entity.name)),
-        (Name::new("partitionId"), Value::from(entity.partition_id)),
-        (Name::new("keys"), Value::from(entity.keys.unwrap_or_default())),
-        (Name::new("transactionHash"), Value::from(entity.transaction_hash)),
-        (
-            Name::new("createdAt"),
-            Value::from(entity.created_at.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)),
-        ),
-    ])
+    let entities: Vec<Entity> = builder.build_query_as().fetch_all(conn).await?;
+    Ok(entities.into_iter().map(EntityObject::value_mapping).collect())
 }

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

@@ -3,11 +3,12 @@ use async_graphql::{Name, Value};
 use chrono::{DateTime, Utc};
 use indexmap::IndexMap;
 use serde::Deserialize;
-use sqlx::pool::PoolConnection;
-use sqlx::{FromRow, Pool, Result, Sqlite};
+use sqlx::{FromRow, Pool, Sqlite};
 
-use super::system_call::system_call_by_id;
+use super::query::{query_all, query_by_id, ID};
+use super::system_call::{SystemCall, SystemCallObject};
 use super::{ObjectTrait, TypeMapping, ValueMapping};
+use crate::graphql::constants::DEFAULT_LIMIT;
 use crate::graphql::types::ScalarType;
 use crate::graphql::utils::extract_value::extract;
 use crate::graphql::utils::remove_quotes;
@@ -38,6 +39,19 @@ impl EventObject {
             ]),
         }
     }
+
+    pub fn value_mapping(event: Event) -> ValueMapping {
+        IndexMap::from([
+            (Name::new("id"), Value::from(event.id)),
+            (Name::new("keys"), Value::from(event.keys)),
+            (Name::new("data"), Value::from(event.data)),
+            (Name::new("systemCallId"), Value::from(event.system_call_id)),
+            (
+                Name::new("createdAt"),
+                Value::from(event.created_at.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)),
+            ),
+        ])
+    }
 }
 
 impl ObjectTrait for EventObject {
@@ -59,12 +73,32 @@ impl ObjectTrait for EventObject {
                 FieldFuture::new(async move {
                     let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
                     let id = remove_quotes(ctx.args.try_get("id")?.string()?);
-                    let event_values = event_by_id(&mut conn, &id).await?;
-
-                    Ok(Some(FieldValue::owned_any(event_values)))
+                    let event = query_by_id(&mut conn, "events", ID::Str(id)).await?;
+                    let result = EventObject::value_mapping(event);
+                    Ok(Some(FieldValue::owned_any(result)))
                 })
             })
             .argument(InputValue::new("id", TypeRef::named_nn(TypeRef::ID))),
+            Field::new("events", TypeRef::named_list(self.type_name()), |ctx| {
+                FieldFuture::new(async move {
+                    let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
+                    let limit = ctx
+                        .args
+                        .try_get("limit")
+                        .and_then(|limit| limit.u64())
+                        .unwrap_or(DEFAULT_LIMIT);
+
+                    let events: Vec<Event> = query_all(&mut conn, "events", limit).await?;
+                    let result: Vec<FieldValue<'_>> = events
+                        .into_iter()
+                        .map(EventObject::value_mapping)
+                        .map(FieldValue::owned_any)
+                        .collect();
+
+                    Ok(Some(FieldValue::list(result)))
+                })
+            })
+            .argument(InputValue::new("limit", TypeRef::named(TypeRef::INT))),
         ]
     }
 
@@ -74,30 +108,11 @@ impl ObjectTrait for EventObject {
                 let mut conn = ctx.data::<Pool<Sqlite>>()?.acquire().await?;
                 let event_values = ctx.parent_value.try_downcast_ref::<ValueMapping>()?;
                 let syscall_id = extract::<i64>(event_values, "system_call_id")?;
-                let system_call = system_call_by_id(&mut conn, syscall_id).await?;
-
-                Ok(Some(FieldValue::owned_any(system_call)))
+                let system_call: SystemCall =
+                    query_by_id(&mut conn, "system_calls", ID::I64(syscall_id)).await?;
+                let result = SystemCallObject::value_mapping(system_call);
+                Ok(Some(FieldValue::owned_any(result)))
             })
         })])
     }
 }
-
-async fn event_by_id(conn: &mut PoolConnection<Sqlite>, id: &str) -> Result<ValueMapping> {
-    let event: Event =
-        sqlx::query_as("SELECT * FROM events WHERE id = $1").bind(id).fetch_one(conn).await?;
-
-    Ok(value_mapping(event))
-}
-
-fn value_mapping(event: Event) -> ValueMapping {
-    IndexMap::from([
-        (Name::new("id"), Value::from(event.id)),
-        (Name::new("keys"), Value::from(event.keys)),
-        (Name::new("data"), Value::from(event.data)),
-        (Name::new("systemCallId"), Value::from(event.system_call_id)),
-        (
-            Name::new("createdAt"),
-            Value::from(event.created_at.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)),
-        ),
-    ])
-}

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

@@ -1,6 +1,7 @@
 pub mod component;
 pub mod entity;
 pub mod event;
+mod query;
 pub mod storage;
 pub mod system;
 pub mod system_call;

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

@@ -0,0 +1,38 @@
+use sqlx::pool::PoolConnection;
+use sqlx::sqlite::SqliteRow;
+use sqlx::{FromRow, QueryBuilder, Result, Sqlite};
+
+pub enum ID {
+    Str(String),
+    I64(i64),
+}
+
+pub async fn query_by_id<T>(
+    conn: &mut PoolConnection<Sqlite>,
+    table_name: &str,
+    id: ID,
+) -> Result<T>
+where
+    T: Send + Unpin + for<'a> FromRow<'a, SqliteRow>,
+{
+    let query = format!("SELECT * FROM {} WHERE id = ?", table_name);
+    let result = match id {
+        ID::Str(id) => sqlx::query_as::<_, T>(&query).bind(id).fetch_one(conn).await?,
+        ID::I64(id) => sqlx::query_as::<_, T>(&query).bind(id).fetch_one(conn).await?,
+    };
+    Ok(result)
+}
+
+pub async fn query_all<T>(
+    conn: &mut PoolConnection<Sqlite>,
+    table_name: &str,
+    limit: u64,
+) -> Result<Vec<T>>
+where
+    T: Send + Unpin + for<'a> FromRow<'a, SqliteRow>,
+{
+    let mut builder: QueryBuilder<'_, Sqlite> = QueryBuilder::new("SELECT * FROM ");
+    builder.push(table_name).push(" ORDER BY created_at DESC LIMIT ").push(limit);
+    let results: Vec<T> = builder.build_query_as().fetch_all(conn).await?;
+    Ok(results)
+}

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

@@ -57,7 +57,7 @@ pub async fn storage_by_name(
     name: &str,
     fields: &TypeMapping,
 ) -> Result<ValueMapping> {
-    let query = format!("SELECT * FROM {}", name);
+    let query = format!("SELECT * FROM {} ORDER BY created_at DESC LIMIT 1", name);
     let storage = sqlx::query(&query).fetch_one(conn).await?;
     let result = value_mapping_from_row(&storage, fields)?;
     Ok(result)