slatedb 0.1.1 → 0.3.1

data/ext/slatedb/src/transaction.rs

@@ -2,8 +2,11 @@ use std::cell::RefCell;
 
 use magnus::prelude::*;
 use magnus::{method, Error, RHash, Ruby};
-use slatedb::config::{DurabilityLevel, PutOptions, ReadOptions, ScanOptions, Ttl, WriteOptions};
-use slatedb::DBTransaction;
+use slatedb::config::{
+    DurabilityLevel, MergeOptions, PutOptions, ReadOptions, ScanOptions, Ttl, WriteOptions,
+};
+use slatedb::DbTransaction;
+use slatedb::IterationOrder;
 
 use crate::errors::{closed_error, invalid_argument_error, map_error};
 use crate::iterator::Iterator;
@@ -16,12 +19,12 @@ use crate::utils::get_optional;
 /// After commit or rollback, the transaction is closed.
 #[magnus::wrap(class = "SlateDb::Transaction", free_immediately, size)]
 pub struct Transaction {
-    inner: RefCell<Option<DBTransaction>>,
+    inner: RefCell<Option<DbTransaction>>,
 }
 
 impl Transaction {
-    /// Create a new Transaction from a DBTransaction.
-    pub fn new(txn: DBTransaction) -> Self {
+    /// Create a new Transaction from a DbTransaction.
+    pub fn new(txn: DbTransaction) -> Self {
         Self {
             inner: RefCell::new(Some(txn)),
         }
@@ -67,13 +70,16 @@ impl Transaction {
             opts.dirty = dirty;
         }
 
+        if let Some(cb) = get_optional::<bool>(&kwargs, "cache_blocks")? {
+            opts.cache_blocks = cb;
+        }
+
         let guard = self.inner.borrow();
         let txn = guard
             .as_ref()
             .ok_or_else(|| closed_error("transaction is closed"))?;
 
-        let result =
-            block_on_result(async { txn.get_with_options(key.as_bytes(), &opts).await })?;
+        let result = block_on_result(async { txn.get_with_options(key.as_bytes(), &opts).await })?;
         Ok(result.map(|b| String::from_utf8_lossy(&b).to_string()))
     }
 
@@ -135,6 +141,53 @@ impl Transaction {
         Ok(())
     }
 
+    /// Merge a value within the transaction.
+    pub fn merge(&self, key: String, value: String) -> Result<(), Error> {
+        if key.is_empty() {
+            return Err(invalid_argument_error("key cannot be empty"));
+        }
+
+        let guard = self.inner.borrow();
+        let txn = guard
+            .as_ref()
+            .ok_or_else(|| closed_error("transaction is closed"))?;
+
+        txn.merge(key.as_bytes(), value.as_bytes())
+            .map_err(map_error)?;
+
+        Ok(())
+    }
+
+    /// Merge a value with options within the transaction.
+    pub fn merge_with_options(
+        &self,
+        key: String,
+        value: String,
+        kwargs: RHash,
+    ) -> Result<(), Error> {
+        if key.is_empty() {
+            return Err(invalid_argument_error("key cannot be empty"));
+        }
+
+        let ttl = get_optional::<u64>(&kwargs, "ttl")?;
+        let merge_opts = MergeOptions {
+            ttl: match ttl {
+                Some(ms) => Ttl::ExpireAfter(ms),
+                None => Ttl::Default,
+            },
+        };
+
+        let guard = self.inner.borrow();
+        let txn = guard
+            .as_ref()
+            .ok_or_else(|| closed_error("transaction is closed"))?;
+
+        txn.merge_with_options(key.as_bytes(), value.as_bytes(), &merge_opts)
+            .map_err(map_error)?;
+
+        Ok(())
+    }
+
     /// Scan a range of keys within the transaction.
     pub fn scan(&self, start: String, end_key: Option<String>) -> Result<Iterator, Error> {
         if start.is_empty() {
@@ -200,6 +253,18 @@ impl Transaction {
         if let Some(mft) = get_optional::<usize>(&kwargs, "max_fetch_tasks")? {
             opts.max_fetch_tasks = mft;
         }
+        if let Some(order) = get_optional::<String>(&kwargs, "order")? {
+            opts.order = match order.as_str() {
+                "ascending" | "asc" => IterationOrder::Ascending,
+                "descending" | "desc" => IterationOrder::Descending,
+                other => {
+                    return Err(invalid_argument_error(&format!(
+                        "invalid order: {} (expected 'asc' or 'desc')",
+                        other
+                    )))
+                }
+            };
+        }
 
         let guard = self.inner.borrow();
         let txn = guard
@@ -219,6 +284,106 @@ impl Transaction {
         Ok(Iterator::new(iter))
     }
 
+    /// Scan all keys with a given prefix within the transaction.
+    pub fn scan_prefix(&self, prefix: String) -> Result<Iterator, Error> {
+        if prefix.is_empty() {
+            return Err(invalid_argument_error("prefix cannot be empty"));
+        }
+
+        let guard = self.inner.borrow();
+        let txn = guard
+            .as_ref()
+            .ok_or_else(|| closed_error("transaction is closed"))?;
+
+        let iter = block_on_result(async { txn.scan_prefix(prefix.as_bytes()).await })?;
+
+        Ok(Iterator::new(iter))
+    }
+
+    /// Scan all keys with a given prefix with options within the transaction.
+    pub fn scan_prefix_with_options(
+        &self,
+        prefix: String,
+        kwargs: RHash,
+    ) -> Result<Iterator, Error> {
+        if prefix.is_empty() {
+            return Err(invalid_argument_error("prefix cannot be empty"));
+        }
+
+        let mut opts = ScanOptions::default();
+
+        if let Some(df) = get_optional::<String>(&kwargs, "durability_filter")? {
+            opts.durability_filter = match df.as_str() {
+                "remote" => DurabilityLevel::Remote,
+                "memory" => DurabilityLevel::Memory,
+                other => {
+                    return Err(invalid_argument_error(&format!(
+                        "invalid durability_filter: {} (expected 'remote' or 'memory')",
+                        other
+                    )))
+                }
+            };
+        }
+
+        if let Some(dirty) = get_optional::<bool>(&kwargs, "dirty")? {
+            opts.dirty = dirty;
+        }
+
+        if let Some(rab) = get_optional::<usize>(&kwargs, "read_ahead_bytes")? {
+            opts.read_ahead_bytes = rab;
+        }
+
+        if let Some(cb) = get_optional::<bool>(&kwargs, "cache_blocks")? {
+            opts.cache_blocks = cb;
+        }
+
+        if let Some(mft) = get_optional::<usize>(&kwargs, "max_fetch_tasks")? {
+            opts.max_fetch_tasks = mft;
+        }
+        if let Some(order) = get_optional::<String>(&kwargs, "order")? {
+            opts.order = match order.as_str() {
+                "ascending" | "asc" => IterationOrder::Ascending,
+                "descending" | "desc" => IterationOrder::Descending,
+                other => {
+                    return Err(invalid_argument_error(&format!(
+                        "invalid order: {} (expected 'asc' or 'desc')",
+                        other
+                    )))
+                }
+            };
+        }
+
+        let guard = self.inner.borrow();
+        let txn = guard
+            .as_ref()
+            .ok_or_else(|| closed_error("transaction is closed"))?;
+
+        let iter = block_on_result(async {
+            txn.scan_prefix_with_options(prefix.as_bytes(), &opts).await
+        })?;
+
+        Ok(Iterator::new(iter))
+    }
+
+    /// Mark keys as read for conflict detection.
+    ///
+    /// This explicitly tracks reads for conflict checking in serializable isolation,
+    /// even when the keys weren't actually read via get().
+    ///
+    /// # Arguments
+    /// * `keys` - Array of keys to mark as read
+    pub fn mark_read(&self, keys: Vec<String>) -> Result<(), Error> {
+        let guard = self.inner.borrow();
+        let txn = guard
+            .as_ref()
+            .ok_or_else(|| closed_error("transaction is closed"))?;
+
+        let key_bytes: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
+        txn.mark_read(&key_bytes).map_err(map_error)?;
+
+        Ok(())
+    }
+
     /// Commit the transaction.
     pub fn commit(&self) -> Result<(), Error> {
         let txn = self
@@ -234,7 +399,11 @@ impl Transaction {
     /// Commit the transaction with options.
     pub fn commit_with_options(&self, kwargs: RHash) -> Result<(), Error> {
         let await_durable = get_optional::<bool>(&kwargs, "await_durable")?.unwrap_or(true);
-        let write_opts = WriteOptions { await_durable };
+        let seqnum = get_optional::<u64>(&kwargs, "seqnum")?.unwrap_or(0);
+        let write_opts = WriteOptions {
+            await_durable,
+            seqnum,
+        };
 
         let txn = self
             .inner
@@ -275,12 +444,23 @@ pub fn define_transaction_class(ruby: &Ruby, module: &magnus::RModule) -> Result
         method!(Transaction::put_with_options, 3),
     )?;
     class.define_method("_delete", method!(Transaction::delete, 1))?;
+    class.define_method("_merge", method!(Transaction::merge, 2))?;
+    class.define_method(
+        "_merge_with_options",
+        method!(Transaction::merge_with_options, 3),
+    )?;
     class.define_method("_scan", method!(Transaction::scan, 2))?;
     class.define_method(
         "_scan_with_options",
         method!(Transaction::scan_with_options, 3),
     )?;
-    class.define_method("commit", method!(Transaction::commit, 0))?;
+    class.define_method("_scan_prefix", method!(Transaction::scan_prefix, 1))?;
+    class.define_method(
+        "_scan_prefix_with_options",
+        method!(Transaction::scan_prefix_with_options, 2),
+    )?;
+    class.define_method("_mark_read", method!(Transaction::mark_read, 1))?;
+    class.define_method("_commit", method!(Transaction::commit, 0))?;
     class.define_method(
         "_commit_with_options",
         method!(Transaction::commit_with_options, 1),

data/ext/slatedb/src/utils.rs

@@ -2,8 +2,8 @@ use std::sync::Arc;
 
 use magnus::value::ReprValue;
 use magnus::{Error, RHash, Ruby, TryConvert};
-use object_store::aws::AmazonS3Builder;
-use object_store::ObjectStoreScheme;
+use slatedb::object_store::aws::AmazonS3Builder;
+use slatedb::object_store::{Error as ObjectStoreError, ObjectStore, ObjectStoreScheme};
 use slatedb::{Db, Error as SlateError};
 use url::Url;
 
@@ -24,7 +24,7 @@ pub fn get_optional<T: TryConvert>(hash: &RHash, key: &str) -> Result<Option<T>,
 }
 
 /// Convert an object_store error to a SlateDB error
-fn to_slate_error(e: object_store::Error) -> SlateError {
+fn to_slate_error(e: ObjectStoreError) -> SlateError {
     SlateError::unavailable(e.to_string())
 }
 
@@ -33,7 +33,7 @@ fn to_slate_error(e: object_store::Error) -> SlateError {
 /// This function handles S3 URLs specially to ensure environment variables
 /// like AWS_ACCESS_KEY_ID are properly recognized (the default object_store
 /// registry only recognizes lowercase variants like aws_access_key_id).
-pub fn resolve_object_store(url: &str) -> Result<Arc<dyn object_store::ObjectStore>, SlateError> {
+pub fn resolve_object_store(url: &str) -> Result<Arc<dyn ObjectStore>, SlateError> {
     let parsed_url: Url = url
         .try_into()
         .map_err(|e: url::ParseError| SlateError::invalid(format!("invalid URL: {}", e)))?;
@@ -44,6 +44,7 @@ pub fn resolve_object_store(url: &str) -> Result<Arc<dyn object_store::ObjectSto
     match scheme {
         ObjectStoreScheme::AmazonS3 => {
             // Use from_env() to properly handle uppercase AWS_* environment variables
+            // (the default object_store registry only recognizes lowercase variants)
             let store = AmazonS3Builder::from_env()
                 .with_url(url)
                 .build()

data/ext/slatedb/src/write_batch.rs

@@ -2,7 +2,7 @@ use std::cell::RefCell;
 
 use magnus::prelude::*;
 use magnus::{function, method, Error, RHash, Ruby};
-use slatedb::config::{PutOptions, Ttl};
+use slatedb::config::{MergeOptions, PutOptions, Ttl};
 use slatedb::WriteBatch as SlateWriteBatch;
 
 use crate::errors::invalid_argument_error;
@@ -72,6 +72,48 @@ impl WriteBatch {
         Ok(())
     }
 
+    /// Add a merge operation to the batch.
+    pub fn merge(&self, key: String, value: String) -> Result<(), Error> {
+        if key.is_empty() {
+            return Err(invalid_argument_error("key cannot be empty"));
+        }
+
+        self.inner
+            .borrow_mut()
+            .merge(key.as_bytes(), value.as_bytes());
+
+        Ok(())
+    }
+
+    /// Add a merge operation with options to the batch.
+    ///
+    /// Options:
+    /// - ttl: Time-to-live in milliseconds
+    pub fn merge_with_options(
+        &self,
+        key: String,
+        value: String,
+        kwargs: RHash,
+    ) -> Result<(), Error> {
+        if key.is_empty() {
+            return Err(invalid_argument_error("key cannot be empty"));
+        }
+
+        let ttl = get_optional::<u64>(&kwargs, "ttl")?;
+        let merge_opts = MergeOptions {
+            ttl: match ttl {
+                Some(ms) => Ttl::ExpireAfter(ms),
+                None => Ttl::Default,
+            },
+        };
+
+        self.inner
+            .borrow_mut()
+            .merge_with_options(key.as_bytes(), value.as_bytes(), &merge_opts);
+
+        Ok(())
+    }
+
     /// Take ownership of the inner WriteBatch (consumes it).
     /// Used internally when writing the batch to the database.
     pub fn take(&self) -> Result<SlateWriteBatch, Error> {
@@ -93,6 +135,11 @@ pub fn define_write_batch_class(ruby: &Ruby, module: &magnus::RModule) -> Result
         method!(WriteBatch::put_with_options, 3),
     )?;
     class.define_method("_delete", method!(WriteBatch::delete, 1))?;
+    class.define_method("_merge", method!(WriteBatch::merge, 2))?;
+    class.define_method(
+        "_merge_with_options",
+        method!(WriteBatch::merge_with_options, 3),
+    )?;
 
     Ok(())
 }