use crate::{internal_error, Result};
use rocksdb::ops::{DropCF, GetColumnFamilys, GetPinnedCF, GetPropertyCF, OpenCF, PutCF};
use rocksdb::{
ColumnFamilyDescriptor, DBPinnableSlice, DBWithTTL as RawDBWithTTL, Options, TTLOpenDescriptor,
};
use std::path::Path;
const PROPERTY_NUM_KEYS: &str = "rocksdb.estimate-num-keys";
const DB_LOG_KEEP_NUM: usize = 10;
#[derive(Debug)]
pub struct DBWithTTL {
pub(crate) inner: RawDBWithTTL,
}
impl DBWithTTL {
pub fn open_cf<P, I, N>(path: P, cf_names: I, ttl: i32) -> Result<Self>
where
P: AsRef<Path>,
I: IntoIterator<Item = N>,
N: Into<String>,
{
let mut opts = Options::default();
opts.create_if_missing(true);
opts.create_missing_column_families(true);
opts.set_keep_log_file_num(DB_LOG_KEEP_NUM);
let cf_descriptors: Vec<_> = cf_names
.into_iter()
.map(|name| ColumnFamilyDescriptor::new(name, Options::default()))
.collect();
let descriptor = TTLOpenDescriptor::by_default(ttl);
let inner = RawDBWithTTL::open_cf_descriptors_with_descriptor(
&opts,
path,
cf_descriptors,
descriptor,
)
.map_err(|err| internal_error(format!("failed to open database: {err}")))?;
Ok(DBWithTTL { inner })
}
pub fn get_pinned(&self, col: &str, key: &[u8]) -> Result<Option<DBPinnableSlice>> {
let cf = self
.inner
.cf_handle(col)
.ok_or_else(|| internal_error(format!("column {col} not found")))?;
self.inner.get_pinned_cf(cf, key).map_err(internal_error)
}
pub fn put<K, V>(&self, col: &str, key: K, value: V) -> Result<()>
where
K: AsRef<[u8]>,
V: AsRef<[u8]>,
{
let cf = self
.inner
.cf_handle(col)
.ok_or_else(|| internal_error(format!("column {col} not found")))?;
self.inner.put_cf(cf, key, value).map_err(internal_error)
}
pub fn create_cf_with_ttl(&mut self, col: &str, ttl: i32) -> Result<()> {
let opts = Options::default();
self.inner
.create_cf_with_ttl(col, &opts, ttl)
.map_err(internal_error)
}
pub fn drop_cf(&mut self, col: &str) -> Result<()> {
self.inner.drop_cf(col).map_err(internal_error)
}
pub fn estimate_num_keys_cf(&self, col: &str) -> Result<Option<u64>> {
let cf = self
.inner
.cf_handle(col)
.ok_or_else(|| internal_error(format!("column {col} not found")))?;
self.inner
.property_int_value_cf(cf, PROPERTY_NUM_KEYS)
.map_err(internal_error)
}
}