refactor generic put_cf method

dan_layer/state_store_rocksdb/src/model/block.rs

@@ -101,23 +101,6 @@ impl RocksdbModel for BlockModel {
         Self::key_from_block_id(item.id())
     }
 
-    fn column_families() -> Vec<&'static str> {
-        vec![ParentIdColumnFamily::name(), EpochHeightColumnFamily::name(), IsCommittedColumnFamily::name()]
-    }
-
-    fn put_cf(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, cf_name: &str, value: &Self::Item) -> Result<(), RocksDbStorageError> {
-        // In each CF value We store the BlockId of the block, so we can reference it back
-        let cf_value = value.id().as_bytes();
-        match cf_name {
-            ParentIdColumnFamily::NAME => ParentIdColumnFamily::put(db, tx, operation,  value, cf_value)?,
-            EpochHeightColumnFamily::NAME => EpochHeightColumnFamily::put(db, tx, operation, value, cf_value)?,
-            IsCommittedColumnFamily::NAME => IsCommittedColumnFamily::put(db, tx, operation, value, cf_value)?,
-            _ => (),
-        }
-
-        Ok(())
-    }
-
     fn get_cf(db: Arc<TransactionDB>, tx: &Transaction<'_, TransactionDB>, cf_name: &str, operation: &'static str, key_prefix: &str, ordering: Ordering) -> Result<Option<Self::Item>, RocksDbStorageError> {
         let cf = db.cf_handle(cf_name).unwrap();
 
@@ -151,10 +134,25 @@ impl RocksdbModel for BlockModel {
         Ok(Some(value))
     }
 
+    fn column_families() -> Vec<&'static str> {
+        vec![ParentIdColumnFamily::name(), EpochHeightColumnFamily::name(), IsCommittedColumnFamily::name()]
+    }
+
+    fn put_in_cfs(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, value: &Self::Item) -> Result<(), RocksDbStorageError> {
+        // In each CF value We store the BlockId of the block, so we can reference it back
+        let cf_value = value.id().as_bytes();
+
+        ParentIdColumnFamily::put(db.clone(), tx, operation,  value, cf_value)?;
+        EpochHeightColumnFamily::put(db.clone(), tx, operation, value, cf_value)?;
+        IsCommittedColumnFamily::put(db, tx, operation, value, cf_value)?;
+
+        Ok(())
+    }
+
     fn delete_from_cfs(db: Arc<TransactionDB>, tx: &Transaction<'_, TransactionDB>, operation: &'static str, item: &Self::Item) -> Result<(), RocksDbStorageError> {
         ParentIdColumnFamily::delete(db.clone(), tx, operation, item)?;
         EpochHeightColumnFamily::delete(db.clone(), tx, operation, item)?;
-        IsCommittedColumnFamily::delete(db.clone(), tx, operation, item)?;
+        IsCommittedColumnFamily::delete(db, tx, operation, item)?;
         Ok(())
     }
 }

dan_layer/state_store_rocksdb/src/model/encoding.rs

@@ -20,7 +20,7 @@
 //  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 //  USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-use serde::{de::DeserializeOwned, Deserialize, Serialize};
+use serde::{de::DeserializeOwned, Serialize};
 
 use crate::error::RocksDbStorageError;
 

dan_layer/state_store_rocksdb/src/model/model.rs

@@ -93,14 +93,16 @@ pub trait RocksdbModel {
         })?;
 
         // insert the main key in each of the column families
-        for cf_name in Self::column_families() {
-            Self::put_cf(db.clone(), tx, operation, &cf_name, &value)?;
-        }
+        Self::put_in_cfs(db.clone(), tx, operation, &value)?;
 
         Ok(())
     }
 
