refactor: use array of models for keysclause

crates/torii/grpc/proto/types.proto

@@ -114,7 +114,7 @@ message EntityKeysClause {
 message KeysClause {
     repeated bytes keys = 1;
     PatternMatching pattern_matching = 2;
-    optional string model = 3;
+    repeated string models = 3;
 }
 
 message HashedKeysClause {

crates/torii/grpc/src/server/mod.rs

@@ -286,7 +286,7 @@ impl DojoWorld {
                 arrays_rows.insert(name, rows);
             }
 
-            entities.push(Self::map_row_to_entity(&row, &arrays_rows, &schemas)?);
+            entities.push(map_row_to_entity(&row, &arrays_rows, &schemas)?);
         }
 
         Ok((entities, total_count))
@@ -327,19 +327,27 @@ impl DojoWorld {
             FROM {table}
             {}
         "#,
-            if let Some(model) = &keys_clause.model {
+            if !keys_clause.models.is_empty() {
+                let model_ids = keys_clause
+                    .models
+                    .iter()
+                    .map(|model| get_selector_from_name(model).map_err(ParseError::NonAsciiName))
+                    .collect::<Result<Vec<_>, _>>()?;
+                let model_ids_str =
+                    model_ids.iter().map(|id| format!("'{:#x}'", id)).collect::<Vec<_>>().join(",");
                 format!(
                     r#"
-                    JOIN {model_relation_table} ON {table}.id = {model_relation_table}.entity_id
-                    WHERE {model_relation_table}.model_id = '{:#x}' AND {table}.keys REGEXP ?
-                "#,
-                    get_selector_from_name(model).map_err(ParseError::NonAsciiName)?
+                JOIN {model_relation_table} ON {table}.id = {model_relation_table}.entity_id
+                WHERE {model_relation_table}.model_id IN ({})
+                AND {table}.keys REGEXP ?
+            "#,
+                    model_ids_str
                 )
             } else {
                 format!(
                     r#"
-                    WHERE {table}.keys REGEXP ?
-                "#
+                WHERE {table}.keys REGEXP ?
+            "#
                 )
             }
         );
@@ -358,19 +366,29 @@ impl DojoWorld {
             JOIN {model_relation_table} ON {table}.id = {model_relation_table}.entity_id
             WHERE {table}.keys REGEXP ?
             GROUP BY {table}.id
-            ORDER BY {table}.event_id DESC
         "#
         );
 
-        if let Some(model) = &keys_clause.model {
+        if !keys_clause.models.is_empty() {
+            // filter by models
             models_query += &format!(
-                r#"
-                HAVING INSTR(model_ids, '{:#x}') > 0
-            "#,
-                get_selector_from_name(model).map_err(ParseError::NonAsciiName)?
+                "HAVING {}",
+                keys_clause
+                    .models
+                    .iter()
+                    .map(|model| {
+                        let model_id =
+                            get_selector_from_name(model).map_err(ParseError::NonAsciiName)?;
+                        Ok(format!("INSTR(model_ids, '{:#x}') > 0", model_id))
+                    })
+                    .collect::<Result<Vec<_>, Error>>()?
+                    .join(" OR ")
+                    .as_str()
             );
         }
 
+        models_query += &format!(" ORDER BY {table}.event_id DESC");
+
         if limit.is_some() {
             models_query += " LIMIT ?";
         }
@@ -405,7 +423,7 @@ impl DojoWorld {
                 arrays_rows.insert(name, rows);
             }
 
-            entities.push(Self::map_row_to_entity(&row, &arrays_rows, &schemas)?);
+            entities.push(map_row_to_entity(&row, &arrays_rows, &schemas)?);
         }
 
         Ok((entities, total_count))
@@ -531,7 +549,7 @@ impl DojoWorld {
 
         let entities_collection = db_entities
             .iter()
-            .map(|row| Self::map_row_to_entity(row, &arrays_rows, &schemas))
+            .map(|row| map_row_to_entity(row, &arrays_rows, &schemas))
             .collect::<Result<Vec<_>, Error>>()?;
         // Since there is not limit and offset, total_count is same as number of entities
         let total_count = entities_collection.len() as u32;
@@ -780,30 +798,6 @@ impl DojoWorld {
     ) -> Result<Receiver<Result<proto::world::SubscribeEventsResponse, tonic::Status>>, Error> {
         self.event_manager.add_subscriber(clause.try_into().unwrap()).await
     }
-
-    fn map_row_to_entity(
-        row: &SqliteRow,
-        arrays_rows: &HashMap<String, Vec<SqliteRow>>,
-        schemas: &[Ty],
-    ) -> Result<proto::types::Entity, Error> {
-        let hashed_keys =
-            FieldElement::from_str(&row.get::<String, _>("id")).map_err(ParseError::FromStr)?;
-        let models = schemas
-            .iter()
-            .map(|schema| {
-                let mut schema = schema.to_owned();
-                map_row_to_ty("", &schema.name(), &mut schema, row, arrays_rows)?;
-                Ok(schema
-                    .as_struct()
-                    .expect("schema should be struct")
-                    .to_owned()
-                    .try_into()
-                    .unwrap())
-            })
-            .collect::<Result<Vec<_>, Error>>()?;
-
-        Ok(proto::types::Entity { hashed_keys: hashed_keys.to_bytes_be().to_vec(), models })
-    }
 }
 
 fn process_event_field(data: &str) -> Result<Vec<Vec<u8>>, Error> {
@@ -825,6 +819,25 @@ fn map_row_to_event(row: &(String, String, String)) -> Result<proto::types::Even
     Ok(proto::types::Event { keys, data, transaction_hash })
 }
 
+fn map_row_to_entity(
+    row: &SqliteRow,
+    arrays_rows: &HashMap<String, Vec<SqliteRow>>,
+    schemas: &[Ty],
+) -> Result<proto::types::Entity, Error> {
+    let hashed_keys =
+        FieldElement::from_str(&row.get::<String, _>("id")).map_err(ParseError::FromStr)?;
+    let models = schemas
+        .iter()
+        .map(|schema| {
+            let mut schema = schema.to_owned();
+            map_row_to_ty("", &schema.name(), &mut schema, row, arrays_rows)?;
+            Ok(schema.as_struct().expect("schema should be struct").to_owned().try_into().unwrap())
+        })
+        .collect::<Result<Vec<_>, Error>>()?;
+
+    Ok(proto::types::Entity { hashed_keys: hashed_keys.to_bytes_be().to_vec(), models })
+}
+
 type ServiceResult<T> = Result<Response<T>, Status>;
 type SubscribeModelsResponseStream =
     Pin<Box<dyn Stream<Item = Result<SubscribeModelsResponse, Status>> + Send>>;

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

@@ -14,14 +14,15 @@ use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::sync::RwLock;
 use torii_core::cache::ModelCache;
 use torii_core::error::{Error, ParseError};
-use torii_core::model::{build_sql_query, map_row_to_ty};
+use torii_core::model::build_sql_query;
 use torii_core::simple_broker::SimpleBroker;
 use torii_core::sql::FELT_DELIMITER;
 use torii_core::types::Entity;
 use tracing::{error, trace};
 
 use crate::proto;
 use crate::proto::world::SubscribeEntityResponse;
+use crate::server::map_row_to_entity;
 use crate::types::{EntityKeysClause, PatternMatching};
 
 pub(crate) const LOG_TARGET: &str = "torii::grpc::server::subscriptions::entity";
@@ -115,11 +116,11 @@ impl Service {
                 Some(EntityKeysClause::Keys(clause)) => {
                     // if we have a model clause, then we need to check that the entity
                     // has an updated model and that the model name matches the clause
-                    if let Some(model) = &clause.model {
-                        if let Some(updated_model) = &entity.updated_model {
-                            if updated_model.name() != model.clone() {
-                                continue;
-                            }
+                    if let Some(updated_model) = &entity.updated_model {
+                        if !clause.models.is_empty()
+                            && !clause.models.contains(&updated_model.name())
+                        {
+                            continue;
                         }
                     }
 
@@ -199,24 +200,8 @@ impl Service {
                 arrays_rows.insert(name, row);
             }
 
-            let models = schemas
-                .into_iter()
-                .map(|mut s| {
-                    map_row_to_ty("", &s.name(), &mut s, &row, &arrays_rows)?;
-
-                    Ok(s.as_struct()
-                        .expect("schema should be a struct")
-                        .to_owned()
-                        .try_into()
-                        .unwrap())
-                })
-                .collect::<Result<Vec<_>, Error>>()?;
-
             let resp = proto::world::SubscribeEntityResponse {
-                entity: Some(proto::types::Entity {
-                    hashed_keys: hashed.to_bytes_be().to_vec(),
-                    models,
-                }),
+                entity: Some(map_row_to_entity(&row, &arrays_rows, &schemas)?),
             };
 
             if sub.sender.send(Ok(resp)).await.is_err() {

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

@@ -14,14 +14,15 @@ use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::sync::RwLock;
 use torii_core::cache::ModelCache;
 use torii_core::error::{Error, ParseError};
-use torii_core::model::{build_sql_query, map_row_to_ty};
+use torii_core::model::build_sql_query;
 use torii_core::simple_broker::SimpleBroker;
 use torii_core::sql::FELT_DELIMITER;
 use torii_core::types::EventMessage;
 use tracing::{error, trace};
 
 use crate::proto;
 use crate::proto::world::SubscribeEntityResponse;
+use crate::server::map_row_to_entity;
 use crate::types::{EntityKeysClause, PatternMatching};
 
 pub(crate) const LOG_TARGET: &str = "torii::grpc::server::subscriptions::event_message";
@@ -114,11 +115,11 @@ impl Service {
                 Some(EntityKeysClause::Keys(clause)) => {
                     // if we have a model clause, then we need to check that the entity
                     // has an updated model and that the model name matches the clause
-                    if let Some(model) = &clause.model {
-                        if let Some(updated_model) = &entity.updated_model {
-                            if updated_model.name() != model.clone() {
-                                continue;
-                            }
+                    if let Some(updated_model) = &entity.updated_model {
+                        if !clause.models.is_empty()
+                            && !clause.models.contains(&updated_model.name())
+                        {
+                            continue;
                         }
                     }
 
@@ -184,23 +185,8 @@ impl Service {
                 arrays_rows.insert(name, rows);
             }
 
-            let models = schemas
-                .into_iter()
-                .map(|mut s| {
-                    map_row_to_ty("", &s.name(), &mut s, &row, &arrays_rows)?;
-                    Ok(s.as_struct()
-                        .expect("schema should be a struct")
-                        .to_owned()
-                        .try_into()
-                        .unwrap())
-                })
-                .collect::<Result<Vec<_>, Error>>()?;
-
             let resp = proto::world::SubscribeEntityResponse {
-                entity: Some(proto::types::Entity {
-                    hashed_keys: hashed.to_bytes_be().to_vec(),
-                    models,
-                }),
+                entity: Some(map_row_to_entity(&row, &arrays_rows, &schemas)?),
             };
 
             if sub.sender.send(Ok(resp)).await.is_err() {

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

@@ -121,7 +121,7 @@ async fn test_entities_queries() {
             KeysClause {
                 keys: vec![account.address().to_bytes_be().to_vec()],
                 pattern_matching: 0,
-                model: None,
+                models: vec![],
             },
             Some(1),
             None,

crates/torii/grpc/src/types/mod.rs

@@ -45,7 +45,7 @@ pub struct ModelKeysClause {
 pub struct KeysClause {
     pub keys: Vec<FieldElement>,
     pub pattern_matching: PatternMatching,
-    pub model: Option<String>,
+    pub models: Vec<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Hash, Eq, Clone)]
@@ -186,7 +186,7 @@ impl From<KeysClause> for proto::types::KeysClause {
         Self {
             keys: value.keys.iter().map(|k| k.to_bytes_be().into()).collect(),
             pattern_matching: value.pattern_matching as i32,
-            model: value.model,
+            models: value.models,
         }
     }
 }
@@ -201,7 +201,7 @@ impl TryFrom<proto::types::KeysClause> for KeysClause {
             .map(|k| FieldElement::from_byte_slice_be(k))
             .collect::<Result<Vec<_>, _>>()?;
 
-        Ok(Self { keys, pattern_matching: value.pattern_matching().into(), model: value.model })
+        Ok(Self { keys, pattern_matching: value.pattern_matching().into(), models: value.models })
     }
 }