Skip to main content

lucistore/
blob_store.rs

1//! Generic blob storage trait for persisting data to external backends.
2//!
3//! This trait abstracts file-level storage so that indexes can be backed
4//! by a database, S3, or any other blob store — not just the local filesystem.
5//!
6//! Implementations:
7//! - [`MemBlobStore`]: in-memory store for testing
8//! - (external) `CypherBlobStore`: rag3db `_index_blobs` table
9//! - (external) `PostgresBlobStore`: Postgres bytea columns
10//! - (external) `S3BlobStore`: S3-compatible object storage
11
12use std::collections::HashMap;
13use std::io;
14use std::sync::{Arc, RwLock};
15
16/// Trait for blob storage backends.
17///
18/// Each blob is identified by `(index_name, file_name)`.
19/// - `index_name`: identifies which index (e.g. "Product", "Article_Index")
20/// - `file_name`: identifies which file within the index (e.g. "meta.json", "{uuid}.idx")
21pub trait BlobStore: Send + Sync + 'static {
22    /// Load a blob. Returns `NotFound` if it doesn't exist.
23    fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>>;
24
25    /// Save a blob (create or overwrite).
26    fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()>;
27
28    /// Delete a blob. Returns Ok(()) even if the blob didn't exist.
29    fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()>;
30
31    /// Check if a blob exists.
32    fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool>;
33
34    /// List all file names for a given index.
35    fn list(&self, index_name: &str) -> io::Result<Vec<String>>;
36
37    /// Size of a blob in bytes WITHOUT loading it, when the backend can
38    /// answer cheaply (`LENGTH(_data)` in SQL, HEAD on S3…). `Ok(None)` —
39    /// the default — means "unknown": lazy consumers then fall back to
40    /// materialising the file on first open instead of first byte read.
41    fn blob_len(&self, _index_name: &str, _file_name: &str) -> io::Result<Option<u64>> {
42        Ok(None)
43    }
44
45    /// Read a byte range of a blob without loading it whole, when the
46    /// backend can (`SUBSTRING(_data FROM … FOR …)` in SQL, ranged GET on
47    /// S3…). `Ok(None)` — the default — means "unsupported": lazy consumers
48    /// then materialise the whole blob instead. Out-of-bounds ranges are the
49    /// caller's bug; backends may truncate or error.
50    fn load_range(
51        &self,
52        _index_name: &str,
53        _file_name: &str,
54        _range: std::ops::Range<u64>,
55    ) -> io::Result<Option<Vec<u8>>> {
56        Ok(None)
57    }
58}
59
60/// Every method forwards, so `Arc<dyn BlobStore>` (and `Arc<Concrete>`) can
61/// be used directly wherever an `S: BlobStore` is expected — e.g.
62/// `BlobShardStorage<Arc<dyn BlobStore>>` — without a hand-written bridge.
63impl<T: BlobStore + ?Sized> BlobStore for std::sync::Arc<T> {
64    fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>> {
65        (**self).load(index_name, file_name)
66    }
67    fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()> {
68        (**self).save(index_name, file_name, data)
69    }
70    fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()> {
71        (**self).delete(index_name, file_name)
72    }
73    fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool> {
74        (**self).exists(index_name, file_name)
75    }
76    fn list(&self, index_name: &str) -> io::Result<Vec<String>> {
77        (**self).list(index_name)
78    }
79    fn blob_len(&self, index_name: &str, file_name: &str) -> io::Result<Option<u64>> {
80        (**self).blob_len(index_name, file_name)
81    }
82    fn load_range(
83        &self,
84        index_name: &str,
85        file_name: &str,
86        range: std::ops::Range<u64>,
87    ) -> io::Result<Option<Vec<u8>>> {
88        (**self).load_range(index_name, file_name, range)
89    }
90}
91
92/// In-memory blob store for testing.
93#[derive(Debug, Clone)]
94pub struct MemBlobStore {
95    /// `index_name -> file_name -> data`
96    data: Arc<RwLock<HashMap<String, HashMap<String, Vec<u8>>>>>,
97}
98
99impl MemBlobStore {
100    pub fn new() -> Self {
101        Self {
102            data: Arc::new(RwLock::new(HashMap::new())),
103        }
104    }
105}
106
107impl Default for MemBlobStore {
108    fn default() -> Self {
109        Self::new()
110    }
111}
112
113impl BlobStore for MemBlobStore {
114    fn blob_len(&self, index_name: &str, file_name: &str) -> io::Result<Option<u64>> {
115        Ok(self.data.read().unwrap()
116            .get(index_name)
117            .and_then(|files| files.get(file_name))
118            .map(|b| b.len() as u64))
119    }
120
121    fn load_range(
122        &self,
123        index_name: &str,
124        file_name: &str,
125        range: std::ops::Range<u64>,
126    ) -> io::Result<Option<Vec<u8>>> {
127        Ok(self.data.read().unwrap()
128            .get(index_name)
129            .and_then(|files| files.get(file_name))
130            .map(|b| b[range.start as usize..(range.end as usize).min(b.len())].to_vec()))
131    }
132
133    fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>> {
134        let guard = self.data.read().map_err(|_| {
135            io::Error::new(io::ErrorKind::Other, "lock poisoned")
136        })?;
137        guard
138            .get(index_name)
139            .and_then(|files| files.get(file_name))
140            .cloned()
141            .ok_or_else(|| {
142                io::Error::new(
143                    io::ErrorKind::NotFound,
144                    format!("{index_name}/{file_name} not found"),
145                )
146            })
147    }
148
149    fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()> {
150        let mut guard = self.data.write().map_err(|_| {
151            io::Error::new(io::ErrorKind::Other, "lock poisoned")
152        })?;
153        guard
154            .entry(index_name.to_string())
155            .or_default()
156            .insert(file_name.to_string(), data.to_vec());
157        Ok(())
158    }
159
160    fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()> {
161        let mut guard = self.data.write().map_err(|_| {
162            io::Error::new(io::ErrorKind::Other, "lock poisoned")
163        })?;
164        if let Some(files) = guard.get_mut(index_name) {
165            files.remove(file_name);
166        }
167        Ok(())
168    }
169
170    fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool> {
171        let guard = self.data.read().map_err(|_| {
172            io::Error::new(io::ErrorKind::Other, "lock poisoned")
173        })?;
174        Ok(guard
175            .get(index_name)
176            .map_or(false, |files| files.contains_key(file_name)))
177    }
178
179    fn list(&self, index_name: &str) -> io::Result<Vec<String>> {
180        let guard = self.data.read().map_err(|_| {
181            io::Error::new(io::ErrorKind::Other, "lock poisoned")
182        })?;
183        Ok(guard
184            .get(index_name)
185            .map(|files| files.keys().cloned().collect())
186            .unwrap_or_default())
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    #[test]
195    fn test_mem_blob_store_roundtrip() {
196        let store = MemBlobStore::new();
197        store.save("idx1", "file.bin", b"hello world").unwrap();
198        assert!(store.exists("idx1", "file.bin").unwrap());
199        assert!(!store.exists("idx1", "other.bin").unwrap());
200        let loaded = store.load("idx1", "file.bin").unwrap();
201        assert_eq!(loaded, b"hello world");
202        store.delete("idx1", "file.bin").unwrap();
203        assert!(!store.exists("idx1", "file.bin").unwrap());
204    }
205
206    #[test]
207    fn test_mem_blob_store_multiple_indexes() {
208        let store = MemBlobStore::new();
209        store.save("idx1", "a.bin", b"aaa").unwrap();
210        store.save("idx2", "a.bin", b"xxx").unwrap();
211        assert_eq!(store.load("idx1", "a.bin").unwrap(), b"aaa");
212        assert_eq!(store.load("idx2", "a.bin").unwrap(), b"xxx");
213    }
214}