Torii grpc entities query (#1196)

* Refactor torii grpc

* error freeeee!

* Refactor torii grpc

* rebase changes

* clean up

* Move model to individual query clause

* Refactor subscription argument to keys

* Add Torii grpc retrieve entities endpoint

* Add cache to store schema ty and query

* map rows to proto model

* add grpc limit offset

* refactor query logic

* clippy & fmt

* remove unused

* fix tests

* move limit offset to query clause

---------

Co-authored-by: Tarrence van As <[email protected]>

crates/dojo-types/src/primitive.rs

@@ -33,6 +33,8 @@ pub enum PrimitiveError {
     NotEnoughFieldElements,
     #[error("Unsupported CairoType for SQL formatting")]
     UnsupportedType,
+    #[error("Set value type mismatch")]
+    TypeMismatch,
     #[error(transparent)]
     ValueOutOfRange(#[from] ValueOutOfRangeError),
 }
@@ -44,97 +46,60 @@ pub enum SqlType {
     Text,
 }
 
-impl Primitive {
-    /// If the `Primitive` is a u8, returns the associated [`u8`]. Returns `None` otherwise.
-    pub fn as_u8(&self) -> Option<u8> {
-        match self {
-            Primitive::U8(value) => *value,
-            _ => None,
-        }
-    }
-
-    /// If the `Primitive` is a u16, returns the associated [`u16`]. Returns `None` otherwise.
-    pub fn as_u16(&self) -> Option<u16> {
-        match self {
-            Primitive::U16(value) => *value,
-            _ => None,
-        }
-    }
-
-    /// If the `Primitive` is a u32, returns the associated [`u32`]. Returns `None` otherwise.
-    pub fn as_u32(&self) -> Option<u32> {
-        match self {
-            Primitive::U32(value) => *value,
-            _ => None,
-        }
-    }
-
-    /// If the `Primitive` is a u64, returns the associated [`u64`]. Returns `None` otherwise.
-    pub fn as_u64(&self) -> Option<u64> {
-        match self {
-            Primitive::U64(value) => *value,
-            _ => None,
-        }
-    }
-
-    /// If the `Primitive` is a u128, returns the associated [`u128`]. Returns `None` otherwise.
-    pub fn as_u128(&self) -> Option<u128> {
-        match self {
-            Primitive::U128(value) => *value,
-            _ => None,
-        }
-    }
-
-    /// If the `Primitive` is a u256, returns the associated [`U256`]. Returns `None` otherwise.
-    pub fn as_u256(&self) -> Option<U256> {
-        match self {
-            Primitive::U256(value) => *value,
-            _ => None,
-        }
-    }
-
-    /// If the `Primitive` is a felt252, returns the associated [`FieldElement`]. Returns `None`
-    /// otherwise.
-    pub fn as_felt252(&self) -> Option<FieldElement> {
-        match self {
-            Primitive::Felt252(value) => *value,
-            _ => None,
-        }
-    }
-
-    /// If the `Primitive` is a ClassHash, returns the associated [`FieldElement`]. Returns `None`
-    /// otherwise.
-    pub fn as_class_hash(&self) -> Option<FieldElement> {
-        match self {
-            Primitive::ClassHash(value) => *value,
-            _ => None,
+/// Macro to generate setter methods for Primitive enum variants.
+macro_rules! set_primitive {
+    ($method_name:ident, $variant:ident, $type:ty) => {
+        /// Sets the inner value of the `Primitive` enum if variant matches.
+        pub fn $method_name(&mut self, value: Option<$type>) -> Result<(), PrimitiveError> {
+            match self {
+                Primitive::$variant(_) => {
+                    *self = Primitive::$variant(value);
+                    Ok(())
+                }
+                _ => Err(PrimitiveError::TypeMismatch),
+            }
         }
-    }
+    };
+}
 
-    /// If the `Primitive` is a ContractAddress, returns the associated [`FieldElement`]. Returns
-    /// `None` otherwise.
-    pub fn as_contract_address(&self) -> Option<FieldElement> {
-        match self {
-            Primitive::ContractAddress(value) => *value,
-            _ => None,
+/// Macro to generate getter methods for Primitive enum variants.
+macro_rules! as_primitive {
+    ($method_name:ident, $variant:ident, $type:ty) => {
+        /// If the `Primitive` is variant type, returns the associated vartiant value. Returns
+        /// `None` otherwise.
+        pub fn $method_name(&self) -> Option<$type> {
+            match self {
+                Primitive::$variant(value) => *value,
+                _ => None,
+            }
         }
-    }
+    };
+}
 
-    /// If the `Primitive` is a usize, returns the associated [`u32`]. Returns `None` otherwise.
-    pub fn as_usize(&self) -> Option<u32> {
-        match self {
-            Primitive::USize(value) => *value,
-            _ => None,
-        }
-    }
+impl Primitive {
+    as_primitive!(as_u8, U8, u8);
+    as_primitive!(as_u16, U16, u16);
+    as_primitive!(as_u32, U32, u32);
+    as_primitive!(as_u64, U64, u64);
+    as_primitive!(as_u128, U128, u128);
+    as_primitive!(as_u256, U256, U256);
+    as_primitive!(as_bool, Bool, bool);
+    as_primitive!(as_usize, USize, u32);
+    as_primitive!(as_felt252, Felt252, FieldElement);
+    as_primitive!(as_class_hash, ClassHash, FieldElement);
+    as_primitive!(as_contract_address, ContractAddress, FieldElement);
 
-    /// If the `Primitive` is a bool, returns the associated [`bool`]. Returns `None` otherwise.
-    pub fn as_bool(&self) -> Option<bool> {
-        match self {
-            Primitive::Bool(value) => *value,
-            _ => None,
-        }
-    }
+    set_primitive!(set_u8, U8, u8);
+    set_primitive!(set_u16, U16, u16);
+    set_primitive!(set_u32, U32, u32);
+    set_primitive!(set_u64, U64, u64);
+    set_primitive!(set_u128, U128, u128);
+    set_primitive!(set_u256, U256, U256);
+    set_primitive!(set_bool, Bool, bool);
+    set_primitive!(set_usize, USize, u32);
+    set_primitive!(set_felt252, Felt252, FieldElement);
+    set_primitive!(set_class_hash, ClassHash, FieldElement);
+    set_primitive!(set_contract_address, ContractAddress, FieldElement);
 
     pub fn to_sql_type(&self) -> SqlType {
         match self {
@@ -333,28 +298,39 @@ mod tests {
     }
 
     #[test]
-    fn as_inner_value() {
-        let primitive = Primitive::U8(Some(1u8));
+    fn inner_value_getter_setter() {
+        let mut primitive = Primitive::U8(None);
+        primitive.set_u8(Some(1u8)).unwrap();
         assert_eq!(primitive.as_u8(), Some(1u8));
-        let primitive = Primitive::U16(Some(1u16));
+        let mut primitive = Primitive::U16(None);
+        primitive.set_u16(Some(1u16)).unwrap();
         assert_eq!(primitive.as_u16(), Some(1u16));
-        let primitive = Primitive::U32(Some(1u32));
+        let mut primitive = Primitive::U32(None);
+        primitive.set_u32(Some(1u32)).unwrap();
         assert_eq!(primitive.as_u32(), Some(1u32));
-        let primitive = Primitive::U64(Some(1u64));
+        let mut primitive = Primitive::U64(None);
+        primitive.set_u64(Some(1u64)).unwrap();
         assert_eq!(primitive.as_u64(), Some(1u64));
-        let primitive = Primitive::U128(Some(1u128));
+        let mut primitive = Primitive::U128(None);
+        primitive.set_u128(Some(1u128)).unwrap();
         assert_eq!(primitive.as_u128(), Some(1u128));
-        let primitive = Primitive::U256(Some(U256::from(1u128)));
+        let mut primitive = Primitive::U256(None);
+        primitive.set_u256(Some(U256::from(1u128))).unwrap();
         assert_eq!(primitive.as_u256(), Some(U256::from(1u128)));
-        let primitive = Primitive::USize(Some(1u32));
+        let mut primitive = Primitive::USize(None);
+        primitive.set_usize(Some(1u32)).unwrap();
         assert_eq!(primitive.as_usize(), Some(1u32));
-        let primitive = Primitive::Bool(Some(true));
+        let mut primitive = Primitive::Bool(None);
+        primitive.set_bool(Some(true)).unwrap();
         assert_eq!(primitive.as_bool(), Some(true));
-        let primitive = Primitive::Felt252(Some(FieldElement::from(1u128)));
+        let mut primitive = Primitive::Felt252(None);
+        primitive.set_felt252(Some(FieldElement::from(1u128))).unwrap();
         assert_eq!(primitive.as_felt252(), Some(FieldElement::from(1u128)));
-        let primitive = Primitive::ClassHash(Some(FieldElement::from(1u128)));
+        let mut primitive = Primitive::ClassHash(None);
+        primitive.set_class_hash(Some(FieldElement::from(1u128))).unwrap();
         assert_eq!(primitive.as_class_hash(), Some(FieldElement::from(1u128)));
-        let primitive = Primitive::ContractAddress(Some(FieldElement::from(1u128)));
+        let mut primitive = Primitive::ContractAddress(None);
+        primitive.set_contract_address(Some(FieldElement::from(1u128))).unwrap();
         assert_eq!(primitive.as_contract_address(), Some(FieldElement::from(1u128)));
     }
 }

crates/dojo-types/src/schema.rs

@@ -269,6 +269,16 @@ impl Enum {
         Ok(self.options[option].name.clone())
     }
 
+    pub fn set_option(&mut self, name: &str) -> Result<(), EnumError> {
+        match self.options.iter().position(|option| option.name == name) {
+            Some(index) => {
+                self.option = Some(index as u8);
+                Ok(())
+            }
+            None => Err(EnumError::OptionInvalid),
+        }
+    }
+
     pub fn to_sql_value(&self) -> Result<String, EnumError> {
         self.option()
     }

crates/torii/core/Cargo.toml

@@ -13,6 +13,7 @@ anyhow.workspace = true
 async-trait.workspace = true
 base64.workspace = true
 chrono.workspace = true
+crypto-bigint = { version = "0.5.3", features = [ "serde" ] }
 dojo-types = { path = "../../dojo-types" }
 dojo-world = { path = "../../dojo-world", features = [ "contracts", "manifest" ] }
 futures-channel = "0.3.0"

crates/torii/core/src/cache.rs

@@ -0,0 +1,56 @@
+use std::collections::HashMap;
+
+use dojo_types::schema::Ty;
+use sqlx::SqlitePool;
+use tokio::sync::RwLock;
+
+use crate::error::{Error, QueryError};
+use crate::model::{parse_sql_model_members, SqlModelMember};
+
+type ModelName = String;
+
+pub struct ModelCache {
+    pool: SqlitePool,
+    schemas: RwLock<HashMap<ModelName, Ty>>,
+}
+
+impl ModelCache {
+    pub fn new(pool: SqlitePool) -> Self {
+        Self { pool, schemas: RwLock::new(HashMap::new()) }
+    }
+
+    pub async fn schema(&self, model: &str) -> Result<Ty, Error> {
+        {
+            let schemas = self.schemas.read().await;
+            if let Some(schema) = schemas.get(model) {
+                return Ok(schema.clone());
+            }
+        }
+
+        self.update_schema(model).await
+    }
+
+    async fn update_schema(&self, model: &str) -> Result<Ty, Error> {
+        let model_members: Vec<SqlModelMember> = sqlx::query_as(
+            "SELECT id, model_idx, member_idx, name, type, type_enum, enum_options, key FROM \
+             model_members WHERE model_id = ? ORDER BY model_idx ASC, member_idx ASC",
+        )
+        .bind(model)
+        .fetch_all(&self.pool)
+        .await?;
+
+        if model_members.is_empty() {
+            return Err(QueryError::ModelNotFound(model.into()).into());
+        }
+
+        let ty = parse_sql_model_members(model, &model_members);
+        let mut schemas = self.schemas.write().await;
+        schemas.insert(model.into(), ty.clone());
+
+        Ok(ty)
+    }
+
+    pub async fn clear(&self) {
+        self.schemas.write().await.clear();
+    }
+}

crates/torii/core/src/error.rs

@@ -1,3 +1,7 @@
+use std::num::ParseIntError;
+
+use dojo_types::primitive::PrimitiveError;
+use dojo_types::schema::EnumError;
 use starknet::core::types::{FromByteSliceError, FromStrError};
 use starknet::core::utils::CairoShortStringToFeltError;
 
@@ -9,6 +13,10 @@ pub enum Error {
     Sql(#[from] sqlx::Error),
     #[error(transparent)]
     QueryError(#[from] QueryError),
+    #[error(transparent)]
+    PrimitiveError(#[from] PrimitiveError),
+    #[error(transparent)]
+    EnumError(#[from] EnumError),
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -19,10 +27,16 @@ pub enum ParseError {
     CairoShortStringToFelt(#[from] CairoShortStringToFeltError),
     #[error(transparent)]
     FromByteSliceError(#[from] FromByteSliceError),
+    #[error(transparent)]
+    ParseIntError(#[from] ParseIntError),
 }
 
 #[derive(Debug, thiserror::Error)]
 pub enum QueryError {
     #[error("unsupported query")]
     UnsupportedQuery,
+    #[error("model not found: {0}")]
+    ModelNotFound(String),
+    #[error("exceeds sqlite `JOIN` limit (64)")]
+    SqliteJoinLimit,
 }

crates/torii/core/src/lib.rs

@@ -3,6 +3,7 @@ use sqlx::FromRow;
 
 use crate::types::SQLFieldElement;
 
+pub mod cache;
 pub mod engine;
 pub mod error;
 pub mod model;