diff --git a/src/satellite/src/db/store.rs b/src/satellite/src/db/store.rs index a51edd8ca..641c04203 100644 --- a/src/satellite/src/db/store.rs +++ b/src/satellite/src/db/store.rs @@ -326,3 +326,19 @@ fn delete_docs_impl( Ok(()) } + +pub fn count_docs(collection: &CollectionKey) -> Result { + let rule = get_state_rule(&collection)?; + + match rule.mem() { + Memory::Heap => STATE.with(|state| { + let state_ref = state.borrow(); + let docs = get_docs_heap(collection, &state_ref.heap.db.db)?; + Ok(docs.len()) + }), + Memory::Stable => STATE.with(|state| { + let stable = get_docs_stable(collection, &state.borrow().stable.db)?; + Ok(stable.len()) + }), + } +} diff --git a/src/satellite/src/lib.rs b/src/satellite/src/lib.rs index f9602820b..59bfe4aac 100644 --- a/src/satellite/src/lib.rs +++ b/src/satellite/src/lib.rs @@ -11,7 +11,10 @@ mod storage; mod types; use crate::controllers::store::get_admin_controllers; -use crate::db::store::{delete_doc, delete_docs, get_doc as get_doc_store, get_docs, insert_doc}; +use crate::db::store::{ + count_docs as count_docs_store, delete_doc, delete_docs, get_doc as get_doc_store, get_docs, + insert_doc, +}; use crate::db::types::interface::{DelDoc, SetDoc}; use crate::db::types::state::Doc; use crate::guards::{caller_is_admin_controller, caller_is_controller}; @@ -28,10 +31,10 @@ use crate::storage::http::response::{ use crate::storage::http::utils::create_token; use crate::storage::routing::get_routing; use crate::storage::store::{ - commit_batch, create_batch, create_chunk, delete_asset, delete_assets, delete_domain, - get_config as get_storage_config, get_content_chunks, get_custom_domains, get_public_asset, - init_certified_assets, list_assets as list_assets_store, set_config as set_storage_config, - set_domain, + commit_batch, count_assets as count_assets_store, create_batch, create_chunk, delete_asset, + delete_assets, delete_domain, get_config as get_storage_config, get_content_chunks, + get_custom_domains, get_public_asset, init_certified_assets, list_assets as list_assets_store, + set_config as set_storage_config, set_domain, }; use crate::storage::types::domain::{CustomDomains, DomainName}; use crate::storage::types::http_request::{ @@ -217,6 +220,16 @@ fn del_docs(collection: CollectionKey) { } } +#[query(guard = "caller_is_admin_controller")] +fn count_docs(collection: CollectionKey) -> usize { + let result = count_docs_store(&collection); + + match result { + Ok(value) => value, + Err(error) => trap(&error), + } +} + /// Rules #[query(guard = "caller_is_admin_controller")] @@ -495,6 +508,16 @@ fn del_assets(collection: CollectionKey) { } } +#[query(guard = "caller_is_admin_controller")] +fn count_assets(collection: CollectionKey) -> usize { + let result = count_assets_store(&collection); + + match result { + Ok(value) => value, + Err(error) => trap(&error), + } +} + /// Mgmt #[update(guard = "caller_is_admin_controller")] diff --git a/src/satellite/src/storage/store.rs b/src/satellite/src/storage/store.rs index 12b547f68..3ceaaab4b 100644 --- a/src/satellite/src/storage/store.rs +++ b/src/satellite/src/storage/store.rs @@ -305,6 +305,22 @@ fn delete_assets_impl( Ok(()) } +pub fn count_assets(collection: &CollectionKey) -> Result { + let rule = get_state_rule(collection)?; + + match rule.mem() { + Memory::Heap => STATE.with(|state| { + let state_ref = state.borrow(); + let assets = get_assets_heap(collection, &state_ref.heap.storage.assets); + Ok(assets.len()) + }), + Memory::Stable => STATE.with(|state| { + let stable = get_assets_stable(collection, &state.borrow().stable.assets); + Ok(stable.len()) + }), + } +} + /// /// Upload batch and chunks ///