Torii client hook up retrieve entities (#1223)

* Torii client hook up retrieve entities

* mirror dojo_types schema in protobuf

* add more protobuf messages

* fix wasm compile error

* fix test

Cargo.toml

@@ -66,6 +66,7 @@ clap_complete = "4.3"
 colored = "2"
 console = "0.15.7"
 convert_case = "0.6.0"
+crypto-bigint = { version = "0.5.3", features = [ "serde" ] }
 env_logger = "0.10.0"
 flate2 = "1.0.24"
 futures = "0.3.28"

crates/dojo-types/src/primitive.rs

@@ -1,10 +1,22 @@
 use crypto_bigint::{Encoding, U256};
 use serde::{Deserialize, Serialize};
 use starknet::core::types::{FieldElement, ValueOutOfRangeError};
+use strum::IntoEnumIterator;
 use strum_macros::{AsRefStr, Display, EnumIter, EnumString};
 
 #[derive(
-    AsRefStr, Display, EnumIter, EnumString, Copy, Clone, Debug, Serialize, Deserialize, PartialEq,
+    AsRefStr,
+    Display,
+    EnumIter,
+    EnumString,
+    Copy,
+    Clone,
+    Debug,
+    Serialize,
+    Deserialize,
+    PartialEq,
+    Hash,
+    Eq,
 )]
 #[serde(tag = "scalar_type", content = "value")]
 #[strum(serialize_all = "lowercase")]
@@ -101,6 +113,14 @@ impl Primitive {
     set_primitive!(set_class_hash, ClassHash, FieldElement);
     set_primitive!(set_contract_address, ContractAddress, FieldElement);
 
+    pub fn to_numeric(&self) -> usize {
+        Self::iter().position(|p| p == *self).unwrap()
+    }
+
+    pub fn from_numeric(value: usize) -> Option<Self> {
+        Self::iter().nth(value)
+    }
+
     pub fn to_sql_type(&self) -> SqlType {
         match self {
             Primitive::U8(_)

crates/dojo-types/src/schema.rs

@@ -6,7 +6,7 @@ use strum_macros::AsRefStr;
 use crate::primitive::{Primitive, PrimitiveError};
 
 /// Represents a model member.
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Hash, Eq)]
 pub struct Member {
     pub name: String,
     #[serde(rename = "member_type")]
@@ -31,7 +31,7 @@ pub struct ModelMetadata {
 }
 
 /// Represents all possible types in Cairo
-#[derive(AsRefStr, Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(AsRefStr, Clone, Debug, Serialize, Deserialize, PartialEq, Hash, Eq)]
 #[serde(tag = "type", content = "content")]
 #[serde(rename_all = "lowercase")]
 pub enum Ty {
@@ -216,7 +216,7 @@ impl std::fmt::Display for Ty {
     }
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Hash, Eq)]
 pub struct Struct {
     pub name: String,
     pub children: Vec<Member>,
@@ -241,14 +241,14 @@ pub enum EnumError {
     OptionInvalid,
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Hash, Eq)]
 pub struct Enum {
     pub name: String,
     pub option: Option<u8>,
     pub options: Vec<EnumOption>,
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Hash, Eq)]
 pub struct EnumOption {
     pub name: String,
     pub ty: Ty,

crates/torii/client/Cargo.toml

@@ -7,7 +7,7 @@ version.workspace = true
 
 [dependencies]
 async-trait.workspace = true
-crypto-bigint = "0.5.3"
+crypto-bigint.workspace = true
 dojo-types = { path = "../../dojo-types" }
 dojo-world = { path = "../../dojo-world", features = [ "contracts" ] }
 futures-util = "0.3.28"

crates/torii/client/src/client/mod.rs

@@ -17,7 +17,8 @@ use starknet::providers::JsonRpcClient;
 use starknet_crypto::FieldElement;
 use tokio::sync::RwLock as AsyncRwLock;
 use torii_grpc::client::EntityUpdateStreaming;
-use torii_grpc::types::KeysClause;
+use torii_grpc::proto::world::RetrieveEntitiesResponse;
+use torii_grpc::types::{Entity, KeysClause, Query};
 
 use self::error::{Error, ParseError};
 use self::storage::ModelStorage;
@@ -98,6 +99,18 @@ impl Client {
         self.subscribed_entities.entities_keys.read()
     }
 
+    /// Retrieves entities matching specified keys and/or model name in query parameter.
+    ///
+    /// The query can include keys and a model name, both optional. Without parameters, it fetches
+    /// all entities, which is less efficient as it requires additional queries for each
+    /// entity's model data. Specifying a model name optimizes the process by limiting the
+    /// retrieval to entities with that model, requiring just one query.
+    pub async fn entities(&self, query: Query) -> Result<Vec<Entity>, Error> {
+        let mut grpc_client = self.inner.write().await;
+        let RetrieveEntitiesResponse { entities } = grpc_client.retrieve_entities(query).await?;
+        Ok(entities.into_iter().map(TryInto::try_into).collect::<Result<Vec<Entity>, _>>()?)
+    }
+
     /// Returns the model value of an entity.
     ///
     /// This function will only return `None`, if `model` doesn't exist. If there is no entity with

crates/torii/core/src/model.rs

@@ -177,14 +177,9 @@ pub fn build_sql_query(model_schemas: &Vec<Ty>) -> Result<String, Error> {
         }
     }
 
