Skip to main content

dk_engine/
storage.rs

1//! Object storage abstraction.
2
3use async_trait::async_trait;
4
5/// Trait for object storage backends.
6#[async_trait]
7pub trait ObjectStore: Send + Sync + 'static {
8    async fn get(&self, key: &str) -> anyhow::Result<Vec<u8>>;
9    async fn put(&self, key: &str, data: Vec<u8>) -> anyhow::Result<()>;
10    async fn delete(&self, key: &str) -> anyhow::Result<()>;
11    async fn list(&self, prefix: &str) -> anyhow::Result<Vec<String>>;
12    async fn exists(&self, key: &str) -> anyhow::Result<bool>;
13}
14
15/// Local filesystem object store.
16pub struct LocalStore {
17    root: std::path::PathBuf,
18}
19
20impl LocalStore {
21    pub fn new(root: std::path::PathBuf) -> Self {
22        Self { root }
23    }
24}
25
26#[async_trait]
27impl ObjectStore for LocalStore {
28    async fn get(&self, key: &str) -> anyhow::Result<Vec<u8>> {
29        let path = self.root.join(key);
30        Ok(tokio::fs::read(path).await?)
31    }
32
33    async fn put(&self, key: &str, data: Vec<u8>) -> anyhow::Result<()> {
34        let path = self.root.join(key);
35        if let Some(parent) = path.parent() {
36            tokio::fs::create_dir_all(parent).await?;
37        }
38        Ok(tokio::fs::write(path, data).await?)
39    }
40
41    async fn delete(&self, key: &str) -> anyhow::Result<()> {
42        let path = self.root.join(key);
43        match tokio::fs::remove_file(path).await {
44            Ok(()) => Ok(()),
45            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
46            Err(e) => Err(e.into()),
47        }
48    }
49
50    async fn list(&self, prefix: &str) -> anyhow::Result<Vec<String>> {
51        let dir = self.root.join(prefix);
52        let mut entries = Vec::new();
53        if dir.exists() {
54            let mut read_dir = tokio::fs::read_dir(&dir).await?;
55            while let Some(entry) = read_dir.next_entry().await? {
56                if let Some(name) = entry.file_name().to_str() {
57                    let key = if prefix.is_empty() {
58                        name.to_string()
59                    } else {
60                        format!("{prefix}/{name}")
61                    };
62                    entries.push(key);
63                }
64            }
65        }
66        Ok(entries)
67    }
68
69    async fn exists(&self, key: &str) -> anyhow::Result<bool> {
70        let path = self.root.join(key);
71        Ok(path.exists())
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[tokio::test]
80    async fn local_store_roundtrip() {
81        let dir = tempfile::tempdir().unwrap();
82        let store = LocalStore::new(dir.path().to_path_buf());
83
84        store.put("test/file.txt", b"hello".to_vec()).await.unwrap();
85        assert!(store.exists("test/file.txt").await.unwrap());
86
87        let data = store.get("test/file.txt").await.unwrap();
88        assert_eq!(data, b"hello");
89
90        let keys = store.list("test").await.unwrap();
91        assert_eq!(keys, vec!["test/file.txt"]);
92
93        store.delete("test/file.txt").await.unwrap();
94        assert!(!store.exists("test/file.txt").await.unwrap());
95    }
96
97    #[tokio::test]
98    async fn local_store_get_not_found() {
99        let dir = tempfile::tempdir().unwrap();
100        let store = LocalStore::new(dir.path().to_path_buf());
101        let result = store.get("nonexistent").await;
102        assert!(result.is_err());
103    }
104
105    #[tokio::test]
106    async fn local_store_delete_idempotent() {
107        let dir = tempfile::tempdir().unwrap();
108        let store = LocalStore::new(dir.path().to_path_buf());
109        store.delete("nonexistent").await.unwrap();
110    }
111}