SchemaBatch and ChangeSet changes immutable sov_state::Storage migration

src/cache/change_set.rs

@@ -25,9 +25,14 @@ impl ChangeSet {
         }
     }
 
+    /// Create new `ChangeSet
+    pub fn new_with_operations(id: SnapshotId, operations: SchemaBatch) -> Self {
+        Self { id, operations }
+    }
+
     /// Get value from its own cache
     pub fn get<S: Schema>(&self, key: &impl KeyCodec<S>) -> anyhow::Result<Option<&Operation>> {
-        self.operations.get(key)
+        self.operations.get_operation(key)
     }
 
     /// Get the ID of this [`ChangeSet`].

src/lib.rs

@@ -313,6 +313,7 @@ pub type SchemaValue = Vec<u8>;
 /// Represents operation written to the database.
 #[cfg_attr(feature = "arbitrary", derive(proptest_derive::Arbitrary))]
 #[derive(Debug, PartialEq, Eq, Hash, Clone)]
+// TODO: Do we want "generic" operation, with 2 options: S::Value and SchemaValue?
 pub enum Operation {
     /// Writing a value to the DB.
     Put {

src/schema_batch.rs

@@ -19,6 +19,7 @@ impl SchemaBatch {
     }
 
     /// Adds an insert/update operation to the batch.
+    // TODO: made this pub(crate) ???
     pub fn put<S: Schema>(
         &mut self,
         key: &impl KeyCodec<S>,
@@ -38,6 +39,7 @@ impl SchemaBatch {
     }
 
     /// Adds a delete operation to the batch.
+    // TODO: Make this pub(crate)
     pub fn delete<S: Schema>(&mut self, key: &impl KeyCodec<S>) -> anyhow::Result<()> {
         let key = key.encode_key()?;
         self.insert_operation::<S>(key, Operation::Delete);
@@ -50,7 +52,8 @@ impl SchemaBatch {
         column_writes.insert(key, operation);
     }
 
-    pub(crate) fn get<S: Schema>(
+    /// Getting the operation from current schema batch if present
+    pub(crate) fn get_operation<S: Schema>(
         &self,
         key: &impl KeyCodec<S>,
     ) -> anyhow::Result<Option<&Operation>> {
@@ -63,6 +66,17 @@ impl SchemaBatch {
         }
     }
 
+    /// Getting value by key if it was written in this batch.
+    /// Deleted operation will return None as well as missing key
+    pub fn get_value<S: Schema>(&self, key: &impl KeyCodec<S>) -> anyhow::Result<Option<S::Value>> {
+        let operation = self.get_operation(key)?;
+        if let Some(operation) = operation {
+            let value = operation.decode_value::<S>()?;
+            return Ok(value);
+        }
+        Ok(None)
+    }
+
     /// Iterator over all values in lexicographic order.
     pub fn iter<S: Schema>(&self) -> btree_map::Iter<SchemaKey, Operation> {
         self.last_writes
@@ -319,7 +333,7 @@ mod tests {
 
             let get_value = |field: &TestField| -> Option<TestField> {
                 batch1
-                    .get::<TestSchema1>(field)
+                    .get_operation::<TestSchema1>(field)
                     .unwrap()
                     .unwrap()
                     .decode_value::<TestSchema1>()