-    let primary_table = model_schemas[0].name();
     let mut global_selections = Vec::new();
-    let mut global_tables = model_schemas
-        .iter()
-        .enumerate()
-        .filter(|(index, _)| *index != 0) // primary_table don't `JOIN` itself
-        .map(|(_, schema)| schema.name())
-        .collect::<Vec<String>>();
+    let mut global_tables =
+        model_schemas.iter().enumerate().map(|(_, schema)| schema.name()).collect::<Vec<String>>();
 
     for ty in model_schemas {
         let schema = ty.as_struct().expect("schema should be struct");
@@ -206,11 +201,11 @@ pub fn build_sql_query(model_schemas: &Vec<Ty>) -> Result<String, Error> {
     let selections_clause = global_selections.join(", ");
     let join_clause = global_tables
         .into_iter()
-        .map(|table| format!(" LEFT JOIN {table} ON {primary_table}.entity_id = {table}.entity_id"))
+        .map(|table| format!(" LEFT JOIN {table} ON entities.id = {table}.entity_id"))
         .collect::<Vec<_>>()
         .join(" ");
 
-    Ok(format!("SELECT {selections_clause} FROM {primary_table}{join_clause}"))
+    Ok(format!("SELECT entities.keys, {selections_clause} FROM entities{join_clause}"))
 }
 
 /// Populate the values of a Ty (schema) from SQLite row.
@@ -509,12 +504,14 @@ mod tests {
         });
 
         let query = build_sql_query(&vec![ty]).unwrap();
+        println!("{query}");
         assert_eq!(
             query,
-            "SELECT Position.external_name AS \"Position.name\", Position.external_age AS \
-             \"Position.age\", Position$Vec2.external_x AS \"Position$Vec2.x\", \
-             Position$Vec2.external_y AS \"Position$Vec2.y\" FROM Position LEFT JOIN \
-             Position$Vec2 ON Position.entity_id = Position$Vec2.entity_id"
+            "SELECT entities.keys, Position.external_name AS \"Position.name\", \
+             Position.external_age AS \"Position.age\", Position$Vec2.external_x AS \
+             \"Position$Vec2.x\", Position$Vec2.external_y AS \"Position$Vec2.y\" FROM entities \
+             LEFT JOIN Position ON entities.id = Position.entity_id  LEFT JOIN Position$Vec2 ON \
+             entities.id = Position$Vec2.entity_id"
         );
     }
 }

crates/torii/grpc/Cargo.toml

@@ -19,6 +19,7 @@ torii-core = { path = "../core", optional = true }
 
 serde.workspace = true
 strum_macros.workspace = true
+crypto-bigint.workspace = true
 
 # server
 hex.workspace = true

crates/torii/grpc/proto/schema.proto

@@ -0,0 +1,62 @@
+syntax = "proto3";
+package types;
+
+enum PrimitiveType {
+    U8 = 0;
+    U16 = 1;
+    U32 = 2;
+    U64 = 3;
+    U128 = 4;
+    U256 = 5;
+    USIZE = 6;
+    BOOL = 7;
+    FELT252 = 8;
+    CLASS_HASH = 9;
+    CONTRACT_ADDRESS = 10;
+}
+
+message EnumOption {
+    string name = 1;
+    Ty ty = 2;
+}
+
+message Enum {
+    string name = 1;
+    uint32 option = 2;
+    repeated EnumOption options = 3;
+}
+
+message Primitive {
+    PrimitiveType type = 1;
+    Value value = 2;
+}
+
+message Struct {
+    string name = 1;
+    repeated Member children = 2;
+}
+
+message Ty {
+    oneof ty_type {
+        Primitive primitive = 2;
+        Enum enum = 3;
+        Struct struct = 4;
+        // TODO: Tuple
+    }
+}
+
+message Member {
+    string name = 1;
+    Ty ty = 2;
+    bool key = 3;
+}
+
+message Value {
+    oneof value_type {
+        string string_value = 2;
+        int64 int_value = 3;
+        uint64 uint_value = 4;
+        bool bool_value = 5;
+        bytes byte_value = 6;
+    }
+}

