From 1622d45c23905d42b3e5e3171cee54b4f4350ffb Mon Sep 17 00:00:00 2001
From: Olivier 'reivilibre <olivier@librepush.net>
Date: Thu, 24 Mar 2022 19:25:33 +0000
Subject: [PATCH] Add get and put methods

---
 src/database/raw.rs     | 37 ++++++++++++++++++++++++++++++++++++-
 src/database/wrapped.rs | 41 +++++++++++++++++++++++++++++++++++++++++
 src/environment.rs      |  2 +-
 src/wrapper.rs          | 11 ++++++++++-
 4 files changed, 88 insertions(+), 3 deletions(-)

diff --git a/src/database/raw.rs b/src/database/raw.rs
index 451ed9a..df0e641 100644
--- a/src/database/raw.rs
+++ b/src/database/raw.rs
@@ -1,5 +1,8 @@
-use libmdbx::{Database, Environment, WriteMap};
+use crate::environment::Txn;
+use crate::wrapper::ZeroCopyByteWrapper;
+use libmdbx::{Database, Environment, TransactionKind, WriteFlags, WriteMap, RW};
 use ouroboros::self_referencing;
+use std::borrow::Cow;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
@@ -14,3 +17,35 @@ pub struct RawTable<K: ?Sized, V: ?Sized> {
     pub(crate) phantom_k: PhantomData<K>,
     pub(crate) phantom_v: PhantomData<V>,
 }
+
+impl<K: ?Sized + ZeroCopyByteWrapper, V: ?Sized + ZeroCopyByteWrapper> RawTable<K, V> {
+    pub fn get<'txn, TK: TransactionKind>(
+        &'txn self,
+        txn: &'txn Txn<'txn, TK>,
+        k: impl AsRef<K>,
+    ) -> anyhow::Result<Option<Cow<'txn, V>>> {
+        let bytes = txn
+            .mdbx_txn
+            .get::<Cow<'txn, [u8]>>(self.borrow_mdbx_db(), k.as_ref().as_byte_slice())?;
+        Ok(match bytes {
+            None => None,
+            Some(Cow::Owned(owned)) => Some(Cow::Owned(V::from_owned_bytes(owned)?)),
+            Some(Cow::Borrowed(borrowed)) => Some(Cow::Borrowed(V::from_byte_slice(borrowed)?)),
+        })
+    }
+
+    pub fn put<'txn>(
+        &self,
+        txn: &Txn<'txn, RW>,
+        k: impl AsRef<K>,
+        v: impl AsRef<V>,
+    ) -> anyhow::Result<()> {
+        txn.mdbx_txn.put(
+            self.borrow_mdbx_db(),
+            k.as_ref().as_byte_slice(),
+            v.as_ref().as_byte_slice(),
+            WriteFlags::empty(),
+        )?;
+        Ok(())
+    }
+}
diff --git a/src/database/wrapped.rs b/src/database/wrapped.rs
index 81d1dab..3de0ef9 100644
--- a/src/database/wrapped.rs
+++ b/src/database/wrapped.rs
@@ -1,7 +1,48 @@
 use crate::database::raw::RawTable;
+use crate::environment::Txn;
+use crate::wrapper::ByteWrapper;
+use anyhow::Context;
+use libmdbx::{TransactionKind, RW};
 
 pub struct WrappedTable<K, V> {
     pub raw: RawTable<[u8], [u8]>,
     pub k_wrapper: K,
     pub v_wrapper: V,
 }
+
+impl<K: ByteWrapper, V: ByteWrapper> WrappedTable<K, V> {
+    pub fn get<'txn, TK: TransactionKind>(
+        &self,
+        txn: &Txn<'txn, TK>,
+        k: impl AsRef<K::Item>,
+    ) -> anyhow::Result<Option<V::Item>> {
+        let k_bytes = self
+            .k_wrapper
+            .dump_to_db_bytes(k.as_ref())
+            .context("whilst converting key to bytes")?;
+
+        self.raw
+            .get(txn, k_bytes)?
+            .map(|v_bytes| self.v_wrapper.load_from_db_bytes(&v_bytes))
+            .transpose()
+    }
+
+    pub fn put<'txn>(
+        &self,
+        txn: &Txn<'txn, RW>,
+        k: impl AsRef<K::Item>,
+        v: impl AsRef<V::Item>,
+    ) -> anyhow::Result<()> {
+        let k_bytes = self
+            .k_wrapper
+            .dump_to_db_bytes(k.as_ref())
+            .context("whilst converting key to bytes")?;
+        let v_bytes = self
+            .v_wrapper
+            .dump_to_db_bytes(v.as_ref())
+            .context("whilst converting value to bytes")?;
+
+        self.raw.put(txn, k_bytes, v_bytes)?;
+        Ok(())
+    }
+}
diff --git a/src/environment.rs b/src/environment.rs
index 12de5b6..a0778e3 100644
--- a/src/environment.rs
+++ b/src/environment.rs
@@ -76,7 +76,7 @@ impl Env {
     pub fn rw_txn<R>(&self, func: impl FnOnce(&Txn<RW>) -> anyhow::Result<R>) -> anyhow::Result<R> {
         let txn = self.begin_rw_txn()?;
         let result = func(&txn)?;
-        txn.commit()?;
+        txn.commit().context("Failed to commit")?;
         Ok(result)
     }
 }
diff --git a/src/wrapper.rs b/src/wrapper.rs
index 23d95af..ff7846d 100644
--- a/src/wrapper.rs
+++ b/src/wrapper.rs
@@ -69,9 +69,10 @@ impl<T: ByteWrapper> ByteWrapper for Arc<T> {
     }
 }
 
-pub trait ZeroCopyByteWrapper {
+pub trait ZeroCopyByteWrapper: ToOwned {
     fn as_byte_slice(&self) -> &[u8];
     fn from_byte_slice(bytes: &[u8]) -> anyhow::Result<&Self>;
+    fn from_owned_bytes(bytes: Vec<u8>) -> anyhow::Result<Self::Owned>;
 }
 
 impl ZeroCopyByteWrapper for [u8] {
@@ -82,6 +83,10 @@ impl ZeroCopyByteWrapper for [u8] {
     fn from_byte_slice(bytes: &[u8]) -> anyhow::Result<&Self> {
         Ok(bytes)
     }
+
+    fn from_owned_bytes(bytes: Vec<u8>) -> anyhow::Result<Self::Owned> {
+        Ok(bytes)
+    }
 }
 
 impl ZeroCopyByteWrapper for str {
@@ -92,4 +97,8 @@ impl ZeroCopyByteWrapper for str {
     fn from_byte_slice(bytes: &[u8]) -> anyhow::Result<&Self> {
         Ok(std::str::from_utf8(bytes)?)
     }
+
+    fn from_owned_bytes(bytes: Vec<u8>) -> anyhow::Result<Self::Owned> {
+        Ok(String::from_utf8(bytes)?)
+    }
 }