feat: use a single usage collection

Author: peterpeterparker

src/libs/collections/src/constants.rs

@@ -19,8 +19,7 @@ pub const DEFAULT_DB_LOG_RULE: SetRule = SetRule {
     rate_config: None,
 };
 
-pub const USER_USAGE_DB_COLLECTION_KEY: &str = "#user_usage_db";
-pub const USER_USAGE_STORAGE_COLLECTION_KEY: &str = "#user_usage_storage";
+pub const USER_USAGE_COLLECTION_KEY: &str = "#user-usage";
 
 pub const DEFAULT_USER_USAGE_RULE: SetRule = SetRule {
     read: Managed,
@@ -34,7 +33,7 @@ pub const DEFAULT_USER_USAGE_RULE: SetRule = SetRule {
     rate_config: None,
 };
 
-pub const DEFAULT_DB_COLLECTIONS: [(&str, SetRule); 4] = [
+pub const DEFAULT_DB_COLLECTIONS: [(&str, SetRule); 3] = [
     (
         "#user",
         SetRule {
@@ -50,8 +49,7 @@ pub const DEFAULT_DB_COLLECTIONS: [(&str, SetRule); 4] = [
         },
     ),
     (LOG_COLLECTION_KEY, DEFAULT_DB_LOG_RULE),
-    (USER_USAGE_DB_COLLECTION_KEY, DEFAULT_USER_USAGE_RULE),
-    (USER_USAGE_STORAGE_COLLECTION_KEY, DEFAULT_USER_USAGE_RULE),
+    (USER_USAGE_COLLECTION_KEY, DEFAULT_USER_USAGE_RULE),
 ];
 
 pub const DB_COLLECTIONS_NO_USER_USAGE: [&str; 1] = [LOG_COLLECTION_KEY];

src/libs/satellite/src/impls.rs

@@ -1,5 +1,6 @@
+use std::fmt::{Display, Formatter, Result as FmtResult};
 use crate::memory::init_stable_state;
-use crate::types::state::{HeapState, RuntimeState, State};
+use crate::types::state::{CollectionType, HeapState, RuntimeState, State};
 
 impl Default for State {
     fn default() -> Self {
@@ -10,3 +11,12 @@ impl Default for State {
         }
     }
 }
+
+impl Display for CollectionType {
+    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
+        write!(f, "{}", match self {
+            CollectionType::Db => "db",
+            CollectionType::Storage => "storage",
+        })
+    }
+}

src/libs/satellite/src/rules/upgrade.rs

@@ -1,34 +1,21 @@
 use crate::memory::STATE;
 use crate::rules::store::set_rule_db;
 use junobuild_collections::constants::{
-    DEFAULT_USER_USAGE_RULE, USER_USAGE_DB_COLLECTION_KEY, USER_USAGE_STORAGE_COLLECTION_KEY,
+    DEFAULT_USER_USAGE_RULE, USER_USAGE_COLLECTION_KEY,
 };
 
 /// One time upgrade
 
 pub fn init_user_usage_collections() {
     let col = STATE.with(|state| {
         let rules = &state.borrow_mut().heap.db.rules;
-        rules.get(USER_USAGE_DB_COLLECTION_KEY).cloned()
+        rules.get(USER_USAGE_COLLECTION_KEY).cloned()
     });
 
     if col.is_none() {
         // We are ignoring potential issues because we are processing during an upgrade.
         let _ = set_rule_db(
-            USER_USAGE_DB_COLLECTION_KEY.to_string(),
-            DEFAULT_USER_USAGE_RULE,
-        );
-    }
-
-    let col_storage = STATE.with(|state| {
-        let rules = &state.borrow_mut().heap.db.rules;
-        rules.get(USER_USAGE_STORAGE_COLLECTION_KEY).cloned()
-    });
-
-    if col_storage.is_none() {
-        // We are ignoring potential issues because we are processing during an upgrade.
-        let _ = set_rule_db(
-            USER_USAGE_STORAGE_COLLECTION_KEY.to_string(),
+            USER_USAGE_COLLECTION_KEY.to_string(),
             DEFAULT_USER_USAGE_RULE,
         );
     }

src/libs/satellite/src/usage/assert.rs

@@ -1,8 +1,6 @@
 use crate::usage::store::increment_usage;
 use crate::usage::utils::{is_db_collection_no_usage, is_storage_collection_no_usage};
-use junobuild_collections::constants::{
-    USER_USAGE_DB_COLLECTION_KEY, USER_USAGE_STORAGE_COLLECTION_KEY,
-};
+use crate::types::state::CollectionType;
 use junobuild_collections::types::core::CollectionKey;
 use junobuild_shared::controllers::is_controller;
 use junobuild_shared::types::state::{Controllers, UserId};
@@ -21,7 +19,7 @@ pub fn increment_and_assert_db_usage(
         caller,
         controllers,
         collection,
-        &USER_USAGE_DB_COLLECTION_KEY.to_string(),
+        &CollectionType::Db,
         max_changes_per_user,
     )
 }
@@ -40,7 +38,7 @@ pub fn increment_and_assert_storage_usage(
         caller,
         controllers,
         collection,
-        &USER_USAGE_STORAGE_COLLECTION_KEY.to_string(),
+        &CollectionType::Storage,
         max_changes_per_user,
     )
 }
@@ -49,15 +47,15 @@ fn increment_and_assert_usage(
     caller: UserId,
     controllers: &Controllers,
     collection: &CollectionKey,
-    user_usage_collection: &CollectionKey,
+    collection_type: &CollectionType,
     max_changes_per_user: Option<u32>,
 ) -> Result<(), String> {
     // We only collect usage for users
     if is_controller(caller, controllers) {
         return Ok(());
     }
 
-    let user_usage = increment_usage(user_usage_collection, collection, &caller)?;
+    let user_usage = increment_usage(collection, collection_type, &caller)?;
 
     if let Some(max_changes_per_user) = max_changes_per_user {
         if user_usage.changes_count > max_changes_per_user {

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

@@ -1,7 +1,9 @@
-use crate::usage::types::state::UserUsage;
+use crate::usage::types::state::{UserUsage, UserUsageKey};
 use ic_cdk::api::time;
-use junobuild_shared::types::state::{Timestamp, Version, Versioned};
+use junobuild_collections::types::core::CollectionKey;
+use junobuild_shared::types::state::{Timestamp, UserId, Version, Versioned};
 use junobuild_shared::version::next_version;
+use crate::types::state::CollectionType;
 
 impl UserUsage {
     pub fn increment(current_user_usage: &Option<UserUsage>) -> Self {
@@ -45,3 +47,23 @@ impl Versioned for UserUsage {
         self.version
     }
 }
+
+impl UserUsageKey {
+    pub fn create(
+        user_id: &UserId,
+        collection_key: &CollectionKey,
+        collection_type: &CollectionType,
+    ) -> Self {
+        Self {
+            user_id: *user_id,
+            collection_key: collection_key.clone(),
+            collection_type: collection_type.clone(),
+        }
+    }
+
+    pub fn to_key(
+        &self
+    ) -> String {
+        format!("{}#{}#{}", self.user_id.to_text(), self.collection_type, self.collection_key)
+    }
+}

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

@@ -1,40 +1,33 @@
-use crate::usage::types::state::UserUsage;
-use crate::usage::utils::build_user_usage_key;
+use crate::usage::types::state::{UserUsage, UserUsageKey};
 use crate::{get_doc_store, set_doc_store, SetDoc};
 use ic_cdk::id;
 use junobuild_collections::types::core::CollectionKey;
 use junobuild_shared::types::state::UserId;
 use junobuild_utils::{decode_doc_data, encode_doc_data};
+use junobuild_collections::constants::USER_USAGE_COLLECTION_KEY;
+use crate::types::state::CollectionType;
 
 pub fn increment_usage(
-    user_usage_collection: &CollectionKey,
-    collection: &CollectionKey,
+    collection_key: &CollectionKey,
+    collection_type: &CollectionType,
     user_id: &UserId,
 ) -> Result<UserUsage, String> {
-    let key = build_user_usage_key(user_id, collection);
+    let user_usage_key = UserUsageKey::create(user_id, collection_key, collection_type);
+    let key = user_usage_key.to_key();
 
-    let doc = get_doc_store(id(), user_usage_collection.to_string(), key.clone())?;
+    let doc = get_doc_store(id(), USER_USAGE_COLLECTION_KEY.to_string(), key.clone())?;
 
-    let current_usage: Option<UserUsage> = match &doc {
-        None => None,
-        Some(doc) => decode_doc_data(&doc.data)?,
-    };
+    let current_usage = doc.as_ref().map(|doc| decode_doc_data(&doc.data)).transpose()?;
 
     let update_usage = UserUsage::increment(&current_usage);
 
-    let update_doc: SetDoc = SetDoc {
+    let update_doc = SetDoc {
         data: encode_doc_data(&update_usage)?,
-        description: match &doc {
-            None => None,
-            Some(doc) => doc.description.clone(),
-        },
-        version: match &doc {
-            None => None,
-            Some(doc) => doc.version.clone(),
-        },
+        description: doc.as_ref().and_then(|d| d.description.clone()),
+        version: doc.as_ref().and_then(|d| d.version.clone()),
     };
 
-    set_doc_store(id(), user_usage_collection.to_string(), key, update_doc)?;
+    set_doc_store(id(), USER_USAGE_COLLECTION_KEY.to_string(), key, update_doc)?;
 
     Ok(update_usage)
 }

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

@@ -1,10 +1,22 @@
 pub mod state {
     use candid::CandidType;
-    use junobuild_shared::types::core::Key;
-    use junobuild_shared::types::state::{Timestamp, Version};
+    use junobuild_shared::types::state::{Timestamp, UserId, Version};
     use serde::{Deserialize, Serialize};
+    use junobuild_collections::types::core::CollectionKey;
+    use crate::types::state::CollectionType;
 
-    pub type UserUsageKey = Key;
+    /// A unique key for identifying user usage within a collection.
+    ///
+    /// It consists of:
+    /// - `user_id`: The unique identifier for the user which is matched to the caller.
+    /// - `collection_key`: The collection where usage is tracked.
+    /// - `collection_type`: The type of collection (`Db` for datastore, `Storage` for assets).
+    #[derive(CandidType, Serialize, Deserialize, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+    pub struct UserUsageKey {
+        pub user_id: UserId,
+        pub collection_key: CollectionKey,
+        pub collection_type: CollectionType,
+    }
 
     /// Tracks the usage (create, set and delete) of a user in a collection.
     ///

src/libs/satellite/src/usage/utils.rs

@@ -1,9 +1,7 @@
-use crate::usage::types::state::UserUsageKey;
 use junobuild_collections::constants::{
     ASSETS_COLLECTIONS_NO_USER_USAGE, DB_COLLECTIONS_NO_USER_USAGE,
 };
 use junobuild_collections::types::core::CollectionKey;
-use junobuild_shared::types::state::UserId;
 
 pub fn is_db_collection_no_usage(collection: &CollectionKey) -> bool {
     DB_COLLECTIONS_NO_USER_USAGE.contains(&collection.as_str())
@@ -12,7 +10,3 @@ pub fn is_db_collection_no_usage(collection: &CollectionKey) -> bool {
 pub fn is_storage_collection_no_usage(collection: &CollectionKey) -> bool {
     ASSETS_COLLECTIONS_NO_USER_USAGE.contains(&collection.as_str())
 }
-
-pub fn build_user_usage_key(user_id: &UserId, collection_key: &CollectionKey) -> UserUsageKey {
-    format!("{}#{}", user_id.to_text(), collection_key)
-}