feat: update torii client and fix tets

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

@@ -20,7 +20,7 @@ use tokio::sync::RwLock as AsyncRwLock;
 use torii_grpc::client::{EntityUpdateStreaming, ModelDiffsStreaming};
 use torii_grpc::proto::world::RetrieveEntitiesResponse;
 use torii_grpc::types::schema::Entity;
-use torii_grpc::types::{KeysClause, Query};
+use torii_grpc::types::{EntityKeysClause, ModelKeysClause, Query};
 use torii_relay::client::EventLoop;
 use torii_relay::types::Message;
 
@@ -56,7 +56,6 @@ impl Client {
         rpc_url: String,
         relay_url: String,
         world: FieldElement,
-        models_keys: Option<Vec<KeysClause>>,
     ) -> Result<Self, Error> {
         let mut grpc_client = torii_grpc::client::WorldClient::new(torii_url, world).await?;
 
@@ -73,23 +72,6 @@ impl Client {
         let provider = JsonRpcClient::new(HttpTransport::new(rpc_url));
         let world_reader = WorldContractReader::new(world, provider);
 
-        if let Some(keys) = models_keys {
-            subbed_models.add_models(keys)?;
-
-            // TODO: change this to querying the gRPC url instead
-            let subbed_models = subbed_models.models_keys.read().clone();
-            for keys in subbed_models {
-                let model_reader = world_reader.model_reader(&keys.model).await?;
-                let values = model_reader.entity_storage(&keys.keys).await?;
-
-                client_storage.set_model_storage(
-                    cairo_short_string_to_felt(&keys.model).unwrap(),
-                    keys.keys,
-                    values,
-                )?;
-            }
-        }
-
         Ok(Self {
             world_reader,
             storage: client_storage,
@@ -123,7 +105,7 @@ impl Client {
         self.metadata.read()
     }
 
-    pub fn subscribed_models(&self) -> RwLockReadGuard<'_, HashSet<KeysClause>> {
+    pub fn subscribed_models(&self) -> RwLockReadGuard<'_, HashSet<ModelKeysClause>> {
         self.subscribed_models.models_keys.read()
     }
 
@@ -151,20 +133,20 @@ impl Client {
     /// A direct stream to grpc subscribe entities
     pub async fn on_entity_updated(
         &self,
-        ids: Vec<FieldElement>,
+        clause: Option<EntityKeysClause>,
     ) -> Result<EntityUpdateStreaming, Error> {
         let mut grpc_client = self.inner.write().await;
-        let stream = grpc_client.subscribe_entities(ids).await?;
+        let stream = grpc_client.subscribe_entities(clause).await?;
         Ok(stream)
     }
 
     /// A direct stream to grpc subscribe event messages
     pub async fn on_event_message_updated(
         &self,
-        ids: Vec<FieldElement>,
+        clause: Option<EntityKeysClause>,
     ) -> Result<EntityUpdateStreaming, Error> {
         let mut grpc_client = self.inner.write().await;
-        let stream = grpc_client.subscribe_event_messages(ids).await?;
+        let stream = grpc_client.subscribe_event_messages(clause).await?;
         Ok(stream)
     }
 
@@ -175,7 +157,7 @@ impl Client {
     ///
     /// If the requested model is not among the synced models, it will attempt to fetch it from
     /// the RPC.
-    pub async fn model(&self, keys: &KeysClause) -> Result<Option<Ty>, Error> {
+    pub async fn model(&self, keys: &ModelKeysClause) -> Result<Option<Ty>, Error> {
         let Some(mut schema) = self.metadata.read().model(&keys.model).map(|m| m.schema.clone())
         else {
             return Ok(None);
@@ -211,7 +193,7 @@ impl Client {
     /// Initiate the model subscriptions and returns a [SubscriptionService] which when await'ed
     /// will execute the subscription service and starts the syncing process.
     pub async fn start_subscription(&self) -> Result<SubscriptionService, Error> {
-        let models_keys: Vec<KeysClause> =
+        let models_keys: Vec<ModelKeysClause> =
             self.subscribed_models.models_keys.read().clone().into_iter().collect();
         let sub_res_stream = self.initiate_subscription(models_keys).await?;
 
@@ -229,7 +211,7 @@ impl Client {
     /// Adds entities to the list of entities to be synced.
     ///
     /// NOTE: This will establish a new subscription stream with the server.
-    pub async fn add_models_to_sync(&self, models_keys: Vec<KeysClause>) -> Result<(), Error> {
+    pub async fn add_models_to_sync(&self, models_keys: Vec<ModelKeysClause>) -> Result<(), Error> {
         for keys in &models_keys {
             self.initiate_model(&keys.model, keys.keys.clone()).await?;
         }
@@ -250,7 +232,10 @@ impl Client {
     /// Removes models from the list of models to be synced.
     ///
     /// NOTE: This will establish a new subscription stream with the server.
-    pub async fn remove_models_to_sync(&self, models_keys: Vec<KeysClause>) -> Result<(), Error> {
+    pub async fn remove_models_to_sync(
+        &self,
+        models_keys: Vec<ModelKeysClause>,
+    ) -> Result<(), Error> {
         self.subscribed_models.remove_models(models_keys)?;
 
         let updated_entities =
@@ -270,7 +255,7 @@ impl Client {
 
     async fn initiate_subscription(
         &self,
-        keys: Vec<KeysClause>,
+        keys: Vec<ModelKeysClause>,
     ) -> Result<ModelDiffsStreaming, Error> {
         let mut grpc_client = self.inner.write().await;
         let stream = grpc_client.subscribe_model_diffs(keys).await?;

crates/torii/client/src/client/subscription.rs

@@ -12,7 +12,7 @@ use starknet::core::types::{StateDiff, StateUpdate};
 use starknet::core::utils::cairo_short_string_to_felt;
 use starknet_crypto::FieldElement;
 use torii_grpc::client::ModelDiffsStreaming;
-use torii_grpc::types::KeysClause;
+use torii_grpc::types::ModelKeysClause;
 
 use crate::client::error::{Error, ParseError};
 use crate::client::storage::ModelStorage;
@@ -24,13 +24,13 @@ pub enum SubscriptionEvent {
 
 pub struct SubscribedModels {
     metadata: Arc<RwLock<WorldMetadata>>,
-    pub(crate) models_keys: RwLock<HashSet<KeysClause>>,
+    pub(crate) models_keys: RwLock<HashSet<ModelKeysClause>>,
     /// All the relevant storage addresses derived from the subscribed models
     pub(crate) subscribed_storage_addresses: RwLock<HashSet<FieldElement>>,
 }
 
 impl SubscribedModels {
-    pub(crate) fn is_synced(&self, keys: &KeysClause) -> bool {
+    pub(crate) fn is_synced(&self, keys: &ModelKeysClause) -> bool {
         self.models_keys.read().contains(keys)
     }
 
@@ -42,21 +42,21 @@ impl SubscribedModels {
         }
     }
 
-    pub(crate) fn add_models(&self, models_keys: Vec<KeysClause>) -> Result<(), Error> {
+    pub(crate) fn add_models(&self, models_keys: Vec<ModelKeysClause>) -> Result<(), Error> {
         for keys in models_keys {
             Self::add_model(self, keys)?;
         }
         Ok(())
     }
 
-    pub(crate) fn remove_models(&self, entities_keys: Vec<KeysClause>) -> Result<(), Error> {
+    pub(crate) fn remove_models(&self, entities_keys: Vec<ModelKeysClause>) -> Result<(), Error> {
         for keys in entities_keys {
             Self::remove_model(self, keys)?;
         }
         Ok(())
     }
 
-    pub(crate) fn add_model(&self, keys: KeysClause) -> Result<(), Error> {
+    pub(crate) fn add_model(&self, keys: ModelKeysClause) -> Result<(), Error> {
         if !self.models_keys.write().insert(keys.clone()) {
             return Ok(());
         }
@@ -83,7 +83,7 @@ impl SubscribedModels {
         Ok(())
     }
 
-    pub(crate) fn remove_model(&self, keys: KeysClause) -> Result<(), Error> {
+    pub(crate) fn remove_model(&self, keys: ModelKeysClause) -> Result<(), Error> {
         if !self.models_keys.write().remove(&keys) {
             return Ok(());
         }
@@ -247,7 +247,7 @@ mod tests {
     use parking_lot::RwLock;
     use starknet::core::utils::cairo_short_string_to_felt;
     use starknet::macros::felt;
-    use torii_grpc::types::KeysClause;
+    use torii_grpc::types::ModelKeysClause;
 
     use crate::utils::compute_all_storage_addresses;
 
@@ -283,7 +283,7 @@ mod tests {
 
         let metadata = self::create_dummy_metadata();
 
-        let keys = KeysClause { model: model_name, keys };
+        let keys = ModelKeysClause { model: model_name, keys };
 
         let subscribed_models = super::SubscribedModels::new(Arc::new(RwLock::new(metadata)));
         subscribed_models.add_models(vec![keys.clone()]).expect("able to add model");

crates/torii/grpc/src/server/subscriptions/entity.rs

@@ -1,4 +1,4 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::future::Future;
 use std::pin::Pin;
 use std::str::FromStr;

crates/torii/grpc/src/server/subscriptions/event_message.rs

@@ -1,4 +1,4 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::future::Future;
 use std::pin::Pin;
 use std::str::FromStr;
@@ -17,7 +17,7 @@ use torii_core::error::{Error, ParseError};
 use torii_core::model::{build_sql_query, map_row_to_ty};
 use torii_core::simple_broker::SimpleBroker;
 use torii_core::sql::FELT_DELIMITER;
-use torii_core::types::{Entity, EventMessage};
+use torii_core::types::EventMessage;
 use tracing::{error, trace};
 
 use crate::proto;

crates/torii/grpc/src/server/tests/entities_test.rs

@@ -26,9 +26,9 @@ use torii_core::processors::register_model::RegisterModelProcessor;
 use torii_core::processors::store_set_record::StoreSetRecordProcessor;
 use torii_core::sql::Sql;
 
+use crate::proto::types::KeysClause;
 use crate::server::DojoWorld;
 use crate::types::schema::Entity;
-use crate::types::KeysClause;
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_entities_queries() {
@@ -116,9 +116,12 @@ async fn test_entities_queries() {
             "entities",
             "entity_model",
             "entity_id",
-            KeysClause { model: "Moves".to_string(), keys: vec![account.address()] }.into(),
-            1,
-            0,
+            KeysClause {
+                keys: vec![account.address()].iter().map(|k| k.to_bytes_be().to_vec()).collect(),
+            }
+            .into(),
+            Some(1),
+            None,
         )
         .await
         .unwrap()