Feat(torii-grpc): Add total_count on RetrieveEntitiesResponse

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

@@ -158,7 +158,8 @@
     /// type of entites matching keys and/or models.
     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?;
+        let RetrieveEntitiesResponse { entities, total_count: _ } =
+            grpc_client.retrieve_entities(query).await?;
         Ok(entities.into_iter().map(TryInto::try_into).collect::<Result<Vec<Entity>, _>>()?)
     }
 

crates/torii/grpc/proto/world.proto

@@ -54,4 +54,5 @@ message RetrieveEntitiesRequest {
 
 message RetrieveEntitiesResponse {
     repeated types.Entity entities = 1;
+    uint32 total_count = 2;
 }

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

@@ -120,7 +120,7 @@
         &self,
         limit: u32,
         offset: u32,
-    ) -> Result<Vec<proto::types::Entity>, Error> {
+    ) -> Result<(Vec<proto::types::Entity>, u32), Error> {
         self.entities_by_hashed_keys(None, limit, offset).await
     }
 
@@ -129,7 +129,7 @@
         hashed_keys: Option<proto::types::HashedKeysClause>,
         limit: u32,
         offset: u32,
-    ) -> Result<Vec<proto::types::Entity>, Error> {
+    ) -> Result<(Vec<proto::types::Entity>, u32), Error> {
         // TODO: use prepared statement for where clause
         let filter_ids = match hashed_keys {
             Some(hashed_keys) => {
@@ -148,6 +148,18 @@
             None => String::new(),
         };
 
+        // count query that matches filter_ids
+        let count_query = format!(
+            r#"
+                    SELECT count(*)
+                    FROM entities
+                    {filter_ids}
+                "#
+        );
+        // total count of rows without limit and offset
+        let total_count: u32 = sqlx::query_scalar(&count_query).fetch_one(&self.pool).await?;
+
+        // query to filter with limit and offset
         let query = format!(
             r#"
             SELECT entities.id, group_concat(entity_model.model_id) as model_names
@@ -188,15 +200,15 @@
             })
         }
 
-        Ok(entities)
+        Ok((entities, total_count))
     }
 
     async fn entities_by_keys(
         &self,
         keys_clause: proto::types::KeysClause,
         limit: u32,
         offset: u32,
-    ) -> Result<Vec<proto::types::Entity>, Error> {
+    ) -> Result<(Vec<proto::types::Entity>, u32), Error> {
         let keys = keys_clause
             .keys
             .iter()
@@ -211,6 +223,20 @@
             .collect::<Result<Vec<_>, Error>>()?;
         let keys_pattern = keys.join("/") + "/%";
 
+        let count_query = format!(
+            r#"
+            SELECT count(*)
+            FROM entities
+            JOIN entity_model ON entities.id = entity_model.entity_id
+            WHERE entity_model.model_id = '{}' and entities.keys LIKE ?
+        "#,
+            keys_clause.model
+        );
+
+        // total count of rows that matches keys_pattern without limit and offset
+        let total_count =
+            sqlx::query_scalar(&count_query).bind(&keys_pattern).fetch_one(&self.pool).await?;
+
         let models_query = format!(
             r#"
             SELECT group_concat(entity_model.model_id) as model_names
@@ -229,6 +255,7 @@
         let model_names = models_str.split(',').collect::<Vec<&str>>();
         let schemas = self.model_cache.schemas(model_names).await?;
 
+        // query to filter with limit and offset
         let entities_query = format!(
             "{} WHERE entities.keys LIKE ? ORDER BY entities.event_id DESC LIMIT ? OFFSET ?",
             build_sql_query(&schemas)?
@@ -240,15 +267,21 @@
             .fetch_all(&self.pool)
             .await?;
 
-        db_entities.iter().map(|row| Self::map_row_to_entity(row, &schemas)).collect()
+        Ok((
+            db_entities
+                .iter()
+                .map(|row| Self::map_row_to_entity(row, &schemas))
+                .collect::<Result<Vec<_>, Error>>()?,
+            total_count,
+        ))
     }
 
     async fn entities_by_member(
         &self,
         member_clause: proto::types::MemberClause,
         _limit: u32,
         _offset: u32,
-    ) -> Result<Vec<proto::types::Entity>, Error> {
+    ) -> Result<(Vec<proto::types::Entity>, u32), Error> {
         let comparison_operator = ComparisonOperator::from_repr(member_clause.operator as usize)
             .expect("invalid comparison operator");
 
@@ -297,16 +330,21 @@
 
         let db_entities =
             sqlx::query(&member_query).bind(comparison_value).fetch_all(&self.pool).await?;
-
-        db_entities.iter().map(|row| Self::map_row_to_entity(row, &schemas)).collect()
+        let entities_collection = db_entities
+            .iter()
+            .map(|row| Self::map_row_to_entity(row, &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;
+        Ok((entities_collection, total_count))
     }
 
     async fn entities_by_composite(
         &self,
         _composite: proto::types::CompositeClause,
         _limit: u32,
         _offset: u32,
-    ) -> Result<Vec<proto::types::Entity>, Error> {
+    ) -> Result<(Vec<proto::types::Entity>, u32), Error> {
         // TODO: Implement
         Err(QueryError::UnsupportedQuery.into())
     }
@@ -373,7 +411,7 @@
         &self,
         query: proto::types::Query,
     ) -> Result<proto::world::RetrieveEntitiesResponse, Error> {
-        let entities = match query.clause {
+        let (entities, total_count) = match query.clause {
             None => self.entities_all(query.limit, query.offset).await?,
             Some(clause) => {
                 let clause_type =
@@ -409,7 +447,7 @@
             }
         };
 
-        Ok(RetrieveEntitiesResponse { entities })
+        Ok(RetrieveEntitiesResponse { entities, total_count })
     }
 
     fn map_row_to_entity(row: &SqliteRow, schemas: &[Ty]) -> Result<proto::types::Entity, Error> {

crates/torii/types-test/Scarb.lock

@@ -15,7 +15,7 @@ source = "git+https://github.com/dojoengine/dojo?tag=v0.3.11#1e651b5d4d3b79b14a7
 
 [[package]]
 name = "types_test"
-version = "0.5.0"
+version = "0.5.1"
 dependencies = [
  "dojo",
 ]

examples/spawn-and-move/Scarb.lock

@@ -10,7 +10,7 @@ dependencies = [
 
 [[package]]
 name = "dojo_examples"
-version = "0.5.0"
+version = "0.5.1"
 dependencies = [
  "dojo",
 ]