crates/torii/grpc/proto/types.proto

@@ -1,6 +1,8 @@
 syntax = "proto3";
 package types;
 
+import "schema.proto";
+
 message WorldMetadata {
     // The hex-encoded address of the world.
     string world_address = 1;
@@ -29,31 +31,10 @@ message ModelMetadata {
     bytes schema = 6;
 }
 
-message Enum {
-    // Variant
-    uint32 option = 1;
-    // Variants of the enum
-    repeated string options = 2;
-}
-
-message Member {
-    // Name of the member
-    string name = 1;
-    // Type of member
-    oneof member_type {
-        Value value = 2;
-        Enum enum = 3;
-        Model struct = 4;
-    }
-}
-
 message Model {
-    // Name of the model
     string name = 1;
-    // Members of the model
     repeated Member members = 2;
 }
-
 message Entity {
     // The entity key
     bytes key = 1;
@@ -85,7 +66,7 @@ message EntityUpdate {
     EntityDiff entity_diff = 2;
 }
 
-message EntityQuery {
+message Query {
     Clause clause = 1;
     uint32 limit = 2;
     uint32 offset = 3;
@@ -129,14 +110,4 @@ enum ComparisonOperator {
     GTE = 3;
     LT = 4;
     LTE = 5;
-}
-
-message Value {
-    oneof value_type {
-        string string_value = 1;
-        int64 int_value = 2;
-        uint64 uint_value = 3;
-        bool bool_value = 4;
-        bytes byte_value = 5;
-    }
-}
+}

crates/torii/grpc/proto/world.proto

@@ -39,7 +39,7 @@ message SubscribeEntitiesResponse {
 
 message RetrieveEntitiesRequest {
     // The entities to retrieve
-    types.EntityQuery query = 1;
+    types.Query query = 1;
 }
 
 message RetrieveEntitiesResponse {

crates/torii/grpc/src/client.rs

@@ -1,22 +1,32 @@
 //! Client implementation for the gRPC service.
+use std::num::ParseIntError;
+
 use futures_util::stream::MapOk;
 use futures_util::{Stream, StreamExt, TryStreamExt};
 use proto::world::{world_client, SubscribeEntitiesRequest};
-use starknet::core::types::{FromStrError, StateUpdate};
+use starknet::core::types::{FromByteSliceError, FromStrError, StateUpdate};
 use starknet_crypto::FieldElement;
 
-use crate::proto::world::{MetadataRequest, SubscribeEntitiesResponse};
+use crate::proto::world::{
+    MetadataRequest, RetrieveEntitiesRequest, RetrieveEntitiesResponse, SubscribeEntitiesResponse,
+};
 use crate::proto::{self};
-use crate::types::KeysClause;
+use crate::types::{KeysClause, Query};
 
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
     #[error(transparent)]
     Grpc(tonic::Status),
     #[error("Missing expected data")]
     MissingExpectedData,
+    #[error("Unsupported type")]
+    UnsupportedType,
+    #[error(transparent)]
+    ParseStr(FromStrError),
+    #[error(transparent)]
+    SliceError(FromByteSliceError),
     #[error(transparent)]
-    Parsing(FromStrError),
+    ParseInt(ParseIntError),
 
     #[cfg(not(target_arch = "wasm32"))]
     #[error(transparent)]
@@ -61,7 +71,15 @@ impl WorldClient {
             .await
             .map_err(Error::Grpc)
             .and_then(|res| res.into_inner().metadata.ok_or(Error::MissingExpectedData))
-            .and_then(|metadata| metadata.try_into().map_err(Error::Parsing))
+            .and_then(|metadata| metadata.try_into().map_err(Error::ParseStr))
+    }
+
+    pub async fn retrieve_entities(
+        &mut self,
+        query: Query,
+    ) -> Result<RetrieveEntitiesResponse, Error> {
+        let request = RetrieveEntitiesRequest { query: Some(query.into()) };
+        self.inner.retrieve_entities(request).await.map_err(Error::Grpc).map(|res| res.into_inner())
     }
 
     /// Subscribe to the state diff for a set of entities of a World.