back to where we were, but proper support for column families

cozo-rocks/include/cozorocks.h

@@ -13,18 +13,31 @@
 #include "rocksdb/utilities/transaction.h"
 #include "rocksdb/utilities/transaction_db.h"
 
-struct Status;
-
-namespace RDB = ROCKSDB_NAMESPACE;
-
-typedef RDB::Status::Code StatusCode;
-typedef RDB::Status::SubCode StatusSubCode;
-typedef RDB::Status::Severity StatusSeverity;
+using namespace ROCKSDB_NAMESPACE;
+using std::unique_ptr;
+using std::shared_ptr;
+using std::make_unique;
+using std::make_shared;
+using std::string;
+using std::vector;
+using std::unordered_map;
+
+struct BridgeStatus;
+
+typedef Status::Code StatusCode;
+typedef Status::SubCode StatusSubCode;
+typedef Status::Severity StatusSeverity;
+
+inline Slice convert_slice(rust::Slice<const uint8_t> d) {
+    return Slice(reinterpret_cast<const char *>(d.data()), d.size());
+}
 
-std::unique_ptr <RDB::DB> new_db();
+inline rust::Slice<const uint8_t> convert_slice_back(const Slice &s) {
+    return rust::Slice(reinterpret_cast<const std::uint8_t *>(s.data()), s.size());
+}
 
 struct ReadOptionsBridge {
-    mutable RDB::ReadOptions inner;
+    mutable ReadOptions inner;
 
     inline void do_set_verify_checksums(bool v) const {
         inner.verify_checksums = v;
@@ -36,7 +49,7 @@ struct ReadOptionsBridge {
 };
 
 struct WriteOptionsBridge {
-    mutable RDB::WriteOptions inner;
+    mutable WriteOptions inner;
 
 public:
     inline void do_set_disable_wal(bool v) const {
@@ -46,12 +59,10 @@ public:
 
 typedef rust::Fn<std::int8_t(rust::Slice<const std::uint8_t>, rust::Slice<const std::uint8_t>)> RustComparatorFn;
 
-class RustComparator : public RDB::Comparator {
+class RustComparator : public Comparator {
 public:
     inline int Compare(const rocksdb::Slice &a, const rocksdb::Slice &b) const {
-        auto ra = rust::Slice(reinterpret_cast<const std::uint8_t *>(a.data()), a.size());
-        auto rb = rust::Slice(reinterpret_cast<const std::uint8_t *>(b.data()), b.size());
-        return int(rust_compare(ra, rb));
+        return int(rust_compare(convert_slice_back(a), convert_slice_back(b)));
     }
 
     const char *Name() const {
@@ -75,10 +86,9 @@ public:
 };
 
 struct OptionsBridge {
-    mutable RDB::Options inner;
+    mutable Options inner;
     mutable RustComparator cmp_obj;
 
-public:
     inline void do_prepare_for_bulk_load() const {
         inner.PrepareForBulkLoad();
     }
@@ -117,40 +127,42 @@ inline std::unique_ptr <OptionsBridge> new_options() {
 
 
 struct PinnableSliceBridge {
-    RDB::PinnableSlice inner;
+    PinnableSlice inner;
 
     inline rust::Slice<const std::uint8_t> as_bytes() const {
-        return rust::Slice(reinterpret_cast<const std::uint8_t *>(inner.data()), inner.size());
+        return convert_slice_back(inner);
     }
 };
 
 struct SliceBridge {
-    RDB::Slice inner;
-    SliceBridge(RDB::Slice&& s) : inner(s) {}
+    Slice inner;
+
+    SliceBridge(Slice &&s) : inner(s) {}
 
     inline rust::Slice<const std::uint8_t> as_bytes() const {
-        return rust::Slice(reinterpret_cast<const std::uint8_t *>(inner.data()), inner.size());
+        return convert_slice_back(inner);
     }
 };
 
-void write_status_impl(Status &status, StatusCode code, StatusSubCode subcode, StatusSeverity severity);
+void write_status_impl(BridgeStatus &status, StatusCode code, StatusSubCode subcode, StatusSeverity severity, int bridge_code);
 
-inline void write_status(RDB::Status &&rstatus, Status &status) {
+inline void write_status(Status &&rstatus, BridgeStatus &status) {
     if (rstatus.code() != StatusCode::kOk || rstatus.subcode() != StatusSubCode::kNoSpace ||
         rstatus.severity() != StatusSeverity::kNoError) {
-        write_status_impl(status, rstatus.code(), rstatus.subcode(), rstatus.severity());
+        write_status_impl(status, rstatus.code(), rstatus.subcode(), rstatus.severity(), 0);
     }
 }
 
-struct WriteBatchBridge {
-    mutable RDB::WriteBatch inner;
-    std::vector<RDB::ColumnFamilyHandle *> *handles;
-};
-
+//
+//struct WriteBatchBridge {
+//    mutable WriteBatch inner;
+//    std::vector<ColumnFamilyHandle *> *handles;
+//};
+//
 struct IteratorBridge {
-    mutable std::unique_ptr <RDB::Iterator> inner;
+    mutable std::unique_ptr <Iterator> inner;
 
-    IteratorBridge(RDB::Iterator *it) : inner(it) {}
+    IteratorBridge(Iterator *it) : inner(it) {}
 
     inline void seek_to_first() const {
         inner->SeekToFirst();
@@ -169,121 +181,151 @@ struct IteratorBridge {
     }
 
     inline void do_seek(rust::Slice<const uint8_t> key) const {
-        auto k = RDB::Slice(reinterpret_cast<const char *>(key.data()), key.size());
+        auto k = Slice(reinterpret_cast<const char *>(key.data()), key.size());
         inner->Seek(k);
     }
 
     inline void do_seek_for_prev(rust::Slice<const uint8_t> key) const {
-        auto k = RDB::Slice(reinterpret_cast<const char *>(key.data()), key.size());
+        auto k = Slice(reinterpret_cast<const char *>(key.data()), key.size());
         inner->SeekForPrev(k);
     }
 
-    inline std::unique_ptr<SliceBridge> key() const {
+    inline std::unique_ptr <SliceBridge> key() const {
         return std::make_unique<SliceBridge>(inner->key());
     }
 
-    inline std::unique_ptr<SliceBridge> value() const {
+    inline std::unique_ptr <SliceBridge> value() const {
         return std::make_unique<SliceBridge>(inner->value());
     }
 
-    Status status() const;
+    BridgeStatus status() const;
 };
 
 struct DBBridge {
-    mutable std::unique_ptr <RDB::DB> inner;
+    mutable unique_ptr <DB> db;
+    mutable unordered_map <string, shared_ptr<ColumnFamilyHandle>> handles;
 
-    mutable std::vector<RDB::ColumnFamilyHandle *> handles;
+    DBBridge(DB *db_,
+             unordered_map <string, shared_ptr<ColumnFamilyHandle>> &&handles_) : db(db_), handles(handles_) {}
 
-    DBBridge(RDB::DB *inner_) : inner(inner_) {}
 
-    inline std::unique_ptr <std::vector<std::string>> cf_names() const {
-        auto ret = std::make_unique < std::vector < std::string >> ();
-        for (auto h: handles) {
-            ret->push_back(h->GetName());
-        }
-        return ret;
+    inline shared_ptr <ColumnFamilyHandle> get_cf_handle_raw(const string &name) const {
+        try {
+            return handles.at(name);
+        } catch (const std::out_of_range &) {
+            return shared_ptr<ColumnFamilyHandle>(nullptr);
         }
-
-    inline std::unique_ptr <WriteBatchBridge> write_batch() const {
-        auto wb = std::make_unique<WriteBatchBridge>();
-        wb->handles = &handles;
-        return wb;
     }
 
-    inline void put(
+//
+//    inline std::unique_ptr <WriteBatchBridge> write_batch() const {
+//        auto wb = std::make_unique<WriteBatchBridge>();
+//        wb->handles = &handles;
+//        return wb;
+//    }
+//
+    inline void put_raw(
             const WriteOptionsBridge &options,
-            std::size_t cf_id,
+            const ColumnFamilyHandle &cf,
             rust::Slice<const uint8_t> key,
             rust::Slice<const uint8_t> val,
-            Status &status
+            BridgeStatus &status
     ) const {
         write_status(
-                inner->Put(options.inner,
-                           handles[cf_id],
-                           RDB::Slice(reinterpret_cast<const char *>(key.data()), key.size()),
-                           RDB::Slice(reinterpret_cast<const char *>(val.data()), val.size())),
+                db->Put(options.inner,
+                        const_cast<ColumnFamilyHandle *>(&cf),
+                        convert_slice(key),
+                        convert_slice(val)),
                 status
         );
     }
 
-    inline std::unique_ptr <PinnableSliceBridge> get(
+    inline std::unique_ptr <PinnableSliceBridge> get_raw(
             const ReadOptionsBridge &options,
-            std::size_t cf_id,
+            const ColumnFamilyHandle &cf,
             rust::Slice<const uint8_t> key,
-            Status &status
+            BridgeStatus &status
     ) const {
         auto pinnable_val = std::make_unique<PinnableSliceBridge>();
         write_status(
-                inner->Get(options.inner,
-                           handles[cf_id],
-                           RDB::Slice(reinterpret_cast<const char *>(key.data()), key.size()),
+                db->Get(options.inner,
+                        const_cast<ColumnFamilyHandle *>(&cf),
+                        convert_slice(key),
                         &pinnable_val->inner),
                 status
         );
         return pinnable_val;
     }
 
-    inline std::unique_ptr <IteratorBridge> iterator(const ReadOptionsBridge &options, std::size_t cf_id) const {
-        return std::make_unique<IteratorBridge>(inner->NewIterator(options.inner, handles[cf_id]));
+    inline std::unique_ptr <IteratorBridge> iterator_raw(
+            const ReadOptionsBridge &options,
+            const ColumnFamilyHandle &cf) const {
+        return std::make_unique<IteratorBridge>(db->NewIterator(options.inner, const_cast<ColumnFamilyHandle *>(&cf)));
+    }
+
+    inline void create_column_family_raw(const OptionsBridge &options, const string &name, BridgeStatus &status) const {
+        if (handles.find(name) != handles.end()) {
+            write_status_impl(status, StatusCode::kMaxCode, StatusSubCode::kMaxSubCode, StatusSeverity::kSoftError, 2);
+            return;
+        }
+        ColumnFamilyHandle *handle;
+        auto s = db->CreateColumnFamily(options.inner, name, &handle);
+        write_status(std::move(s), status);
+        handles[name] = shared_ptr<ColumnFamilyHandle>(handle);
+    }
+
+    inline void drop_column_family_raw(const string &name, BridgeStatus &status) const {
+        auto cf_it = handles.find(name);
+        if (cf_it != handles.end()) {
+            auto s = db->DropColumnFamily(cf_it->second.get());
+            handles.erase(cf_it);
+            write_status(std::move(s), status);
+        } else {
+            write_status_impl(status, StatusCode::kMaxCode, StatusSubCode::kMaxSubCode, StatusSeverity::kSoftError, 3);
+        }
+        // When should we call DestroyColumnFamilyHandle?
+    }
+
+    inline unique_ptr<vector<string>> get_column_family_names_raw() const {
+        auto ret = make_unique<vector<string>>();
+        for (auto entry : handles) {
+            ret->push_back(entry.first);
+        }
+        return ret;
     }
 };
 
-inline std::unique_ptr <std::vector<std::string>> list_column_families(const OptionsBridge &options,
-                                                                       const rust::Slice<const uint8_t> path) {
-    auto column_families = std::make_unique < std::vector < std::string >> ();
-    RDB::DB::ListColumnFamilies(options.inner,
-                                std::string(reinterpret_cast<const char *>(path.data()), path.size()),
-                                &*column_families);
-    return column_families;
-}
 
 inline std::unique_ptr <DBBridge>
-open_db(const OptionsBridge &options, const rust::Slice<const uint8_t> path, Status &status) {
-    auto old_column_families = std::vector<std::string>();
-    RDB::DB::ListColumnFamilies(options.inner,
-                                std::string(reinterpret_cast<const char *>(path.data()), path.size()),
-                                &old_column_families);
-    if (old_column_families.empty()) {
-        old_column_families.push_back(RDB::kDefaultColumnFamilyName);
+open_db_raw(const OptionsBridge &options,
+            const string &path,
+            BridgeStatus &status) {
+    auto cf_names = std::vector<std::string>();
+    DB::ListColumnFamilies(options.inner, path, &cf_names);
+    if (cf_names.empty()) {
+        cf_names.push_back(kDefaultColumnFamilyName);
     }
 
-    std::vector <RDB::ColumnFamilyDescriptor> column_families;
+    std::vector <ColumnFamilyDescriptor> column_families;
 
-    for (auto el: old_column_families) {
-        column_families.push_back(RDB::ColumnFamilyDescriptor(
+    for (auto el: cf_names) {
+        column_families.push_back(ColumnFamilyDescriptor(
                 el, options.inner));
     }
 
-    std::vector < RDB::ColumnFamilyHandle * > handles;
+    std::vector < ColumnFamilyHandle * > handles;
 
-    RDB::DB *db_ptr;
-    RDB::Status s = RDB::DB::Open(options.inner,
-                                  std::string(reinterpret_cast<const char *>(path.data()), path.size()),
-                                  column_families,
-                                  &handles,
-                                  &db_ptr);
+    DB *db_ptr;
+    Status s = DB::Open(options.inner, path, column_families, &handles, &db_ptr);
+
+    auto ok = s.ok();
     write_status(std::move(s), status);
-    auto ret = std::unique_ptr<DBBridge>(new DBBridge(db_ptr));
-    ret->handles = std::move(handles);
-    return ret;
+    unordered_map <string, shared_ptr<ColumnFamilyHandle>> handle_map;
+    if (ok) {
+        assert(handles.size() == cf_names.size());
+        for (size_t i = 0; i < handles.size(); ++i) {
+            handle_map[cf_names[i]] = shared_ptr<ColumnFamilyHandle>(handles[i]);
+        }
+    }
+    return std::make_unique<DBBridge>(db_ptr, std::move(handle_map));
 }

cozo-rocks/src/cozorocks.cc

@@ -1,18 +1,19 @@
+////
+//// Created by Ziyang Hu on 2022/4/13.
+////
 //
-// Created by Ziyang Hu on 2022/4/13.
-//
-
 #include "../include/cozorocks.h"
 #include "cozo-rocks/src/lib.rs.h"
 
-void write_status_impl(Status &status, StatusCode code, StatusSubCode subcode, StatusSeverity severity) {
+void write_status_impl(BridgeStatus &status, StatusCode code, StatusSubCode subcode, StatusSeverity severity, int bridge_code) {
     status.code = code;
     status.subcode = subcode;
     status.severity = severity;
+    status.bridge_code = static_cast<StatusBridgeCode>(bridge_code);
 }
 
-Status IteratorBridge::status() const {
-    Status s;
+BridgeStatus IteratorBridge::status() const {
+    BridgeStatus s;
     write_status(inner->status(), s);
     return s;
 }

cozo-rocks/src/lib.rs

@@ -1,10 +1,19 @@
 #[cxx::bridge]
 mod ffi {
     #[derive(Copy, Clone, Debug, Eq, PartialEq)]
-    pub struct Status {
+    pub enum StatusBridgeCode {
+        OK = 0,
+        LOCK_ERROR = 1,
+        EXISTING_ERROR = 2,
+        NOT_FOUND_ERROR = 3,
+    }
+
+    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
+    pub struct BridgeStatus {
         pub code: StatusCode,
         pub subcode: StatusSubCode,
         pub severity: StatusSeverity,
+        pub bridge_code: StatusBridgeCode,
     }
 
     #[derive(Copy, Clone, Debug, Eq, PartialEq)]
@@ -66,60 +75,66 @@ mod ffi {
         type StatusSeverity;
 
         type PinnableSliceBridge;
-        pub fn as_bytes(self: &PinnableSliceBridge) -> &[u8];
+        fn as_bytes(self: &PinnableSliceBridge) -> &[u8];
 
         type SliceBridge;
-        pub fn as_bytes(self: &SliceBridge) -> &[u8];
+        fn as_bytes(self: &SliceBridge) -> &[u8];
 
         type ReadOptionsBridge;
         fn new_read_options() -> UniquePtr<ReadOptionsBridge>;
-        pub fn do_set_verify_checksums(self: &ReadOptionsBridge, v: bool);
-        pub fn do_set_total_order_seek(self: &ReadOptionsBridge, v: bool);
+        fn do_set_verify_checksums(self: &ReadOptionsBridge, v: bool);
+        fn do_set_total_order_seek(self: &ReadOptionsBridge, v: bool);
 
         type WriteOptionsBridge;
         fn new_write_options() -> UniquePtr<WriteOptionsBridge>;
-        pub fn do_set_disable_wal(self: &WriteOptionsBridge, v: bool);
+        fn do_set_disable_wal(self: &WriteOptionsBridge, v: bool);
 
         type OptionsBridge;
         fn new_options() -> UniquePtr<OptionsBridge>;
-        pub fn do_prepare_for_bulk_load(self: &OptionsBridge);
-        pub fn do_increase_parallelism(self: &OptionsBridge);
-        pub fn do_optimize_level_style_compaction(self: &OptionsBridge);
-        pub fn do_set_create_if_missing(self: &OptionsBridge, v: bool);
-        pub fn do_set_comparator(self: &OptionsBridge, name: &str, compare: fn(&[u8], &[u8]) -> i8);
+        fn do_prepare_for_bulk_load(self: &OptionsBridge);
+        fn do_increase_parallelism(self: &OptionsBridge);
+        fn do_optimize_level_style_compaction(self: &OptionsBridge);
+        fn do_set_create_if_missing(self: &OptionsBridge, v: bool);
+        fn do_set_comparator(self: &OptionsBridge, name: &str, compare: fn(&[u8], &[u8]) -> i8);
 
+        pub type ColumnFamilyHandle;
         type DBBridge;
-        fn list_column_families(options: &OptionsBridge, path: &[u8]) -> UniquePtr<CxxVector<CxxString>>;
-        fn open_db(options: &OptionsBridge, path: &[u8], status: &mut Status) -> UniquePtr<DBBridge>;
-        fn cf_names(self: &DBBridge) -> UniquePtr<CxxVector<CxxString>>;
-        fn put(self: &DBBridge, options: &WriteOptionsBridge, cf_id: usize, key: &[u8], val: &[u8], status: &mut Status);
-        fn get(self: &DBBridge, options: &ReadOptionsBridge, cf_id: usize, key: &[u8], status: &mut Status) -> UniquePtr<PinnableSliceBridge>;
-        fn write_batch(self: &DBBridge) -> UniquePtr<WriteBatchBridge>;
-        fn iterator(self: &DBBridge, options: &ReadOptionsBridge, cf_id: usize) -> UniquePtr<IteratorBridge>;
-
-        type WriteBatchBridge;
-
+        fn open_db_raw(options: &OptionsBridge, path: &CxxString, status: &mut BridgeStatus) -> UniquePtr<DBBridge>;
+        fn get_cf_handle_raw(self: &DBBridge, name: &CxxString) -> SharedPtr<ColumnFamilyHandle>;
+        fn put_raw(self: &DBBridge, options: &WriteOptionsBridge, cf: &ColumnFamilyHandle, key: &[u8], val: &[u8], status: &mut BridgeStatus);
+        fn get_raw(self: &DBBridge, options: &ReadOptionsBridge, cf: &ColumnFamilyHandle, key: &[u8], status: &mut BridgeStatus) -> UniquePtr<PinnableSliceBridge>;
+        fn iterator_raw(self: &DBBridge, options: &ReadOptionsBridge, cf: &ColumnFamilyHandle) -> UniquePtr<IteratorBridge>;
+        fn create_column_family_raw(self: &DBBridge, options: &OptionsBridge, name: &CxxString, status: &mut BridgeStatus);
+        fn drop_column_family_raw(self: &DBBridge, name: &CxxString, status: &mut BridgeStatus);
+        fn get_column_family_names_raw(self: &DBBridge) -> UniquePtr<CxxVector<CxxString>>;
+
+        //         fn write_batch(self: &DBBridge) -> UniquePtr<WriteBatchBridge>;
+//
+//         type WriteBatchBridge;
+//
         type IteratorBridge;
-        pub fn seek_to_first(self: &IteratorBridge);
-        pub fn seek_to_last(self: &IteratorBridge);
-        pub fn next(self: &IteratorBridge);
-        pub fn is_valid(self: &IteratorBridge) -> bool;
+        fn seek_to_first(self: &IteratorBridge);
+        fn seek_to_last(self: &IteratorBridge);
+        fn next(self: &IteratorBridge);
+        fn is_valid(self: &IteratorBridge) -> bool;
         fn do_seek(self: &IteratorBridge, key: &[u8]);
         fn do_seek_for_prev(self: &IteratorBridge, key: &[u8]);
-        pub fn key(self: &IteratorBridge) -> UniquePtr<SliceBridge>;
-        pub fn value(self: &IteratorBridge) -> UniquePtr<SliceBridge>;
-        pub fn status(self: &IteratorBridge) -> Status;
+        fn key(self: &IteratorBridge) -> UniquePtr<SliceBridge>;
+        fn value(self: &IteratorBridge) -> UniquePtr<SliceBridge>;
+        fn status(self: &IteratorBridge) -> BridgeStatus;
     }
 }
 
-
-pub use ffi::{Status, StatusCode, StatusSubCode, StatusSeverity};
-use std::collections::BTreeMap;
-use cxx::UniquePtr;
+use std::path::Path;
+use cxx::{UniquePtr, SharedPtr, let_cxx_string};
 use ffi::*;
 
+type Result<T> = std::result::Result<T, BridgeStatus>;
+
 pub type Options = UniquePtr<OptionsBridge>;
 
+type ColumnFamilyHandle = SharedPtr<ffi::ColumnFamilyHandle>;
+
 pub trait OptionsTrait {
     fn prepare_for_bulk_load(self) -> Self;
     fn increase_parallelism(self) -> Self;
@@ -228,12 +243,12 @@ impl AsRef<[u8]> for Slice {
 
 pub type Iterator = UniquePtr<IteratorBridge>;
 
-pub trait IteratorTrait {
+pub trait IteratorImpl {
     fn seek(&self, key: impl AsRef<[u8]>);
     fn seek_for_prev(&self, key: impl AsRef<[u8]>);
 }
 
-impl IteratorTrait for Iterator {
+impl IteratorImpl for IteratorBridge {
     fn seek(&self, key: impl AsRef<[u8]>) {
         self.do_seek(key.as_ref());
     }
@@ -242,14 +257,6 @@ impl IteratorTrait for Iterator {
     }
 }
 
-pub struct DB {
-    bridge: UniquePtr<DBBridge>,
-    pub options: Options,
-    pub default_read_options: ReadOptions,
-    pub default_write_options: WriteOptions,
-    pub column_families: BTreeMap<String, usize>,
-}
-
 fn get_path_bytes(path: &std::path::Path) -> &[u8] {
     #[cfg(unix)]
     {
@@ -261,40 +268,106 @@ fn get_path_bytes(path: &std::path::Path) -> &[u8] {
     { path.to_string_lossy().to_string().as_bytes() }
 }
 
-impl DB {
+//
+//     #[inline]
+//     pub fn write_batch(&self) -> UniquePtr<WriteBatchBridge> {
+//         self.bridge.write_batch()
+//     }
+//
+//     // #[inline]
+//     // pub fn get_column_family_id(&self, name: impl AsRef<str>) -> Result<Option<usize>, Status> {
+//     //     let handles = self.cf_map.read()
+//     //         .map_err(|_| Status::bridge(StatusBridgeCode::LOCK_ERROR))?;
+//     //     Ok(handles.get(name.as_ref()).copied())
+//     // }
+//
+//     // #[inline]
+//     // pub fn create_column_family(&self, name: impl AsRef<str>) -> Result<(), Status> {
+//     //     let mut s = Status::default();
+//     //     let mut cf_map = self.cf_map.write()
+//     //         .map_err(|_| Status::bridge(StatusBridgeCode::LOCK_ERROR))?;
+//     //     let mut cfs = self.cfs.write()
+//     //         .map_err(|_| Status::bridge(StatusBridgeCode::LOCK_ERROR))?;
+//     //     let v = self.bridge.create_column_family(&self.options, name.as_ref(), &mut s);
+//     //     if v > 0 {
+//     //         assert_eq!(v as usize, cfs.len());
+//     //         cf_map.insert(name.as_ref().to_string(), v as usize);
+//     //         cfs.push(name.as_ref().to_string());
+//     //         Ok(())
+//     //     } else {
+//     //         Err(s)
+//     //     }
+//     // }
+//
+//     // #[inline]
+//     // pub fn drop_column_family(&self, _name: impl AsRef<str>) -> Result<(), Status> {
+//     //     unimplemented!()
+//     // }
+//
+//     pub fn destroy_data(self) -> Result<(), Status> {
+//         unimplemented!()
+//     }
+// }
+//
+impl Default for BridgeStatus {
     #[inline]
-    pub fn list_column_families(options: &Options, path: impl AsRef<std::path::Path>) -> Vec<String> {
-        let results = list_column_families(&options, get_path_bytes(path.as_ref()));
-        results.iter().map(|s| s.to_string_lossy().into_owned()).collect()
+    fn default() -> Self {
+        Self {
+            code: StatusCode::kOk,
+            subcode: StatusSubCode::kNone,
+            severity: StatusSeverity::kNoError,
+            bridge_code: StatusBridgeCode::OK,
+        }
     }
+}
 
+impl BridgeStatus {
     #[inline]
-    pub fn open(options: Options, path: impl AsRef<std::path::Path>) -> Result<Self, Status> {
-        let mut status = Status::default();
-        let bridge = open_db(
-            &options,
-            get_path_bytes(path.as_ref()),
-            &mut status,
-        );
+    fn bridge(c: StatusBridgeCode) -> Self {
+        Self {
+            code: StatusCode::kMaxCode,
+            subcode: StatusSubCode::kMaxSubCode,
+            severity: StatusSeverity::kMaxSeverity,
+            bridge_code: c,
+        }
+    }
+}
 
-        if status.code == StatusCode::kOk {
-            let column_families = bridge.cf_names().iter().enumerate().map(|(i, v)| (v.to_string_lossy().into_owned(), i)).collect();
-            Ok(Self {
-                bridge,
-                default_read_options: ReadOptions::default(),
-                default_write_options: WriteOptions::default(),
-                options,
-                column_families,
-            })
-        } else {
-            Err(status)
+pub trait DBRead {
+    fn get(&self, key: impl AsRef<[u8]>, cf: &ColumnFamilyHandle, options: Option<&ReadOptions>)
+           -> Result<Option<PinnableSlice>>;
+}
+
+pub trait DBWrite {
+    fn put(&self, key: impl AsRef<[u8]>, val: impl AsRef<[u8]>, cf: &ColumnFamilyHandle, options: Option<&WriteOptions>)
+           -> Result<BridgeStatus>;
+}
+
+pub struct DB {
+    inner: UniquePtr<DBBridge>,
+    pub options: Options,
+    pub default_read_options: ReadOptions,
+    pub default_write_options: WriteOptions,
+}
+
+impl DBRead for DB {
+    #[inline]
+    fn get(&self, key: impl AsRef<[u8]>, cf: &ColumnFamilyHandle, options: Option<&ReadOptions>) -> Result<Option<PinnableSlice>> {
+        let mut status = BridgeStatus::default();
+        let slice = self.inner.get_raw(options.unwrap_or(&self.default_read_options), cf, key.as_ref(), &mut status);
+        match status.code {
+            StatusCode::kOk => Ok(Some(PinnableSlice(slice))),
+            StatusCode::kNotFound => Ok(None),
+            _ => Err(status)
         }
     }
+}
 
+impl DBWrite for DB {
     #[inline]
-    pub fn put(&self, key: impl AsRef<[u8]>, val: impl AsRef<[u8]>, cf: usize, options: Option<&WriteOptions>) -> Result<Status, Status> {
-        let mut status = Status::default();
-        self.bridge.put(options.unwrap_or(&self.default_write_options), cf,
+    fn put(&self, key: impl AsRef<[u8]>, val: impl AsRef<[u8]>, cf: &ColumnFamilyHandle, options: Option<&WriteOptions>) -> Result<BridgeStatus> {
+        let mut status = BridgeStatus::default();
+        self.inner.put_raw(options.unwrap_or(&self.default_write_options), cf,
                            key.as_ref(), val.as_ref(),
                            &mut status);
         if status.code == StatusCode::kOk {
@@ -303,52 +376,83 @@ impl DB {
             Err(status)
         }
     }
+}
 
-    #[inline]
-    pub fn get(&self, key: impl AsRef<[u8]>, cf: usize, options: Option<&ReadOptions>) -> Result<Option<PinnableSlice>, Status> {
-        let mut status = Status::default();
-        let slice = self.bridge.get(
-            options.unwrap_or(&self.default_read_options), cf,
-            key.as_ref(), &mut status);
-        match status.code {
-            StatusCode::kOk => Ok(Some(PinnableSlice(slice))),
-            StatusCode::kNotFound => Ok(None),
-            _ => Err(status)
+
+pub trait DBImpl {
+    fn open(options: Options, path: &Path) -> Result<DB>;
+    fn get_cf_handle(&self, name: impl AsRef<str>) -> Result<ColumnFamilyHandle>;
+    fn iterator(&self, cf: &ColumnFamilyHandle, options: Option<&ReadOptions>) -> Iterator;
+    fn create_column_family(&self, name: impl AsRef<str>) -> Result<()>;
+    fn drop_column_family(&self, name: impl AsRef<str>) -> Result<()>;
+    fn get_column_family_names(&self) -> Vec<String>;
+}
+
+impl DBImpl for DB {
+    fn open(options: Options, path: &Path) -> Result<DB> {
+        let_cxx_string!(path = get_path_bytes(path));
+        let mut status = BridgeStatus::default();
+        let bridge = open_db_raw(
+            &options,
+            &path,
+            &mut status,
+        );
+
+        if status.code == StatusCode::kOk {
+            Ok(DB {
+                inner: bridge,
+                options,
+                default_read_options: ReadOptions::default(),
+                default_write_options: WriteOptions::default(),
+            })
+        } else {
+            Err(status)
         }
     }
 
-    #[inline]
-    pub fn iterator(&self, cf: usize, options: Option<&ReadOptions>) -> Iterator {
-        self.bridge.iterator(options.unwrap_or(&self.default_read_options), cf)
+    fn get_cf_handle(&self, name: impl AsRef<str>) -> Result<ColumnFamilyHandle> {
+        let_cxx_string!(name = name.as_ref());
+        let ret = self.inner.get_cf_handle_raw(&name);
+        if ret.is_null() {
+            Err(BridgeStatus {
+                code: StatusCode::kMaxCode,
+                subcode: StatusSubCode::kMaxSubCode,
+                severity: StatusSeverity::kSoftError,
+                bridge_code: StatusBridgeCode::NOT_FOUND_ERROR
+            })
+        } else {
+            Ok(ret)
         }
-
-    #[inline]
-    pub fn write_batch(&self) -> UniquePtr<WriteBatchBridge> {
-        self.bridge.write_batch()
     }
 
     #[inline]
-    pub fn create_column_family(&self, _name: impl AsRef<str>) -> Result<usize, Status> {
-        unimplemented!()
+    fn iterator(&self, cf: &ColumnFamilyHandle, options: Option<&ReadOptions>) -> Iterator {
+        self.inner.iterator_raw(options.unwrap_or(&self.default_read_options), cf)
     }
 
-    #[inline]
-    pub fn drop_column_family(&self, _name: impl AsRef<str>) -> Result<(), Status> {
-        unimplemented!()
+    fn create_column_family(&self, name: impl AsRef<str>) -> Result<()> {
+        let_cxx_string!(name = name.as_ref());
+        let mut status = BridgeStatus::default();
+        self.inner.create_column_family_raw(&self.options, &name, &mut status);
+        if status.code == StatusCode::kOk {
+            Ok(())
+        } else {
+            Err(status)
         }
-
-    pub fn destroy_data(self) -> Result<(), Status> {
-        unimplemented!()
     }
-}
 
-impl Default for Status {
-    #[inline]
-    fn default() -> Self {
-        Self {
-            code: StatusCode::kOk,
-            subcode: StatusSubCode::kNone,
-            severity: StatusSeverity::kNoError,
+    fn drop_column_family(&self, name: impl AsRef<str>) -> Result<()> {
+        let_cxx_string!(name = name.as_ref());
+        let mut status = BridgeStatus::default();
+        self.inner.drop_column_family_raw(&name, &mut status);
+        if status.code == StatusCode::kOk {
+            Ok(())
+        } else {
+            Err(status)
         }
     }
+
+    fn get_column_family_names(&self) -> Vec<String> {
+        self.inner.get_column_family_names_raw().iter().map(|v| v.to_string_lossy().to_string()).collect()
+    }
 }

src/storage.rs

@@ -58,6 +58,7 @@ impl Storage {
 #[cfg(test)]
 mod tests {
     use std::str::from_utf8;
+    use cozo_rocks::DBImpl;
     use crate::value::{ByteArrayBuilder, cozo_comparator_v1, Value};
 
     #[test]
@@ -70,7 +71,7 @@ mod tests {
             .set_comparator("cozo_comparator_v1", cozo_comparator_v1);
 
         let db = DB::open(options,
-                          "xxyyzz.db").unwrap();
+                          "xxyyzz.db".as_ref()).unwrap();
 
         let mut builder = ByteArrayBuilder::default();
         builder.build_value(&Value::RefString("A key"));
@@ -79,20 +80,22 @@ mod tests {
         let mut builder = ByteArrayBuilder::default();
         builder.build_value(&Value::RefString("Another key"));
         let key2 = builder;
+        let cf = db.get_cf_handle("default").unwrap();
+        println!("{:?}", db.get_column_family_names());
 
-        let val = db.get(&key, 0, None).unwrap();
+        let val = db.get(&key, &cf, None).unwrap();
         println!("before anything {}", val.is_none());
 
-        db.put(&key, "A motherfucking value!!! 👋👋👋", 0, None).unwrap();
-        db.put(&key2, "Another motherfucking value!!! 👋👋👋", 0, None).unwrap();
+        db.put(&key, "A motherfucking value!!! 👋👋👋", &cf, None).unwrap();
+        db.put(&key2, "Another motherfucking value!!! 👋👋👋", &cf, None).unwrap();
         // db.put("Yes man", "A motherfucking value!!! 👋👋👋", None).unwrap();
-        let val = db.get(&key, 0, None).unwrap().unwrap();
+        let val = db.get(&key, &cf, None).unwrap().unwrap();
         println!("1 {}", from_utf8(val.as_ref()).unwrap());
-        let val = db.get(&key2, 0, None).unwrap().unwrap();
+        let val = db.get(&key2, &cf, None).unwrap().unwrap();
         // let val = val.as_bytes();
         println!("2 {}", from_utf8(val.as_ref()).unwrap());
-        let val = db.get(&key, 0, None).unwrap().unwrap();
+        let val = db.get(&key, &cf, None).unwrap().unwrap();
         println!("3 {}", from_utf8(val.as_ref()).unwrap());
-        println!("4 {}", from_utf8(db.get(&key, 0, None).unwrap().unwrap().as_ref()).unwrap());
+        println!("4 {}", from_utf8(db.get(&key, &cf, None).unwrap().unwrap().as_ref()).unwrap());
     }
 }