slatedb 0.2.0 → 0.4.4

data/ext/slatedb/src/database.rs

@@ -1,4 +1,5 @@
-use std::sync::Arc;
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
 
 use magnus::prelude::*;
 use magnus::{function, method, Error, RHash, Ruby};
@@ -6,15 +7,16 @@ use slatedb::config::{
     DurabilityLevel, MergeOptions, PutOptions, ReadOptions, ScanOptions, Ttl, WriteOptions,
 };
 use slatedb::object_store::memory::InMemory;
-use slatedb::{Db, IsolationLevel};
+use slatedb::{Db, IsolationLevel, IterationOrder, KeyValue};
 
 use crate::errors::invalid_argument_error;
 use crate::iterator::Iterator;
 use crate::merge_ops::{parse_merge_operator, parse_merge_operator_proc};
+use crate::metrics::Metrics;
 use crate::runtime::block_on_result;
 use crate::snapshot::Snapshot;
 use crate::transaction::Transaction;
-use crate::utils::{get_optional, resolve_object_store};
+use crate::utils::{get_optional, prefix_subrange_from_kwargs, resolve_object_store};
 use crate::write_batch::WriteBatch;
 
 /// Ruby wrapper for SlateDB database.
@@ -23,9 +25,69 @@ use crate::write_batch::WriteBatch;
 #[magnus::wrap(class = "SlateDb::Database", free_immediately, size)]
 pub struct Database {
     inner: Arc<Db>,
+    metrics: Arc<Mutex<HashMap<String, i64>>>,
 }
 
 impl Database {
+    fn increment_metric(&self, name: &str) {
+        let mut metrics = self.metrics.lock().expect("metrics mutex poisoned");
+        *metrics.entry(name.to_string()).or_insert(0) += 1;
+    }
+
+    fn key_value_to_hash(kv: KeyValue) -> Result<RHash, Error> {
+        let ruby = Ruby::get().expect("Ruby runtime not available");
+        let hash = ruby.hash_new();
+        hash.aset(
+            ruby.to_symbol("key"),
+            String::from_utf8_lossy(&kv.key).to_string(),
+        )?;
+        hash.aset(
+            ruby.to_symbol("value"),
+            String::from_utf8_lossy(&kv.value).to_string(),
+        )?;
+        hash.aset(ruby.to_symbol("seq"), kv.seq)?;
+        hash.aset(ruby.to_symbol("create_ts"), kv.create_ts)?;
+        hash.aset(ruby.to_symbol("expire_ts"), kv.expire_ts)?;
+        Ok(hash)
+    }
+
+    fn read_options_from_kwargs(kwargs: &RHash) -> Result<ReadOptions, Error> {
+        let mut opts = ReadOptions::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(cb) = get_optional::<bool>(kwargs, "cache_blocks")? {
+            opts.cache_blocks = cb;
+        }
+
+        Ok(opts)
+    }
+
+    fn write_options_from_kwargs(kwargs: &RHash) -> Result<WriteOptions, Error> {
+        let await_durable = get_optional::<bool>(kwargs, "await_durable")?.unwrap_or(true);
+        let seqnum = get_optional::<u64>(kwargs, "seqnum")?.unwrap_or(0);
+
+        Ok(WriteOptions {
+            await_durable,
+            seqnum,
+        })
+    }
+
     /// Open a database at the given path.
     ///
     /// # Arguments
@@ -37,8 +99,7 @@ impl Database {
     /// A new Database instance
     pub fn open(path: String, url: Option<String>, kwargs: RHash) -> Result<Self, Error> {
         // Try string-based merge operator first, then proc-based
-        let merge_operator = parse_merge_operator(&kwargs)?
-            .or(parse_merge_operator_proc(&kwargs)?);
+        let merge_operator = parse_merge_operator(&kwargs)?.or(parse_merge_operator_proc(&kwargs)?);
 
         let db = block_on_result(async {
             let object_store: Arc<dyn slatedb::object_store::ObjectStore> =
@@ -58,6 +119,7 @@ impl Database {
 
         Ok(Self {
             inner: Arc::new(db),
+            metrics: Arc::new(Mutex::new(HashMap::new())),
         })
     }
 
@@ -77,6 +139,7 @@ impl Database {
 
         let result =
             block_on_result(async { self.inner.get_with_options(key.as_bytes(), &opts).await })?;
+        self.increment_metric("db.get.count");
 
         Ok(result.map(|b| String::from_utf8_lossy(&b).to_string()))
     }
@@ -94,31 +157,64 @@ impl Database {
             return Err(invalid_argument_error("key cannot be empty"));
         }
 
-        let mut opts = ReadOptions::default();
+        let opts = Self::read_options_from_kwargs(&kwargs)?;
 
-        // Parse durability_filter
-        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
-                    )))
-                }
-            };
+        let result =
+            block_on_result(async { self.inner.get_with_options(key.as_bytes(), &opts).await })?;
+        self.increment_metric("db.get_with_options.count");
+
+        Ok(result.map(|b| String::from_utf8_lossy(&b).to_string()))
+    }
+
+    /// Get a key-value pair with metadata by key.
+    ///
+    /// # Arguments
+    /// * `key` - The key to look up
+    ///
+    /// # Returns
+    /// A Hash with key, value, seq, create_ts, and expire_ts, or nil if not found
+    pub fn get_key_value(&self, key: String) -> Result<Option<RHash>, Error> {
+        if key.is_empty() {
+            return Err(invalid_argument_error("key cannot be empty"));
         }
 
-        // Parse dirty
-        if let Some(dirty) = get_optional::<bool>(&kwargs, "dirty")? {
-            opts.dirty = dirty;
+        let opts = ReadOptions::default();
+        let result = block_on_result(async {
+            self.inner
+                .get_key_value_with_options(key.as_bytes(), &opts)
+                .await
+        })?;
+        self.increment_metric("db.get_key_value.count");
+
+        result.map(Self::key_value_to_hash).transpose()
+    }
+
+    /// Get a key-value pair with metadata by key with options.
+    ///
+    /// # Arguments
+    /// * `key` - The key to look up
+    /// * `kwargs` - Keyword arguments (durability_filter, dirty, cache_blocks)
+    ///
+    /// # Returns
+    /// A Hash with key, value, seq, create_ts, and expire_ts, or nil if not found
+    pub fn get_key_value_with_options(
+        &self,
+        key: String,
+        kwargs: RHash,
+    ) -> Result<Option<RHash>, Error> {
+        if key.is_empty() {
+            return Err(invalid_argument_error("key cannot be empty"));
         }
 
-        let result =
-            block_on_result(async { self.inner.get_with_options(key.as_bytes(), &opts).await })?;
+        let opts = Self::read_options_from_kwargs(&kwargs)?;
+        let result = block_on_result(async {
+            self.inner
+                .get_key_value_with_options(key.as_bytes(), &opts)
+                .await
+        })?;
+        self.increment_metric("db.get_key_value_with_options.count");
 
-        Ok(result.map(|b| String::from_utf8_lossy(&b).to_string()))
+        result.map(Self::key_value_to_hash).transpose()
     }
 
     /// Get a value by key as raw bytes.
@@ -137,6 +233,7 @@ impl Database {
 
         let result =
             block_on_result(async { self.inner.get_with_options(key.as_bytes(), &opts).await })?;
+        self.increment_metric("db.get_bytes.count");
 
         Ok(result.map(|b| b.to_vec()))
     }
@@ -155,6 +252,7 @@ impl Database {
 
         let write_opts = WriteOptions {
             await_durable: true,
+            seqnum: 0,
         };
 
         block_on_result(async {
@@ -162,6 +260,7 @@ impl Database {
                 .put_with_options(key.as_bytes(), value.as_bytes(), &put_opts, &write_opts)
                 .await
         })?;
+        self.increment_metric("db.put.count");
 
         Ok(())
     }
@@ -171,7 +270,7 @@ impl Database {
     /// # Arguments
     /// * `key` - The key to store
     /// * `value` - The value to store
-    /// * `kwargs` - Keyword arguments (ttl, await_durable)
+    /// * `kwargs` - Keyword arguments (ttl, await_durable, seqnum)
     pub fn put_with_options(&self, key: String, value: String, kwargs: RHash) -> Result<(), Error> {
         if key.is_empty() {
             return Err(invalid_argument_error("key cannot be empty"));
@@ -187,14 +286,14 @@ impl Database {
         };
 
         // Parse await_durable
-        let await_durable = get_optional::<bool>(&kwargs, "await_durable")?.unwrap_or(true);
-        let write_opts = WriteOptions { await_durable };
+        let write_opts = Self::write_options_from_kwargs(&kwargs)?;
 
         block_on_result(async {
             self.inner
                 .put_with_options(key.as_bytes(), value.as_bytes(), &put_opts, &write_opts)
                 .await
         })?;
+        self.increment_metric("db.put_with_options.count");
 
         Ok(())
     }
@@ -210,6 +309,7 @@ impl Database {
 
         let write_opts = WriteOptions {
             await_durable: true,
+            seqnum: 0,
         };
 
         block_on_result(async {
@@ -217,6 +317,7 @@ impl Database {
                 .delete_with_options(key.as_bytes(), &write_opts)
                 .await
         })?;
+        self.increment_metric("db.delete.count");
 
         Ok(())
     }
@@ -225,20 +326,20 @@ impl Database {
     ///
     /// # Arguments
     /// * `key` - The key to delete
-    /// * `kwargs` - Keyword arguments (await_durable)
+    /// * `kwargs` - Keyword arguments (await_durable, seqnum)
     pub fn delete_with_options(&self, key: String, kwargs: RHash) -> Result<(), Error> {
         if key.is_empty() {
             return Err(invalid_argument_error("key cannot be empty"));
         }
 
-        let await_durable = get_optional::<bool>(&kwargs, "await_durable")?.unwrap_or(true);
-        let write_opts = WriteOptions { await_durable };
+        let write_opts = Self::write_options_from_kwargs(&kwargs)?;
 
         block_on_result(async {
             self.inner
                 .delete_with_options(key.as_bytes(), &write_opts)
                 .await
         })?;
+        self.increment_metric("db.delete_with_options.count");
 
         Ok(())
     }
@@ -325,6 +426,18 @@ impl Database {
         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 start_bytes = start.into_bytes();
         let end_bytes = end_key.map(|e| e.into_bytes());
@@ -354,7 +467,7 @@ impl Database {
         let opts = ScanOptions::default();
         let iter = block_on_result(async {
             self.inner
-                .scan_prefix_with_options(prefix.as_bytes(), &opts)
+                .scan_prefix_with_options(prefix.as_bytes(), .., &opts)
                 .await
         })?;
 
@@ -408,10 +521,23 @@ impl Database {
         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 subrange = prefix_subrange_from_kwargs(&kwargs)?;
         let iter = block_on_result(async {
             self.inner
-                .scan_prefix_with_options(prefix.as_bytes(), &opts)
+                .scan_prefix_with_options(prefix.as_bytes(), subrange, &opts)
                 .await
         })?;
 
@@ -432,10 +558,9 @@ impl Database {
     ///
     /// # Arguments
     /// * `batch` - The WriteBatch to write
-    /// * `kwargs` - Keyword arguments (await_durable)
+    /// * `kwargs` - Keyword arguments (await_durable, seqnum)
     pub fn write_with_options(&self, batch: &WriteBatch, kwargs: RHash) -> Result<(), Error> {
-        let await_durable = get_optional::<bool>(&kwargs, "await_durable")?.unwrap_or(true);
-        let write_opts = WriteOptions { await_durable };
+        let write_opts = Self::write_options_from_kwargs(&kwargs)?;
 
         let batch_inner = batch.take()?;
 
@@ -462,6 +587,7 @@ impl Database {
 
         let write_opts = WriteOptions {
             await_durable: true,
+            seqnum: 0,
         };
 
         block_on_result(async {
@@ -478,8 +604,13 @@ impl Database {
     /// # Arguments
     /// * `key` - The key to merge into
     /// * `value` - The merge operand to apply
-    /// * `kwargs` - Keyword arguments (ttl, await_durable)
-    pub fn merge_with_options(&self, key: String, value: String, kwargs: RHash) -> Result<(), Error> {
+    /// * `kwargs` - Keyword arguments (ttl, await_durable, seqnum)
+    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"));
         }
@@ -492,8 +623,7 @@ impl Database {
             },
         };
 
-        let await_durable = get_optional::<bool>(&kwargs, "await_durable")?.unwrap_or(true);
-        let write_opts = WriteOptions { await_durable };
+        let write_opts = Self::write_options_from_kwargs(&kwargs)?;
 
         block_on_result(async {
             self.inner
@@ -548,8 +678,8 @@ impl Database {
     pub fn create_checkpoint(&self, kwargs: RHash) -> Result<RHash, Error> {
         use slatedb::config::{CheckpointOptions, CheckpointScope};
 
-        let lifetime = get_optional::<u64>(&kwargs, "lifetime")?
-            .map(std::time::Duration::from_millis);
+        let lifetime =
+            get_optional::<u64>(&kwargs, "lifetime")?.map(std::time::Duration::from_millis);
         let name = get_optional::<String>(&kwargs, "name")?;
 
         let options = CheckpointOptions {
@@ -578,6 +708,11 @@ impl Database {
         Ok(())
     }
 
+    /// Return the database metrics registry.
+    pub fn metrics(&self) -> Result<Metrics, Error> {
+        Ok(Metrics::new(self.metrics.clone()))
+    }
+
     /// Close the database.
     pub fn close(&self) -> Result<(), Error> {
         block_on_result(async { self.inner.close().await })?;
@@ -595,6 +730,11 @@ pub fn define_database_class(ruby: &Ruby, module: &magnus::RModule) -> Result<()
     // Instance methods - simple versions
     class.define_method("_get", method!(Database::get, 1))?;
     class.define_method("_get_with_options", method!(Database::get_with_options, 2))?;
+    class.define_method("_get_key_value", method!(Database::get_key_value, 1))?;
+    class.define_method(
+        "_get_key_value_with_options",
+        method!(Database::get_key_value_with_options, 2),
+    )?;
     class.define_method("get_bytes", method!(Database::get_bytes, 1))?;
     class.define_method("_put", method!(Database::put, 2))?;
     class.define_method("_put_with_options", method!(Database::put_with_options, 3))?;
@@ -633,6 +773,7 @@ pub fn define_database_class(ruby: &Ruby, module: &magnus::RModule) -> Result<()
         method!(Database::create_checkpoint, 1),
     )?;
     class.define_method("flush", method!(Database::flush, 0))?;
+    class.define_method("_metrics", method!(Database::metrics, 0))?;
     class.define_method("close", method!(Database::close, 0))?;
 
     Ok(())

data/ext/slatedb/src/iterator.rs

@@ -98,7 +98,10 @@ impl Iterator {
         let result = block_on(async {
             let mut guard = inner.lock().await;
             match guard.as_mut() {
-                Some(iter) => iter.seek(key.as_bytes()).await.map_err(IteratorError::Slate),
+                Some(iter) => iter
+                    .seek(key.as_bytes())
+                    .await
+                    .map_err(IteratorError::Slate),
                 None => Err(IteratorError::Closed),
             }
         });

data/ext/slatedb/src/lib.rs

@@ -21,6 +21,7 @@ mod database;
 mod errors;
 mod iterator;
 mod merge_ops;
+mod metrics;
 mod reader;
 mod runtime;
 mod snapshot;
@@ -46,6 +47,7 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     snapshot::define_snapshot_class(ruby, &module)?;
     reader::define_reader_class(ruby, &module)?;
     admin::define_admin_class(ruby, &module)?;
+    metrics::define_metrics_class(ruby, &module)?;
 
     Ok(())
 }

data/ext/slatedb/src/merge_ops.rs

@@ -149,9 +149,7 @@ impl RubyProcMergeOperator {
             "Ruby merge operator called from non-Ruby thread, using fallback concatenation. \
              This can happen during background compaction."
         );
-        let mut result = existing_value
-            .map(|v| v.to_vec())
-            .unwrap_or_default();
+        let mut result = existing_value.map(|v| v.to_vec()).unwrap_or_default();
         result.extend_from_slice(new_value);
         Ok(Bytes::from(result))
     }

data/ext/slatedb/src/metrics.rs

@@ -0,0 +1,47 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+use magnus::prelude::*;
+use magnus::{method, Error, Ruby};
+/// Ruby wrapper for SlateDB metrics registry.
+///
+/// This struct is exposed to Ruby as `SlateDb::Metrics`.
+#[magnus::wrap(class = "SlateDb::Metrics", free_immediately, size)]
+pub struct Metrics {
+    inner: Arc<Mutex<HashMap<String, i64>>>,
+}
+
+impl Metrics {
+    pub fn new(inner: Arc<Mutex<HashMap<String, i64>>>) -> Self {
+        Self { inner }
+    }
+
+    /// Return a list of metric names.
+    pub fn names(&self) -> Result<magnus::RArray, Error> {
+        let metrics = self.inner.lock().expect("metrics mutex poisoned");
+        let ruby = Ruby::get().expect("Ruby runtime not available");
+        let result = ruby.ary_new_capa(metrics.len());
+
+        for name in metrics.keys() {
+            result.push(ruby.str_new(name))?;
+        }
+
+        Ok(result)
+    }
+
+    /// Get the current value of a metric by name.
+    pub fn get(&self, name: String) -> Result<Option<i64>, Error> {
+        let metrics = self.inner.lock().expect("metrics mutex poisoned");
+        Ok(metrics.get(&name).copied())
+    }
+}
+
+/// Define the Metrics class on the SlateDb module.
+pub fn define_metrics_class(ruby: &Ruby, module: &magnus::RModule) -> Result<(), Error> {
+    let class = module.define_class("Metrics", ruby.class_object())?;
+
+    class.define_method("names", method!(Metrics::names, 0))?;
+    class.define_method("get", method!(Metrics::get, 1))?;
+
+    Ok(())
+}