diff --git a/merk/src/element/get.rs b/merk/src/element/get.rs index 9b71827d4..73c2540ae 100644 --- a/merk/src/element/get.rs +++ b/merk/src/element/get.rs @@ -609,8 +609,9 @@ mod tests { .unwrap() .unwrap(); + let batch2 = StorageBatch::new(); let ctx = storage - .get_transactional_storage_context(SubtreePath::empty(), None, &transaction) + .get_transactional_storage_context(SubtreePath::empty(), Some(&batch2), &transaction) .unwrap(); let mut merk = Merk::open_base( ctx, diff --git a/merk/src/merk/mod.rs b/merk/src/merk/mod.rs index 5507b85d5..6aa85b292 100644 --- a/merk/src/merk/mod.rs +++ b/merk/src/merk/mod.rs @@ -1291,10 +1291,15 @@ mod test { let storage = RocksDbStorage::default_rocksdb_with_path(tmp_dir.path()) .expect("cannot open rocksdb storage"); let transaction = storage.start_transaction(); + let storage_batch = StorageBatch::new(); let mut merk = Merk::open_base( storage - .get_transactional_storage_context(SubtreePath::empty(), None, &transaction) + .get_transactional_storage_context( + SubtreePath::empty(), + Some(&storage_batch), + &transaction, + ) .unwrap(), TreeType::NormalTree, None::<&fn(&[u8], &GroveVersion) -> Option>, @@ -1319,10 +1324,15 @@ mod test { let storage = RocksDbStorage::default_rocksdb_with_path(tmp_dir.path()) .expect("cannot open rocksdb storage"); let transaction = storage.start_transaction(); + let storage_batch = StorageBatch::new(); let mut merk = Merk::open_base( storage - .get_transactional_storage_context(SubtreePath::empty(), None, &transaction) + .get_transactional_storage_context( + SubtreePath::empty(), + Some(&storage_batch), + &transaction, + ) .unwrap(), TreeType::NormalTree, None::<&fn(&[u8], &GroveVersion) -> Option>, @@ -1605,9 +1615,14 @@ mod test { .unwrap() .expect("cannot commit batch"); + let batch2 = StorageBatch::new(); let mut merk = Merk::open_base( storage - .get_transactional_storage_context(SubtreePath::empty(), None, &transaction) + .get_transactional_storage_context( + SubtreePath::empty(), + Some(&batch2), + &transaction, + ) .unwrap(), TreeType::NormalTree, None::<&fn(&[u8], &GroveVersion) -> Option>, diff --git a/storage/src/rocksdb_storage/storage_context/context_tx.rs b/storage/src/rocksdb_storage/storage_context/context_tx.rs index 671f3be5c..53b399da2 100644 --- a/storage/src/rocksdb_storage/storage_context/context_tx.rs +++ b/storage/src/rocksdb_storage/storage_context/context_tx.rs @@ -124,14 +124,21 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { children_sizes: ChildrenSizesWithIsSumTree, cost_info: Option, ) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.put( - make_prefixed_key(&self.prefix, key), - value.to_vec(), - children_sizes, - cost_info, - ); - } + let batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted put operation on transactional context without a batch".to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + batch.put( + make_prefixed_key(&self.prefix, key), + value.to_vec(), + children_sizes, + cost_info, + ); Ok(()).wrap_with_cost(OperationCost::default()) } @@ -141,13 +148,21 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { value: &[u8], cost_info: Option, ) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.put_aux( - make_prefixed_key(&self.prefix, key), - value.to_vec(), - cost_info, - ); - } + let batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted put_aux operation on transactional context without a batch" + .to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + batch.put_aux( + make_prefixed_key(&self.prefix, key), + value.to_vec(), + cost_info, + ); Ok(()).wrap_with_cost(OperationCost::default()) } @@ -157,13 +172,21 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { value: &[u8], cost_info: Option, ) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.put_root( - make_prefixed_key(&self.prefix, key), - value.to_vec(), - cost_info, - ); - } + let batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted put_root operation on transactional context without a batch" + .to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + batch.put_root( + make_prefixed_key(&self.prefix, key), + value.to_vec(), + cost_info, + ); Ok(()).wrap_with_cost(OperationCost::default()) } @@ -173,13 +196,21 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { value: &[u8], cost_info: Option, ) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.put_meta( - make_prefixed_key(&self.prefix, key), - value.to_vec(), - cost_info, - ); - } + let batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted put_meta operation on transactional context without a batch" + .to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + batch.put_meta( + make_prefixed_key(&self.prefix, key), + value.to_vec(), + cost_info, + ); Ok(()).wrap_with_cost(OperationCost::default()) } @@ -188,10 +219,17 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { key: K, cost_info: Option, ) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.delete(make_prefixed_key(&self.prefix, key), cost_info); - } - + let batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted delete operation on transactional context without a batch" + .to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + batch.delete(make_prefixed_key(&self.prefix, key), cost_info); Ok(()).wrap_with_cost(OperationCost::default()) } @@ -200,10 +238,17 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { key: K, cost_info: Option, ) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.delete_aux(make_prefixed_key(&self.prefix, key), cost_info); - } - + let batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted delete_aux operation on transactional context without a batch" + .to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + batch.delete_aux(make_prefixed_key(&self.prefix, key), cost_info); Ok(()).wrap_with_cost(OperationCost::default()) } @@ -212,10 +257,17 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { key: K, cost_info: Option, ) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.delete_root(make_prefixed_key(&self.prefix, key), cost_info); - } - + let batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted delete_root operation on transactional context without a batch" + .to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + batch.delete_root(make_prefixed_key(&self.prefix, key), cost_info); Ok(()).wrap_with_cost(OperationCost::default()) } @@ -224,10 +276,17 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { key: K, cost_info: Option, ) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.delete_meta(make_prefixed_key(&self.prefix, key), cost_info); - } - + let batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted delete_meta operation on transactional context without a batch" + .to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + batch.delete_meta(make_prefixed_key(&self.prefix, key), cost_info); Ok(()).wrap_with_cost(OperationCost::default()) } @@ -303,10 +362,17 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> { } fn commit_batch(&self, batch: Self::Batch) -> CostResult<(), Error> { - if let Some(existing_batch) = self.batch { - existing_batch.merge(batch.batch); - } - + let existing_batch = match self.batch { + Some(b) => b, + None => { + return Err(Error::StorageError( + "attempted commit_batch operation on transactional context without a batch" + .to_string(), + )) + .wrap_with_cost(OperationCost::default()) + } + }; + existing_batch.merge(batch.batch); Ok(()).wrap_with_cost(OperationCost::default()) } diff --git a/storage/src/rocksdb_storage/tests.rs b/storage/src/rocksdb_storage/tests.rs index f80b4f9d0..caff56162 100644 --- a/storage/src/rocksdb_storage/tests.rs +++ b/storage/src/rocksdb_storage/tests.rs @@ -1446,3 +1446,125 @@ mod storage_management { ); } } + +mod transactional_context_without_batch { + use super::*; + use crate::{Storage, StorageBatch, StorageContext}; + + #[test] + fn test_write_operations_error_when_batch_is_none() { + let storage = TempStorage::new(); + let transaction = storage.start_transaction(); + + // Create a transactional context with batch = None (read-only context) + let context = storage + .get_transactional_storage_context([b"test"].as_ref().into(), None, &transaction) + .unwrap(); + + // All write operations should return an error instead of silently succeeding + + // put + let result = context.put(b"key", b"value", None, None).unwrap(); + assert!( + result.is_err(), + "put should fail on transactional context without a batch" + ); + + // put_aux + let result = context.put_aux(b"key", b"value", None).unwrap(); + assert!( + result.is_err(), + "put_aux should fail on transactional context without a batch" + ); + + // put_root + let result = context.put_root(b"key", b"value", None).unwrap(); + assert!( + result.is_err(), + "put_root should fail on transactional context without a batch" + ); + + // put_meta + let result = context.put_meta(b"key", b"value", None).unwrap(); + assert!( + result.is_err(), + "put_meta should fail on transactional context without a batch" + ); + + // delete + let result = context.delete(b"key", None).unwrap(); + assert!( + result.is_err(), + "delete should fail on transactional context without a batch" + ); + + // delete_aux + let result = context.delete_aux(b"key", None).unwrap(); + assert!( + result.is_err(), + "delete_aux should fail on transactional context without a batch" + ); + + // delete_root + let result = context.delete_root(b"key", None).unwrap(); + assert!( + result.is_err(), + "delete_root should fail on transactional context without a batch" + ); + + // delete_meta + let result = context.delete_meta(b"key", None).unwrap(); + assert!( + result.is_err(), + "delete_meta should fail on transactional context without a batch" + ); + + // commit_batch + let new_batch = context.new_batch(); + let result = context.commit_batch(new_batch).unwrap(); + assert!( + result.is_err(), + "commit_batch should fail on transactional context without a batch" + ); + } + + #[test] + fn test_read_operations_succeed_when_batch_is_none() { + let storage = TempStorage::new(); + let transaction = storage.start_transaction(); + + // First, write some data using a context WITH a batch + let batch = StorageBatch::new(); + let write_context = storage + .get_transactional_storage_context( + [b"test"].as_ref().into(), + Some(&batch), + &transaction, + ) + .unwrap(); + write_context + .put(b"key", b"value", None, None) + .unwrap() + .expect("put with batch should succeed"); + storage + .commit_multi_context_batch(batch, Some(&transaction)) + .unwrap() + .expect("commit batch should succeed"); + + // Now create a read-only context (batch = None) and verify reads work + let read_context = storage + .get_transactional_storage_context([b"test"].as_ref().into(), None, &transaction) + .unwrap(); + + let result = read_context.get(b"key").unwrap(); + assert!( + result.is_ok(), + "get should succeed on transactional context without a batch" + ); + assert_eq!( + result.unwrap(), + Some(b"value".to_vec()), + "get should return the correct value" + ); + } +}