feat: collection type in user usage key

src/libs/satellite/src/lib.rs

@@ -19,7 +19,7 @@ mod version;
 use crate::auth::types::config::AuthenticationConfig;
 use crate::db::types::config::DbConfig;
 use crate::guards::{caller_is_admin_controller, caller_is_controller};
-use crate::types::interface::{Config, RulesType};
+use crate::types::interface::{CollectionType, Config};
 use crate::version::SATELLITE_VERSION;
 use ic_cdk::api::trap;
 use ic_cdk_macros::{init, post_upgrade, pre_upgrade, query, update};
@@ -175,25 +175,25 @@ pub fn count_collection_docs(collection: CollectionKey) -> usize {
 
 #[doc(hidden)]
 #[query(guard = "caller_is_admin_controller")]
-pub fn get_rule(rules_type: RulesType, collection: CollectionKey) -> Option<Rule> {
+pub fn get_rule(rules_type: CollectionType, collection: CollectionKey) -> Option<Rule> {
     satellite::get_rule(&rules_type, &collection)
 }
 
 #[doc(hidden)]
 #[query(guard = "caller_is_admin_controller")]
-pub fn list_rules(rules_type: RulesType) -> Vec<(CollectionKey, Rule)> {
+pub fn list_rules(rules_type: CollectionType) -> Vec<(CollectionKey, Rule)> {
     satellite::list_rules(rules_type)
 }
 
 #[doc(hidden)]
 #[update(guard = "caller_is_admin_controller")]
-pub fn set_rule(rules_type: RulesType, collection: CollectionKey, rule: SetRule) -> Rule {
+pub fn set_rule(rules_type: CollectionType, collection: CollectionKey, rule: SetRule) -> Rule {
     satellite::set_rule(rules_type, collection, rule)
 }
 
 #[doc(hidden)]
 #[update(guard = "caller_is_admin_controller")]
-pub fn del_rule(rules_type: RulesType, collection: CollectionKey, rule: DelRule) {
+pub fn del_rule(rules_type: CollectionType, collection: CollectionKey, rule: DelRule) {
     satellite::del_rule(rules_type, collection, rule)
 }
 
@@ -401,8 +401,12 @@ pub fn get_many_assets(
 
 #[doc(hidden)]
 #[query]
-pub fn get_user_usage(collection: CollectionKey, user_id: Option<UserId>) -> Option<UserUsage> {
-    satellite::get_user_usage(&collection, &user_id)
+pub fn get_user_usage(
+    collection_key: CollectionKey,
+    collection_type: CollectionType,
+    user_id: Option<UserId>,
+) -> Option<UserUsage> {
+    satellite::get_user_usage(&collection_key, &collection_type, &user_id)
 }
 
 // ---------------------------------------------------------

src/libs/satellite/src/satellite.rs

@@ -36,11 +36,12 @@ use crate::storage::store::{
     set_domain_store,
 };
 use crate::storage::strategy_impls::StorageState;
-use crate::types::interface::{Config, RulesType};
+use crate::types::interface::{CollectionType, Config};
 use crate::types::state::{HeapState, RuntimeState, State};
 use crate::usage::types::state::UserUsage;
 use crate::usage::user_usage::{
-    decrease_user_usage, decrease_user_usage_by, get_user_usage_by_id, increase_user_usage,
+    decrease_db_usage, decrease_db_usage_by, get_db_usage_by_id, get_storage_usage_by_id,
+    increase_db_usage,
 };
 use ciborium::{from_reader, into_writer};
 use ic_cdk::api::call::{arg_data, ArgDecoderConfig};
@@ -128,7 +129,7 @@ pub fn set_doc(collection: CollectionKey, key: Key, doc: SetDoc) -> Doc {
 
     match result {
         Ok(doc) => {
-            increase_user_usage(&caller, &collection);
+            increase_db_usage(&caller, &collection);
 
             invoke_on_set_doc(&caller, &doc);
 
@@ -155,7 +156,7 @@ pub fn del_doc(collection: CollectionKey, key: Key, doc: DelDoc) {
     let deleted_doc =
         delete_doc_store(caller, collection.clone(), key, doc).unwrap_or_else(|e| trap(&e));
 
-    decrease_user_usage(&caller, &collection);
+    decrease_db_usage(&caller, &collection);
 
     invoke_on_delete_doc(&caller, &deleted_doc);
 }
@@ -201,7 +202,7 @@ pub fn set_many_docs(docs: Vec<(CollectionKey, Key, SetDoc)>) -> Vec<(Key, Doc)>
         let result = set_doc_store(caller, collection.clone(), key.clone(), doc)
             .unwrap_or_else(|e| trap(&e));
 
-        increase_user_usage(&caller, &collection);
+        increase_db_usage(&caller, &collection);
 
         results.push((result.key.clone(), result.data.after.clone()));
 
@@ -222,7 +223,7 @@ pub fn del_many_docs(docs: Vec<(CollectionKey, Key, DelDoc)>) {
         let deleted_doc = delete_doc_store(caller, collection.clone(), key.clone(), doc)
             .unwrap_or_else(|e| trap(&e));
 
-        decrease_user_usage(&caller, &collection);
+        decrease_db_usage(&caller, &collection);
 
         results.push(deleted_doc);
     }
@@ -236,7 +237,7 @@ pub fn del_filtered_docs(collection: CollectionKey, filter: ListParams) {
     let results = delete_filtered_docs_store(caller, collection.clone(), &filter)
         .unwrap_or_else(|e| trap(&e));
 
-    decrease_user_usage_by(&caller, &collection, results.len() as u32);
+    decrease_db_usage_by(&caller, &collection, results.len() as u32);
 
     invoke_on_delete_filtered_docs(&caller, &results);
 }
@@ -263,31 +264,31 @@ pub fn count_collection_docs(collection: CollectionKey) -> usize {
 // Rules
 // ---------------------------------------------------------
 
-pub fn get_rule(rules_type: &RulesType, collection: &CollectionKey) -> Option<Rule> {
+pub fn get_rule(rules_type: &CollectionType, collection: &CollectionKey) -> Option<Rule> {
     match rules_type {
-        RulesType::Db => get_rule_db(collection),
-        RulesType::Storage => get_rule_storage(collection),
+        CollectionType::Db => get_rule_db(collection),
+        CollectionType::Storage => get_rule_storage(collection),
     }
 }
 
-pub fn list_rules(rules_type: RulesType) -> Vec<(CollectionKey, Rule)> {
+pub fn list_rules(rules_type: CollectionType) -> Vec<(CollectionKey, Rule)> {
     match rules_type {
-        RulesType::Db => get_rules_db(),
-        RulesType::Storage => get_rules_storage(),
+        CollectionType::Db => get_rules_db(),
+        CollectionType::Storage => get_rules_storage(),
     }
 }
 
-pub fn set_rule(rules_type: RulesType, collection: CollectionKey, rule: SetRule) -> Rule {
+pub fn set_rule(rules_type: CollectionType, collection: CollectionKey, rule: SetRule) -> Rule {
     match rules_type {
-        RulesType::Db => set_rule_db(collection, rule).unwrap_or_else(|e| trap(&e)),
-        RulesType::Storage => set_rule_storage(collection, rule).unwrap_or_else(|e| trap(&e)),
+        CollectionType::Db => set_rule_db(collection, rule).unwrap_or_else(|e| trap(&e)),
+        CollectionType::Storage => set_rule_storage(collection, rule).unwrap_or_else(|e| trap(&e)),
     }
 }
 
-pub fn del_rule(rules_type: RulesType, collection: CollectionKey, rule: DelRule) {
+pub fn del_rule(rules_type: CollectionType, collection: CollectionKey, rule: DelRule) {
     match rules_type {
-        RulesType::Db => del_rule_db(collection, rule).unwrap_or_else(|e| trap(&e)),
-        RulesType::Storage => del_rule_storage(collection, rule).unwrap_or_else(|e| trap(&e)),
+        CollectionType::Db => del_rule_db(collection, rule).unwrap_or_else(|e| trap(&e)),
+        CollectionType::Storage => del_rule_storage(collection, rule).unwrap_or_else(|e| trap(&e)),
     }
 }
 
@@ -549,7 +550,16 @@ pub fn get_many_assets(
 // User usage
 // ---------------------------------------------------------
 
-pub fn get_user_usage(collection: &CollectionKey, user_id: &Option<UserId>) -> Option<UserUsage> {
+pub fn get_user_usage(
+    collection: &CollectionKey,
+    collection_type: &CollectionType,
+    user_id: &Option<UserId>,
+) -> Option<UserUsage> {
     let caller = caller();
-    get_user_usage_by_id(collection, &user_id.unwrap_or(caller), caller)
+    let user_id_or_caller = user_id.unwrap_or(caller);
+
+    match collection_type {
+        CollectionType::Db => get_db_usage_by_id(collection, &user_id_or_caller, caller),
+        CollectionType::Storage => get_storage_usage_by_id(collection, &user_id_or_caller, caller),
+    }
 }

src/libs/satellite/src/types.rs

@@ -53,8 +53,8 @@ pub mod interface {
     use junobuild_storage::types::config::StorageConfig;
     use serde::Deserialize;
 
-    #[derive(CandidType, Deserialize)]
-    pub enum RulesType {
+    #[derive(CandidType, Deserialize, Serialize, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+    pub enum CollectionType {
         Db,
         Storage,
     }

src/libs/satellite/src/usage/impls.rs

@@ -1,13 +1,14 @@
+use crate::types::interface::CollectionType;
 use crate::usage::types::interface::ModificationType;
 use crate::usage::types::state::{UserUsage, UserUsageKey};
 use ic_cdk::api::time;
 use ic_stable_structures::storable::Bound;
 use ic_stable_structures::Storable;
+use junobuild_collections::types::core::CollectionKey;
 use junobuild_shared::constants::INITIAL_VERSION;
 use junobuild_shared::serializers::{deserialize_from_bytes, serialize_to_bytes};
 use junobuild_shared::types::state::{Timestamp, UserId, Version};
 use std::borrow::Cow;
-use junobuild_collections::types::core::CollectionKey;
 
 impl Storable for UserUsage {
     fn to_bytes(&self) -> Cow<[u8]> {
@@ -76,10 +77,15 @@ impl UserUsage {
 }
 
 impl UserUsageKey {
-    pub fn new(user_id: &UserId, collection: &CollectionKey) -> Self {
+    pub fn new(
+        user_id: &UserId,
+        collection_key: &CollectionKey,
+        collection_type: &CollectionType,
+    ) -> Self {
         Self {
             user_id: *user_id,
-            collection: collection.clone(),
+            collection_key: collection_key.clone(),
+            collection_type: collection_type.clone(),
         }
     }
-}
+}

src/libs/satellite/src/usage/store.rs

@@ -1,23 +1,37 @@
 use crate::memory::STATE;
+use crate::types::interface::CollectionType;
 use crate::usage::types::interface::ModificationType;
 use crate::usage::types::state::{UserUsage, UserUsageKey, UserUsageStable};
 use junobuild_collections::types::core::CollectionKey;
 use junobuild_shared::types::state::UserId;
 
-pub fn get_user_usage(user_id: &UserId, collection: &CollectionKey) -> Option<UserUsage> {
-    STATE.with(|state| get_user_usage_impl(user_id, collection, &state.borrow().stable.user_usage))
+pub fn get_user_usage(
+    user_id: &UserId,
+    collection_key: &CollectionKey,
+    collection_type: &CollectionType,
+) -> Option<UserUsage> {
+    STATE.with(|state| {
+        get_user_usage_impl(
+            user_id,
+            collection_key,
+            collection_type,
+            &state.borrow().stable.user_usage,
+        )
+    })
 }
 
 pub fn update_user_usage(
     user_id: &UserId,
-    collection: &CollectionKey,
+    collection_key: &CollectionKey,
+    collection_type: &CollectionType,
     modification: &ModificationType,
     count: Option<u32>,
 ) {
     STATE.with(|state| {
         update_user_usage_impl(
             user_id,
-            collection,
+            collection_key,
+            collection_type,
             modification,
             count,
             &mut state.borrow_mut().stable.user_usage,
@@ -27,22 +41,24 @@ pub fn update_user_usage(
 
 fn get_user_usage_impl(
     user_id: &UserId,
-    collection: &CollectionKey,
+    collection_key: &CollectionKey,
+    collection_type: &CollectionType,
     state: &UserUsageStable,
 ) -> Option<UserUsage> {
-    let key = UserUsageKey::new(user_id, collection);
+    let key = UserUsageKey::new(user_id, collection_key, collection_type);
 
     state.get(&key)
 }
 
 fn update_user_usage_impl(
     user_id: &UserId,
-    collection: &CollectionKey,
+    collection_key: &CollectionKey,
+    collection_type: &CollectionType,
     modification: &ModificationType,
     count: Option<u32>,
     state: &mut UserUsageStable,
 ) {
-    let key = UserUsageKey::new(user_id, collection);
+    let key = UserUsageKey::new(user_id, collection_key, collection_type);
 
     let current_usage = state.get(&key);
 

src/libs/satellite/src/usage/types.rs

@@ -1,4 +1,5 @@
 pub mod state {
+    use crate::types::interface::CollectionType;
     use candid::{CandidType, Deserialize};
     use ic_stable_structures::StableBTreeMap;
     use junobuild_collections::types::core::CollectionKey;
@@ -11,7 +12,8 @@ pub mod state {
     #[derive(CandidType, Serialize, Deserialize, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
     pub struct UserUsageKey {
         pub user_id: UserId,
-        pub collection: CollectionKey,
+        pub collection_key: CollectionKey,
+        pub collection_type: CollectionType,
     }
 
     #[derive(CandidType, Serialize, Deserialize, Clone)]

src/libs/satellite/src/usage/user_usage.rs

@@ -1,4 +1,5 @@
 use crate::get_controllers;
+use crate::types::interface::CollectionType;
 use crate::usage::store;
 use crate::usage::store::get_user_usage as get_user_usage_store;
 use crate::usage::types::interface::ModificationType;
@@ -9,28 +10,63 @@ use junobuild_shared::controllers::is_controller;
 use junobuild_shared::types::state::{Controllers, UserId};
 use junobuild_shared::utils::principal_not_anonymous_and_equal;
 
-pub fn get_user_usage_by_id(
+pub fn get_db_usage_by_id(
     collection: &CollectionKey,
     user_id: &UserId,
     caller: Principal,
+) -> Option<UserUsage> {
+    get_user_usage_by_id(collection, &CollectionType::Db, user_id, caller)
+}
+
+pub fn get_storage_usage_by_id(
+    collection: &CollectionKey,
+    user_id: &UserId,
+    caller: Principal,
+) -> Option<UserUsage> {
+    get_user_usage_by_id(collection, &CollectionType::Storage, user_id, caller)
+}
+
+fn get_user_usage_by_id(
+    collection_key: &CollectionKey,
+    collection_type: &CollectionType,
+    user_id: &UserId,
+    caller: Principal,
 ) -> Option<UserUsage> {
     let controllers: Controllers = get_controllers();
 
     if principal_not_anonymous_and_equal(*user_id, caller) || is_controller(caller, &controllers) {
-        return get_user_usage_store(user_id, collection);
+        return get_user_usage_store(user_id, collection_key, collection_type);
     }
 
     None
 }
 
-pub fn increase_user_usage(user_id: &UserId, collection: &CollectionKey) {
-    store::update_user_usage(user_id, collection, &ModificationType::Set, None);
+pub fn increase_db_usage(user_id: &UserId, collection: &CollectionKey) {
+    store::update_user_usage(
+        user_id,
+        collection,
+        &CollectionType::Db,
+        &ModificationType::Set,
+        None,
+    );
 }
 
-pub fn decrease_user_usage(user_id: &UserId, collection: &CollectionKey) {
-    store::update_user_usage(user_id, collection, &ModificationType::Delete, None);
+pub fn decrease_db_usage(user_id: &UserId, collection: &CollectionKey) {
+    store::update_user_usage(
+        user_id,
+        collection,
+        &CollectionType::Db,
+        &ModificationType::Delete,
+        None,
+    );
 }
 
-pub fn decrease_user_usage_by(user_id: &UserId, collection: &CollectionKey, count: u32) {
-    store::update_user_usage(user_id, collection, &ModificationType::Delete, Some(count));
+pub fn decrease_db_usage_by(user_id: &UserId, collection: &CollectionKey, count: u32) {
+    store::update_user_usage(
+        user_id,
+        collection,
+        &CollectionType::Db,
+        &ModificationType::Delete,
+        Some(count),
+    );
 }