Skip to main content

feox_vector/
storage.rs

1use std::path::PathBuf;
2use std::sync::Arc;
3
4use feoxdb::{FeoxError, FeoxStore};
5
6#[derive(Debug, Clone)]
7pub struct StoreConfig {
8    pub device_path: Option<PathBuf>,
9    pub file_size: Option<u64>,
10    pub max_memory: Option<usize>,
11}
12
13impl Default for StoreConfig {
14    fn default() -> Self {
15        Self {
16            device_path: None,
17            file_size: None,
18            max_memory: Some(512 * 1024 * 1024),
19        }
20    }
21}
22
23#[derive(Debug, Clone)]
24pub struct KeyValue {
25    pub key: String,
26    pub value: Vec<u8>,
27}
28
29#[derive(Debug, Clone)]
30pub struct StoreStats {
31    pub records: usize,
32    pub memory_bytes: usize,
33}
34
35#[derive(Debug, thiserror::Error)]
36pub enum StoreError {
37    #[error("feox error: {0}")]
38    Feox(#[from] FeoxError),
39    #[error("stored key is not utf-8")]
40    InvalidKey,
41}
42
43pub type Result<T> = std::result::Result<T, StoreError>;
44
45#[derive(Clone)]
46pub struct Store {
47    inner: Arc<FeoxStore>,
48}
49
50impl Store {
51    pub fn open(config: StoreConfig) -> Result<Self> {
52        let mut builder = FeoxStore::builder();
53
54        if let Some(path) = config.device_path {
55            builder = builder.device_path(path.to_string_lossy().to_string());
56        }
57        if let Some(file_size) = config.file_size {
58            builder = builder.file_size(file_size);
59        }
60        if let Some(max_memory) = config.max_memory {
61            builder = builder.max_memory(max_memory);
62        } else {
63            builder = builder.no_memory_limit();
64        }
65
66        Ok(Self {
67            inner: Arc::new(builder.build()?),
68        })
69    }
70
71    pub fn put_bytes(&self, key: &str, value: &[u8]) -> Result<bool> {
72        Ok(self.inner.insert(key.as_bytes(), value)?)
73    }
74
75    pub fn get_bytes(&self, key: &str) -> Result<Option<Vec<u8>>> {
76        match self.inner.get(key.as_bytes()) {
77            Ok(value) => Ok(Some(value)),
78            Err(FeoxError::KeyNotFound) => Ok(None),
79            Err(error) => Err(StoreError::Feox(error)),
80        }
81    }
82
83    pub fn delete(&self, key: &str) -> Result<bool> {
84        match self.inner.delete(key.as_bytes()) {
85            Ok(()) => Ok(true),
86            Err(FeoxError::KeyNotFound) => Ok(false),
87            Err(error) => Err(StoreError::Feox(error)),
88        }
89    }
90
91    pub fn list_prefix(&self, prefix: &str, limit: usize) -> Result<Vec<KeyValue>> {
92        self.list_prefix_after(prefix, None, limit)
93    }
94
95    pub fn list_prefix_after(
96        &self,
97        prefix: &str,
98        start_after: Option<&str>,
99        limit: usize,
100    ) -> Result<Vec<KeyValue>> {
101        if limit == 0 {
102            return Ok(Vec::new());
103        }
104
105        let start = match start_after {
106            Some(key) => {
107                let mut value = key.as_bytes().to_vec();
108                value.push(0);
109                value
110            }
111            None => prefix.as_bytes().to_vec(),
112        };
113        let mut end = prefix.as_bytes().to_vec();
114        end.push(0xff);
115
116        let pairs = self
117            .inner
118            .range_query(&start, &end, limit)?
119            .into_iter()
120            .filter(|(key, _)| key.starts_with(prefix.as_bytes()))
121            .map(|(key, value)| {
122                let key = String::from_utf8(key).map_err(|_| StoreError::InvalidKey)?;
123                Ok(KeyValue { key, value })
124            })
125            .collect::<Result<Vec<_>>>()?;
126
127        Ok(pairs)
128    }
129
130    pub fn stats(&self) -> StoreStats {
131        StoreStats {
132            records: self.inner.len(),
133            memory_bytes: self.inner.memory_usage(),
134        }
135    }
136}