From a0f3ad3aa4056e0b0e2706bcd17df9fc1da98781 Mon Sep 17 00:00:00 2001
From: David Greenberg <dsg123456789@gmail.com>
Date: Sat, 11 Jul 2015 21:22:40 -0400
Subject: [PATCH] Add support for writebatch

---
 src/ffi.rs     |  4 +++
 src/lib.rs     |  1 +
 src/rocksdb.rs | 88 +++++++++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 88 insertions(+), 5 deletions(-)

diff --git a/src/ffi.rs b/src/ffi.rs
index 5a15cfd..1fc2dfe 100644
--- a/src/ffi.rs
+++ b/src/ffi.rs
@@ -272,6 +272,10 @@ extern {
     pub fn rocksdb_iter_get_error(iter: RocksDBIterator,
                                   err: *const *const u8);
     // Write batch
+    pub fn rocksdb_write(db: RocksDBInstance,
+                         writeopts: RocksDBWriteOptions,
+                         batch : RocksDBWriteBatch,
+                         err: *mut i8);
     pub fn rocksdb_writebatch_create() -> RocksDBWriteBatch;
     pub fn rocksdb_writebatch_create_from(rep: *const u8,
                                           size: size_t) -> RocksDBWriteBatch;
diff --git a/src/lib.rs b/src/lib.rs
index e423b02..2d2fe95 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -34,6 +34,7 @@ pub use rocksdb::{
     RocksDB,
     RocksDBResult,
     RocksDBVector,
+    WriteBatch,
     Writable,
 };
 pub use rocksdb_options::{
diff --git a/src/rocksdb.rs b/src/rocksdb.rs
index a504db7..c0e4c49 100644
--- a/src/rocksdb.rs
+++ b/src/rocksdb.rs
@@ -128,11 +128,17 @@ impl RocksDB {
         }
     }
 
-    ////pub fn create_writebatch(self) -> WriteBatch {
-    ////    WriteBatch(unsafe {
-    ////        rocksdb_ffi::rocksdb_writebatch_create()
-    ////    })
-    ////}
+    pub fn write(&self, batch: WriteBatch) -> Result<(), &str> {
+        unsafe {
+            let writeopts = rocksdb_ffi::rocksdb_writeoptions_create();
+            let err = 0 as *mut i8;
+            rocksdb_ffi::rocksdb_write(self.inner, writeopts, batch.inner, err);
+            if !err.is_null() {
+                return Err(error_message(err));
+            }
+            return Ok(())
+        }
+    }
 
     pub fn get(&self, key: &[u8]) -> RocksDBResult<RocksDBVector, &str> {
         unsafe {
@@ -209,6 +215,52 @@ impl Writable for RocksDB {
     }
 }
 
+impl WriteBatch {
+    pub fn new() -> WriteBatch {
+        WriteBatch {
+            inner: unsafe {
+                       rocksdb_ffi::rocksdb_writebatch_create()
+                   }
+        }
+    }
+}
+
+impl Drop for WriteBatch {
+    fn drop(&mut self) {
+        unsafe {
+            rocksdb_ffi::rocksdb_writebatch_destroy(self.inner)
+        }
+    }
+}
+
+impl Writable for WriteBatch {
+    fn put(&self, key: &[u8], value: &[u8]) -> Result<(), &str> {
+        unsafe {
+            rocksdb_ffi::rocksdb_writebatch_put(self.inner, key.as_ptr(),
+                        key.len() as size_t, value.as_ptr(),
+                        value.len() as size_t);
+            return Ok(())
+        }
+    }
+
+    fn merge(&self, key: &[u8], value: &[u8]) -> Result<(), &str> {
+        unsafe {
+            rocksdb_ffi::rocksdb_writebatch_merge(self.inner, key.as_ptr(),
+                        key.len() as size_t, value.as_ptr(),
+                        value.len() as size_t);
+            return Ok(())
+        }
+    }
+
+    fn delete(&self, key: &[u8]) -> Result<(),&str> {
+        unsafe {
+            rocksdb_ffi::rocksdb_writebatch_delete(self.inner, key.as_ptr(),
+                        key.len() as size_t);
+            return Ok(())
+        }
+    }
+}
+
 pub struct RocksDBVector {
     base: Unique<u8>,
     len: usize,
@@ -329,3 +381,29 @@ fn external() {
     let opts = RocksDBOptions::new();
     assert!(RocksDB::destroy(opts, path).is_ok());
 }
+
+#[test]
+fn writebatch_works() {
+    let path = "_rust_rocksdb_writebacktest";
+    let db = RocksDB::open_default(path).unwrap();
+    { // test put
+        let batch = WriteBatch::new();
+        assert!(db.get(b"k1").is_none());
+        batch.put(b"k1", b"v1111");
+        assert!(db.get(b"k1").is_none());
+        let p = db.write(batch);
+        assert!(p.is_ok());
+        let r: RocksDBResult<RocksDBVector, &str> = db.get(b"k1");
+        assert!(r.unwrap().to_utf8().unwrap() == "v1111");
+    }
+    { // test delete
+        let batch = WriteBatch::new();
+        batch.delete(b"k1");
+        let p = db.write(batch);
+        assert!(p.is_ok());
+        assert!(db.get(b"k1").is_none());
+    }
+    db.close();
+    let opts = RocksDBOptions::new();
+    assert!(RocksDB::destroy(opts, path).is_ok());
+}