diff --git a/crates/torii/client/src/client/mod.rs b/crates/torii/client/src/client/mod.rs index 32984c6292..07c96bf77c 100644 --- a/crates/torii/client/src/client/mod.rs +++ b/crates/torii/client/src/client/mod.rs @@ -158,7 +158,8 @@ impl Client { /// type of entites matching keys and/or models. pub async fn entities(&self, query: Query) -> Result, 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::, _>>()?) } diff --git a/crates/torii/grpc/proto/world.proto b/crates/torii/grpc/proto/world.proto index 0539e4e493..a28ebaf0c7 100644 --- a/crates/torii/grpc/proto/world.proto +++ b/crates/torii/grpc/proto/world.proto @@ -54,4 +54,5 @@ message RetrieveEntitiesRequest { message RetrieveEntitiesResponse { repeated types.Entity entities = 1; + uint32 total_count = 2; } diff --git a/crates/torii/grpc/src/server/mod.rs b/crates/torii/grpc/src/server/mod.rs index a6c3625a28..0cecf11365 100644 --- a/crates/torii/grpc/src/server/mod.rs +++ b/crates/torii/grpc/src/server/mod.rs @@ -120,7 +120,7 @@ impl DojoWorld { &self, limit: u32, offset: u32, - ) -> Result, Error> { + ) -> Result<(Vec, u32), Error> { self.entities_by_hashed_keys(None, limit, offset).await } @@ -129,7 +129,7 @@ impl DojoWorld { hashed_keys: Option, limit: u32, offset: u32, - ) -> Result, Error> { + ) -> Result<(Vec, u32), Error> { // TODO: use prepared statement for where clause let filter_ids = match hashed_keys { Some(hashed_keys) => { @@ -148,6 +148,18 @@ impl DojoWorld { 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,7 +200,7 @@ impl DojoWorld { }) } - Ok(entities) + Ok((entities, total_count)) } async fn entities_by_keys( @@ -196,7 +208,7 @@ impl DojoWorld { keys_clause: proto::types::KeysClause, limit: u32, offset: u32, - ) -> Result, Error> { + ) -> Result<(Vec, u32), Error> { let keys = keys_clause .keys .iter() @@ -211,6 +223,20 @@ impl DojoWorld { .collect::, 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 @@ impl DojoWorld { let model_names = models_str.split(',').collect::>(); 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,7 +267,13 @@ impl DojoWorld { .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::, Error>>()?, + total_count, + )) } async fn entities_by_member( @@ -248,7 +281,7 @@ impl DojoWorld { member_clause: proto::types::MemberClause, _limit: u32, _offset: u32, - ) -> Result, Error> { + ) -> Result<(Vec, u32), Error> { let comparison_operator = ComparisonOperator::from_repr(member_clause.operator as usize) .expect("invalid comparison operator"); @@ -297,8 +330,13 @@ impl DojoWorld { 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::, 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( @@ -306,7 +344,7 @@ impl DojoWorld { _composite: proto::types::CompositeClause, _limit: u32, _offset: u32, - ) -> Result, Error> { + ) -> Result<(Vec, u32), Error> { // TODO: Implement Err(QueryError::UnsupportedQuery.into()) } @@ -373,7 +411,7 @@ impl DojoWorld { &self, query: proto::types::Query, ) -> Result { - 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 @@ impl DojoWorld { } }; - Ok(RetrieveEntitiesResponse { entities }) + Ok(RetrieveEntitiesResponse { entities, total_count }) } fn map_row_to_entity(row: &SqliteRow, schemas: &[Ty]) -> Result { diff --git a/crates/torii/types-test/Scarb.lock b/crates/torii/types-test/Scarb.lock index 0081e547e0..69570413ff 100644 --- a/crates/torii/types-test/Scarb.lock +++ b/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", ] diff --git a/examples/spawn-and-move/Scarb.lock b/examples/spawn-and-move/Scarb.lock index 0a3638b88d..af4e4b139f 100644 --- a/examples/spawn-and-move/Scarb.lock +++ b/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", ]