From d8e3d48caa403470969fa219e8069c371f22051e Mon Sep 17 00:00:00 2001
From: Mahmoud-Emad <mahmmoud.hassanein@gmail.com>
Date: Tue, 27 May 2025 21:15:08 +0300
Subject: [PATCH] feat: Improve database handling of model indices

- Ensure all model types always have a primary key index entry.
- Simplify `get_all()` implementation by relying on primary key index.
- Remove redundant code in `get_all()` for finding all object IDs.
- Improve error handling and clarity in database operations.
---
 heromodels/src/db/hero.rs | 64 +++++++++++++++++++++------------------
 1 file changed, 35 insertions(+), 29 deletions(-)

diff --git a/heromodels/src/db/hero.rs b/heromodels/src/db/hero.rs
index dff28b1..47fd7c4 100644
--- a/heromodels/src/db/hero.rs
+++ b/heromodels/src/db/hero.rs
@@ -323,6 +323,16 @@ where
                 assigned_id
             };
 
+            // Always create a primary key index entry for this model type
+            // This ensures get_all() can find all objects of this type, even if they have no explicit indexed fields
+            let primary_index_key = format!("{}::primary", M::db_prefix());
+            let mut primary_ids: HashSet<u32> =
+                Self::get_tst_value(&mut index_db, &primary_index_key)?
+                    .unwrap_or_else(HashSet::new);
+            primary_ids.insert(assigned_id);
+            let raw_primary_ids = bincode::serde::encode_to_vec(&primary_ids, BINCODE_CONFIG)?;
+            index_db.set(&primary_index_key, raw_primary_ids)?;
+
             // Now add the new indices
             for index_key in indices_to_add {
                 let key = Self::index_key(M::db_prefix(), index_key.name, &index_key.value);
@@ -420,6 +430,22 @@ where
             }
         }
 
+        // Also remove from the primary key index
+        let primary_index_key = format!("{}::primary", M::db_prefix());
+        if let Some(mut primary_ids) =
+            Self::get_tst_value::<HashSet<u32>>(&mut index_db, &primary_index_key)?
+        {
+            primary_ids.remove(&id);
+            if primary_ids.is_empty() {
+                // This was the last object of this type, remove the primary index entirely
+                index_db.delete(&primary_index_key)?;
+            } else {
+                // There are still other objects of this type, write back updated set
+                let raw_primary_ids = bincode::serde::encode_to_vec(&primary_ids, BINCODE_CONFIG)?;
+                index_db.set(&primary_index_key, raw_primary_ids)?;
+            }
+        }
+
         // Finally delete the object itself
         Ok(data_db.delete(id)?)
     }
@@ -428,36 +454,16 @@ where
         let mut index_db = self.index.lock().expect("can lock index DB");
         let mut data_db = self.data.lock().expect("can lock data DB");
 
-        let prefix = M::db_prefix();
-        let mut all_object_ids: HashSet<u32> = HashSet::new();
-
-        // Use getall to find all index entries (values are serialized HashSet<u32>) for the given model prefix.
-        match index_db.getall(prefix) {
-            Ok(list_of_raw_ids_set_bytes) => {
-                for raw_ids_set_bytes in list_of_raw_ids_set_bytes {
-                    // Each item in the list is a bincode-serialized HashSet<u32> of object IDs.
-                    match bincode::serde::decode_from_slice::<HashSet<u32>, _>(&raw_ids_set_bytes, BINCODE_CONFIG) {
-                        Ok((ids_set, _)) => { // Destructure the tuple (HashSet<u32>, usize)
-                            all_object_ids.extend(ids_set);
-                        }
-                        Err(e) => {
-                            // If deserialization of an ID set fails, propagate as a decode error.
-                            return Err(super::Error::Decode(e));
-                        }
-                    }
+        // Look for the primary key index entry for this model type
+        let primary_index_key = format!("{}::primary", M::db_prefix());
+        let all_object_ids: HashSet<u32> =
+            match Self::get_tst_value(&mut index_db, &primary_index_key)? {
+                Some(ids) => ids,
+                None => {
+                    // No primary index found, meaning no objects of this type exist
+                    return Ok(Vec::new());
                 }
-            }
-            Err(tst::Error::PrefixNotFound(_)) => {
-                // No index entries found for this prefix, meaning no objects of this type exist.
-                // Note: tst::getall might return Ok(vec![]) in this case instead of PrefixNotFound.
-                // Depending on tst implementation, this arm might be redundant if getall returns empty vec.
-                return Ok(Vec::new());
-            }
-            Err(e) => {
-                // Other TST errors.
-                return Err(super::Error::DB(e));
-            }
-        }
+            };
 
         let mut results: Vec<M> = Vec::with_capacity(all_object_ids.len());
         for obj_id in all_object_ids {