From 6febf63f61bc775324b96dc1d1a0290b8a3d9e41 Mon Sep 17 00:00:00 2001
From: Nikolai Golub <nikolai@sovlabs.io>
Date: Wed, 8 May 2024 13:28:49 +0200
Subject: [PATCH] get_value

---
 src/cache/change_set.rs |  2 +-
 src/lib.rs              |  1 +
 src/schema_batch.rs     | 20 +++++++++++++++++---
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/cache/change_set.rs b/src/cache/change_set.rs
index 9da375e..bbf36f1 100644
--- a/src/cache/change_set.rs
+++ b/src/cache/change_set.rs
@@ -32,7 +32,7 @@ impl ChangeSet {
 
     /// 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`].
diff --git a/src/lib.rs b/src/lib.rs
index 65e6220..94d4ecb 100644
--- a/src/lib.rs
+++ b/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 {
diff --git a/src/schema_batch.rs b/src/schema_batch.rs
index 0fca7d5..1c830ca 100644
--- a/src/schema_batch.rs
+++ b/src/schema_batch.rs
@@ -52,8 +52,11 @@ impl SchemaBatch {
         column_writes.insert(key, operation);
     }
 
-    /// Getting the
-    pub fn get<S: Schema>(&self, key: &impl KeyCodec<S>) -> anyhow::Result<Option<&Operation>> {
+    /// 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>> {
         let key = key.encode_key()?;
 
         if let Some(column_writes) = self.last_writes.get(&S::COLUMN_FAMILY_NAME) {
@@ -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>()