documentation

cozo-core/src/lib.rs

@@ -10,6 +10,19 @@
 //! * [Installation and first queries](https://github.com/cozodb/cozo#install)
 //! * [Tutorial](https://nbviewer.org/github/cozodb/cozo-docs/blob/main/tutorial/tutorial.ipynb)
 //! * [Manual for CozoScript](https://cozodb.github.io/current/manual/)
+//!
+//! Example usage:
+//! ```
+//! use cozo::*;
+//!
+//! let db = new_cozo_mem().unwrap();
+//! let script = "?[a] := a in [1, 2, 3]";
+//! let result = db.run_script(script, &Default::default()).unwrap();
+//! println!("{:?}", result);
+//! ```
+//! We created an in-memory database with [`new_cozo_mem`](crate::new_cozo_mem) above.
+//! Persistent options include [`new_cozo_rocksdb`](crate::new_cozo_rocksdb),
+//! [`new_cozo_sqlite`](crate::new_cozo_sqlite) and others.
 
 #![warn(rust_2018_idioms, future_incompatible)]
 #![warn(missing_docs)]
@@ -19,11 +32,13 @@
 pub use miette::Error;
 
 pub use runtime::db::Db;
+pub use runtime::relation::decode_tuple_from_kv;
 pub use storage::mem::{new_cozo_mem, MemStorage};
 pub use storage::rocks::{new_cozo_rocksdb, RocksDbStorage};
 pub use storage::sled::{new_cozo_sled, SledStorage};
 pub use storage::sqlite::{new_cozo_sqlite, SqliteStorage};
 pub use storage::tikv::{new_cozo_tikv, TiKvStorage};
+pub use storage::{Storage, StoreTx};
 
 // pub use storage::re::{new_cozo_redb, ReStorage};
 

cozo-core/src/runtime/db.rs

