kiromi-ai-memory 0.2.2

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! Local-filesystem storage backend.

use std::path::{Path, PathBuf};

use async_trait::async_trait;
use bytes::Bytes;
use tokio::fs;
use tokio::io::AsyncWriteExt;

use crate::error::{Error, Result};
use crate::storage::{Storage, StorageEntry, StorageKey};

/// Local-FS storage backend rooted at a directory. Keys map to relative paths
/// underneath the root.
#[derive(Debug, Clone)]
pub struct LocalFsBackend {
    root: PathBuf,
}

impl LocalFsBackend {
    /// Construct rooted at `root`. Creates the directory if missing.
    pub async fn new(root: impl Into<PathBuf>) -> Result<Self> {
        let root = root.into();
        fs::create_dir_all(&root).await.map_err(Error::Io)?;
        Ok(Self { root })
    }

    fn path(&self, key: &StorageKey) -> PathBuf {
        self.root.join(key.as_str())
    }

    fn rel_key(&self, abs: &Path) -> StorageKey {
        let rel = abs.strip_prefix(&self.root).unwrap_or(abs);
        StorageKey::new(rel.to_string_lossy().replace('\\', "/"))
    }
}

#[async_trait]
impl Storage for LocalFsBackend {
    fn id(&self) -> String {
        format!("local-fs:{}", self.root.display())
    }

    async fn put(&self, key: &StorageKey, body: Bytes) -> Result<()> {
        let path = self.path(key);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).await.map_err(Error::Io)?;
        }
        let tmp = path.with_extension("tmp");
        let mut f = fs::File::create(&tmp).await.map_err(Error::Io)?;
        f.write_all(&body).await.map_err(Error::Io)?;
        f.sync_data().await.map_err(Error::Io)?;
        drop(f);
        fs::rename(tmp, path).await.map_err(Error::Io)?;
        Ok(())
    }

    async fn get(&self, key: &StorageKey) -> Result<Bytes> {
        let bytes = fs::read(self.path(key)).await.map_err(Error::Io)?;
        Ok(Bytes::from(bytes))
    }

    async fn exists(&self, key: &StorageKey) -> Result<bool> {
        match fs::metadata(self.path(key)).await {
            Ok(_) => Ok(true),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
            Err(e) => Err(Error::Io(e)),
        }
    }

    async fn delete(&self, key: &StorageKey) -> Result<()> {
        match fs::remove_file(self.path(key)).await {
            Ok(()) => Ok(()),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
            Err(e) => Err(Error::Io(e)),
        }
    }

    fn root_hint(&self) -> Option<std::path::PathBuf> {
        Some(self.root.clone())
    }

    async fn rename_within(&self, src: &StorageKey, dst: &StorageKey) -> Result<()> {
        if src == dst {
            return Ok(());
        }
        let src_path = self.path(src);
        match fs::metadata(&src_path).await {
            Ok(_) => {}
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
            Err(e) => return Err(Error::Io(e)),
        }
        let dst_path = self.path(dst);
        if let Some(parent) = dst_path.parent() {
            fs::create_dir_all(parent).await.map_err(Error::Io)?;
        }
        fs::rename(&src_path, &dst_path).await.map_err(Error::Io)?;
        Ok(())
    }

    async fn list_prefix(&self, prefix: &StorageKey) -> Result<Vec<StorageEntry>> {
        let abs_prefix = self.path(prefix);
        let mut out = Vec::new();
        let mut stack = vec![abs_prefix.clone()];
        while let Some(dir) = stack.pop() {
            let mut rd = match fs::read_dir(&dir).await {
                Ok(rd) => rd,
                Err(e) if e.kind() == std::io::ErrorKind::NotFound => continue,
                Err(e) => return Err(Error::Io(e)),
            };
            while let Some(ent) = rd.next_entry().await.map_err(Error::Io)? {
                let path = ent.path();
                let ft = ent.file_type().await.map_err(Error::Io)?;
                if ft.is_dir() {
                    stack.push(path);
                } else if ft.is_file() {
                    let md = ent.metadata().await.map_err(Error::Io)?;
                    out.push(StorageEntry {
                        key: self.rel_key(&path),
                        size: md.len(),
                    });
                }
            }
        }
        out.sort_by(|a, b| a.key.cmp(&b.key));
        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[tokio::test]
    async fn put_get_exists_delete() {
        let td = TempDir::new().unwrap();
        let s = LocalFsBackend::new(td.path()).await.unwrap();
        let k = StorageKey::new("a/b.md");
        assert!(!s.exists(&k).await.unwrap());
        s.put(&k, Bytes::from_static(b"hi")).await.unwrap();
        assert!(s.exists(&k).await.unwrap());
        assert_eq!(s.get(&k).await.unwrap(), Bytes::from_static(b"hi"));
        s.delete(&k).await.unwrap();
        assert!(!s.exists(&k).await.unwrap());
    }

    #[tokio::test]
    async fn list_prefix_returns_relative_keys() {
        let td = TempDir::new().unwrap();
        let s = LocalFsBackend::new(td.path()).await.unwrap();
        s.put(&StorageKey::new("a/x"), Bytes::from_static(b"1"))
            .await
            .unwrap();
        s.put(&StorageKey::new("a/y"), Bytes::from_static(b"22"))
            .await
            .unwrap();
        s.put(&StorageKey::new("b/z"), Bytes::from_static(b"333"))
            .await
            .unwrap();
        let got: Vec<_> = s
            .list_prefix(&StorageKey::new("a"))
            .await
            .unwrap()
            .into_iter()
            .map(|e| e.key.as_str().to_string())
            .collect();
        assert_eq!(got, vec!["a/x".to_string(), "a/y".to_string()]);
    }

    #[tokio::test]
    async fn delete_missing_is_idempotent() {
        let td = TempDir::new().unwrap();
        let s = LocalFsBackend::new(td.path()).await.unwrap();
        s.delete(&StorageKey::new("nope")).await.unwrap();
    }
}