From be96f049c99e7f50bf9fbacc7943e1fbaa90e26e Mon Sep 17 00:00:00 2001
From: David Calavera <david.calavera@gmail.com>
Date: Fri, 8 May 2020 06:45:43 -0700
Subject: [PATCH] Expose `open_with_ttl`. (#425)

---
 src/db.rs        | 27 +++++++++++++++++++++++++++
 tests/test_db.rs | 17 +++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/src/db.rs b/src/db.rs
index 17759e5..1b69521 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -31,6 +31,7 @@ use std::path::PathBuf;
 use std::ptr;
 use std::slice;
 use std::str;
+use std::time::Duration;
 
 /// A RocksDB database.
 ///
@@ -55,6 +56,7 @@ enum AccessType<'a> {
     ReadWrite,
     ReadOnly { error_if_log_file_exist: bool },
     Secondary { secondary_path: &'a Path },
+    WithTTL { ttl: Duration },
 }
 
 impl DB {
@@ -88,6 +90,25 @@ impl DB {
         DB::open_cf_as_secondary(opts, primary_path, secondary_path, None::<&str>)
     }
 
+    /// Opens the database with a Time to Live compaction filter.
+    pub fn open_with_ttl<P: AsRef<Path>>(
+        opts: &Options,
+        path: P,
+        ttl: Duration,
+    ) -> Result<DB, Error> {
+        let c_path = to_cpath(&path)?;
+        let db = DB::open_raw(opts, &c_path, &AccessType::WithTTL { ttl })?;
+        if db.is_null() {
+            return Err(Error::new("Could not initialize database.".to_owned()));
+        }
+
+        Ok(DB {
+            inner: db,
+            cfs: BTreeMap::new(),
+            path: path.as_ref().to_path_buf(),
+        })
+    }
+
     /// Opens a database with the given database options and column family names.
     ///
     /// Column families opened using this function will be created with default `Options`.
@@ -275,6 +296,11 @@ impl DB {
                         to_cpath(secondary_path)?.as_ptr() as *const _,
                     ))
                 }
+                AccessType::WithTTL { ttl } => ffi_try!(ffi::rocksdb_open_with_ttl(
+                    opts.inner,
+                    cpath.as_ptr() as *const _,
+                    ttl.as_secs() as c_int,
+                )),
             }
         };
         Ok(db)
@@ -321,6 +347,7 @@ impl DB {
                         cfhandles.as_mut_ptr(),
                     ))
                 }
+                _ => return Err(Error::new("Unsupported access type".to_owned())),
             }
         };
         Ok(db)
diff --git a/tests/test_db.rs b/tests/test_db.rs
index 524d304..cae1b85 100644
--- a/tests/test_db.rs
+++ b/tests/test_db.rs
@@ -16,6 +16,7 @@ mod util;
 
 use rocksdb::{Error, IteratorMode, Options, Snapshot, WriteBatch, DB};
 use std::sync::Arc;
+use std::time::Duration;
 use std::{mem, thread};
 use util::DBPath;
 
@@ -357,3 +358,19 @@ fn test_open_as_secondary() {
     let result = secondary.get(b"key1").unwrap().unwrap();
     assert_eq!(get_byte_slice(&result), b"value2");
 }
+
+#[test]
+fn test_open_with_ttl() {
+    let path = DBPath::new("_rust_rocksdb_test_open_with_ttl");
+
+    let mut opts = Options::default();
+    opts.create_if_missing(true);
+    let db = DB::open_with_ttl(&opts, &path, Duration::from_secs(1)).unwrap();
+    db.put(b"key1", b"value1").unwrap();
+
+    thread::sleep(Duration::from_secs(2));
+    // Trigger a manual compaction, this will check the TTL filter
+    // in the database and drop all expired entries.
+    db.compact_range(None::<&[u8]>, None::<&[u8]>);
+    assert!(db.get(b"key1").unwrap().is_none());
+}