@@ -89,7 +89,9 @@ lazy_static! {
 }
 
 impl<'s, S: Storage<'s>> Db<S> {
-    /// create a new database with the specified storage
+    /// Create a new database object with the given storage.
+    /// You must call [`initialize`](Self::initialize) immediately after creation.
+    /// Due to lifetime restrictions we are not able to call that for you automatically.
     pub fn new(storage: S) -> Result<Self> {
         let ret = Self {
             db: storage,
@@ -100,7 +102,7 @@ impl<'s, S: Storage<'s>> Db<S> {
         Ok(ret)
     }
 
-    /// should be called after creation of the database to initialize the runtime data.
+    /// Must be called after creation of the database to initialize the runtime state.
     pub fn initialize(&'s self) -> Result<()> {
         self.load_last_ids()?;
         Ok(())

cozo-core/src/runtime/relation.rs

@@ -240,8 +240,10 @@ impl RelationHandle {
     }
 }
 
+/// Decode tuple from key-value pairs. Used for customizing storage
+/// in traite [`StoreTx`](crate::StoreTx).
 #[inline]
-pub(crate) fn decode_tuple_from_kv(key: &[u8], val: &[u8]) -> Tuple {
+pub fn decode_tuple_from_kv(key: &[u8], val: &[u8]) -> Tuple {
     let mut tup = Tuple::decode_from_key(key);
     if !val.is_empty() {
         let vals: Vec<DataValue> = rmp_serde::from_slice(&val[ENCODED_KEY_MIN_LEN..]).unwrap();

cozo-core/src/storage/mem.rs

@@ -18,7 +18,9 @@ use crate::runtime::relation::decode_tuple_from_kv;
 use crate::storage::{Storage, StoreTx};
 use crate::utils::swap_option_result;
 
-/// create a database backed by memory
+/// Create a database backed by memory.
+/// This is the fastest storage, but non-persistent.
+/// Supports concurrent readers but only a single writer.
 pub fn new_cozo_mem() -> Result<crate::Db<MemStorage>> {
     let ret = crate::Db::new(MemStorage::default())?;
 
@@ -186,6 +188,24 @@ impl<'s> StoreTx<'s> for MemTx<'s> {
             }),
         }
     }
+
+    fn batch_put(
+        &mut self,
+        data: Box<dyn Iterator<Item = Result<(Vec<u8>, Vec<u8>)>>>,
+    ) -> Result<()> {
+        match self {
+            MemTx::Reader(_) => {
+                bail!("write in read transaction")
+            }
+            MemTx::Writer(_, cache) => {
+                for pair in data {
+                    let (k, v) = pair?;
+                    cache.insert(k, Some(v));
+                }
+                Ok(())
+            }
+        }
+    }
 }
 
 struct CacheIterRaw<'a> {

cozo-core/src/storage/mod.rs

@@ -6,27 +6,58 @@ use miette::Result;
 
 use crate::data::tuple::Tuple;
 
+pub(crate) mod mem;
 pub(crate) mod rocks;
 pub(crate) mod sled;
-pub(crate) mod tikv;
 pub(crate) mod sqlite;
-pub(crate) mod mem;
+pub(crate) mod tikv;
 // pub(crate) mod re;
 
+/// Swappable storage trait for Cozo's storage engine
 pub trait Storage<'s> {
+    /// The associated transaction type used by this engine
     type Tx: StoreTx<'s>;
 
+    /// Create a transaction object. Write ops will only be called when `write == true`.
     fn transact(&'s self, write: bool) -> Result<Self::Tx>;
+
+    /// Delete a range. It is ok to return immediately and do the deletion in
+    /// the background. It is guaranteed that no keys within the deleted range
+    /// will be accessed in any way by any transaction again.
     fn del_range(&'s self, lower: &[u8], upper: &[u8]) -> Result<()>;
+
+    /// Compact the key range. Can be a no-op if the storage engine does not
+    /// have the concept of compaction.
     fn range_compact(&'s self, lower: &[u8], upper: &[u8]) -> Result<()>;
 }
 
+/// Trait for the associated transaction type of a storage engine.
+/// A transaction needs to guarantee MVCC semantics for all operations.
 pub trait StoreTx<'s> {
+    /// Get a key. If `for_update` is `true` (only possible in a write transaction),
+    /// then the database needs to guarantee that `commit()` can only succeed if
+    /// the key has not been modified outside the transaction.
     fn get(&self, key: &[u8], for_update: bool) -> Result<Option<Vec<u8>>>;
+
+    /// Put a key-value pair into the storage. In case of existing key,
+    /// the storage engine needs to overwrite the old value.
     fn put(&mut self, key: &[u8], val: &[u8]) -> Result<()>;
+
+    /// Delete a key-value pair from the storage.
     fn del(&mut self, key: &[u8]) -> Result<()>;
+
+    /// Check if a key exists. If `for_update` is `true` (only possible in a write transaction),
+    /// then the database needs to guarantee that `commit()` can only succeed if
+    /// the key has not been modified outside the transaction.
     fn exists(&self, key: &[u8], for_update: bool) -> Result<bool>;
+
+    /// Commit a transaction. Must return an `Err` if MVCC consistency cannot be guaranteed,
+    /// and discard all changes introduced by this transaction.
     fn commit(&mut self) -> Result<()>;
+
+    /// Scan on a range. `lower` is inclusive whereas `upper` is exclusive.
+    /// The implementation must call [`decode_tuple_from_kv`](crate::decode_tuple_from_kv) to obtain
+    /// a decoded tuple in the loop of the iterator.
     fn range_scan<'a>(
         &'a self,
         lower: &[u8],
@@ -34,6 +65,9 @@ pub trait StoreTx<'s> {
     ) -> Box<dyn Iterator<Item = Result<Tuple>> + 'a>
     where
         's: 'a;
+
+    /// Scan on a range and return the raw results.
+    /// `lower` is inclusive whereas `upper` is exclusive.
     fn range_scan_raw<'a>(
         &'a self,
         lower: &[u8],
@@ -41,4 +75,18 @@ pub trait StoreTx<'s> {
     ) -> Box<dyn Iterator<Item = Result<(Vec<u8>, Vec<u8>)>> + 'a>
     where
         's: 'a;
+
+    /// Put multiple key-value pairs into the database.
+    /// The default implementation just calls `put` repeatedly.
+    /// Implement if there is a more efficient way.
+    fn batch_put(
+        &mut self,
+        data: Box<dyn Iterator<Item = Result<(Vec<u8>, Vec<u8>)>>>,
+    ) -> Result<()> {
+        for pair in data {
+            let (k, v) = pair?;
+            self.put(&k, &v)?;
+        }
+        Ok(())
+    }
 }

cozo-core/src/storage/rocks.rs

@@ -17,6 +17,9 @@ use crate::utils::swap_option_result;
 use crate::Db;
 
 /// Creates a RocksDB database object.
+/// This is currently the fastest persistent storage and it can
+/// sustain huge concurrency.
+/// Supports concurrent readers and writers.
 pub fn new_cozo_rocksdb(path: impl AsRef<str>) -> Result<Db<RocksDbStorage>> {
     let builder = DbBuilder::default().path(path.as_ref());
     let path = builder.opts.db_path;

cozo-core/src/storage/sled.rs

@@ -16,7 +16,9 @@ use crate::runtime::relation::decode_tuple_from_kv;
 use crate::storage::{Storage, StoreTx};
 use crate::utils::swap_option_result;
 
-/// Creates a Sled database object.
+/// Creates a Sled database object. Experimental.
+/// You should use [`new_cozo_rocksdb`](crate::new_cozo_rocksdb) or
+/// [`new_cozo_sqlite`](crate::new_cozo_sqlite) instead.
 pub fn new_cozo_sled(path: impl AsRef<Path>) -> Result<crate::Db<SledStorage>> {
     let db = sled::open(path).into_diagnostic()?;
     let ret = crate::Db::new(SledStorage { db })?;

cozo-core/src/storage/sqlite.rs

@@ -6,7 +6,7 @@ use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 use ::sqlite::Connection;
 use either::{Either, Left, Right};
-use miette::{bail, miette, IntoDiagnostic, Result, WrapErr};
+use miette::{bail, miette, IntoDiagnostic, Result};
 use sqlite::{State, Statement};
 
 use crate::data::tuple::Tuple;
@@ -19,7 +19,14 @@ pub struct SqliteStorage {
     name: String,
 }
 
-/// create a sqlite backed database. `:memory:` is not OK.
+/// Create a sqlite backed database.
+/// This is slower than [`new_cozo_rocksdb`](crate::new_cozo_rocksdb)
+/// but uses way less resources and is much easier to compile for exotic
+/// environments.
+/// Supports concurrent readers but only a single writer.
+///
+/// You must provide a disk-based path: `:memory:` is not OK.
+/// If you want a pure memory storage, use [`new_cozo_mem`](crate::new_cozo_mem).
 pub fn new_cozo_sqlite(path: String) -> Result<crate::Db<SqliteStorage>> {
     let conn = sqlite::open(&path).into_diagnostic()?;
     let query = r#"
@@ -129,12 +136,7 @@ impl<'s> StoreTx<'s> for SqliteTx<'s> {
         let mut statement = self.conn.prepare(query).unwrap();
         statement.bind((1, key)).unwrap();
         statement.bind((2, val)).unwrap();
-        while statement
-            .next()
-            .into_diagnostic()
-            .with_context(|| format!("{:x?} {:?} {:x?}", key, val, Tuple::decode_from_key(key)))?
-            != State::Done
-        {}
+        while statement.next().into_diagnostic()? != State::Done {}
         Ok(())
     }
 
@@ -210,6 +212,26 @@ impl<'s> StoreTx<'s> for SqliteTx<'s> {
         statement.bind((2, upper)).unwrap();
         Box::new(RawIter(statement))
     }
+
+    fn batch_put(
+        &mut self,
+        data: Box<dyn Iterator<Item = Result<(Vec<u8>, Vec<u8>)>>>,
+    ) -> Result<()> {
+        let query = r#"
+                insert into cozo(k, v) values (?, ?)
+                on conflict(k) do update set v=excluded.v;
+            "#;
+
+        let mut statement = self.conn.prepare(query).unwrap();
+        for pair in data {
+            let (key, val) = pair?;
+            statement.bind((1, key.as_slice())).unwrap();
+            statement.bind((2, val.as_slice())).unwrap();
+            while statement.next().into_diagnostic()? != State::Done {}
+            statement.reset().into_diagnostic()?;
+        }
+        Ok(())
+    }
 }
 
 struct TupleIter<'l>(Statement<'l>);

cozo-core/src/storage/tikv.rs

@@ -18,7 +18,8 @@ use crate::storage::{Storage, StoreTx};
 use crate::utils::swap_option_result;
 use crate::Db;
 
-/// connect to a Storage engine backed by TiKV
+/// Connect to a Storage engine backed by TiKV.
+/// Experimental and very slow.
 pub fn new_cozo_tikv(pd_endpoints: Vec<String>, optimistic: bool) -> Result<Db<TiKvStorage>> {
     let raw_client = RT
         .block_on(RawClient::new(pd_endpoints.clone()))