-    fn put_cf(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, cf_name: &str, value: &Self::Item) -> Result<(), RocksDbStorageError>;
+    fn put_in_cfs(_db: Arc<TransactionDB>, _tx: &mut Transaction<'_, TransactionDB>, _operation: &'static str, _value: &Self::Item) -> Result<(), RocksDbStorageError> {
+        // It's up to concrete models to override this method
+        // We provide a default implementation to simplify all the models that do not have column families
+        Ok(())
+    }
 
     fn key_exists(tx: &Transaction<'_, TransactionDB>, operation: &'static str, key: &str) -> Result<bool, RocksDbStorageError> {
         let value = tx.get(&key)

dan_layer/state_store_rocksdb/src/model/state_transition.rs

@@ -72,15 +72,19 @@ impl RocksdbModel for StateTransitionModel {
         vec![ShardColumnFamily::name()]
     }
 
-    fn put_cf(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, cf_name: &str, value: &Self::Item) -> Result<(), RocksDbStorageError> {
+    fn put_in_cfs(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, value: &Self::Item) -> Result<(), RocksDbStorageError> {
         // In each CF value We store the key to the main collection, so we can retrieve the actual value
         let main_key = Self::key(value);
         let main_key_bytes = main_key.as_bytes();
-        match cf_name {
-            ShardColumnFamily::NAME => ShardColumnFamily::put(db, tx, operation,  value, main_key_bytes)?,
-            _ => (),
-        }
 
+        ShardColumnFamily::put(db, tx, operation,  value, main_key_bytes)?;
+
+        Ok(())
+    }
+
+    fn delete_from_cfs(db: Arc<TransactionDB>, tx: &Transaction<'_, TransactionDB>, operation: &'static str, item: &Self::Item) -> Result<(), RocksDbStorageError> {
+        ShardColumnFamily::delete(db.clone(), tx, operation, item)?;
+
         Ok(())
     }
 }

dan_layer/state_store_rocksdb/src/model/state_tree_shard_versions.rs

@@ -20,15 +20,10 @@
 //  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 //  USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-use std::sync::Arc;
-
-use rocksdb::{Transaction, TransactionDB};
 use serde::{Deserialize, Serialize};
 use tari_dan_common_types::shard::Shard;
 use tari_state_tree::Version;
 
-use crate::error::RocksDbStorageError;
-
 use super::model::RocksdbModel;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -59,9 +54,4 @@ impl RocksdbModel for StateTreeShardVersionModel {
     fn column_families() -> Vec<&'static str> {
         vec![]
     }
-
-    fn put_cf(_db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, _operation: &'static str, _cf_name: &str, _value: &Self::Item) -> Result<(), RocksDbStorageError> {
-        // No column families for this model
-        Ok(())
-    }
 }

dan_layer/state_store_rocksdb/src/model/substate.rs

@@ -133,17 +133,23 @@ impl RocksdbModel for SubstateModel {
         vec![VersionColumnFamily::name(), CreatedByTxColumnFamily::name(), DestroyedByTxColumnFamily::name()]
     }
 
-    fn put_cf(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, cf_name: &str, value: &Self::Item) -> Result<(), RocksDbStorageError> {
+    fn put_in_cfs(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, value: &Self::Item) -> Result<(), RocksDbStorageError> {
         // In each CF value We store the key to the main collection, so we can retrieve the actual value
         let main_key = Self::key(value);
         let main_key_bytes = main_key.as_bytes();
-        match cf_name {
-            VersionColumnFamily::NAME => VersionColumnFamily::put(db, tx, operation,  value, main_key_bytes)?,
-            CreatedByTxColumnFamily::NAME => CreatedByTxColumnFamily::put(db, tx, operation, value, main_key_bytes)?,
-            DestroyedByTxColumnFamily::NAME => DestroyedByTxColumnFamily::put(db, tx, operation, value, main_key_bytes)?,
-            _ => (),
-        }
 
+        VersionColumnFamily::put(db.clone(), tx, operation,  value, main_key_bytes)?;
+        CreatedByTxColumnFamily::put(db.clone(), tx, operation, value, main_key_bytes)?;
+        DestroyedByTxColumnFamily::put(db, tx, operation, value, main_key_bytes)?;
+
+        Ok(())
+    }
+
+    fn delete_from_cfs(db: Arc<TransactionDB>, tx: &Transaction<'_, TransactionDB>, operation: &'static str, item: &Self::Item) -> Result<(), RocksDbStorageError> {
+        VersionColumnFamily::delete(db.clone(), tx, operation, item)?;
+        CreatedByTxColumnFamily::delete(db.clone(), tx, operation, item)?;
+        DestroyedByTxColumnFamily::delete(db, tx, operation, item)?;
+
         Ok(())
     }
 }