diskann_providers/storage/
file_storage_provider.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5use std::{
6    fs::{self, File, OpenOptions},
7    io::{BufReader, BufWriter, Result},
8};
9
10use super::{StorageReadProvider, StorageWriteProvider};
11
12/// FileStorage implements both StorageReadProvider and StorageWriteProvider.
13#[derive(Default)]
14pub struct FileStorageProvider;
15
16impl StorageReadProvider for FileStorageProvider {
17    type Reader = BufReader<File>;
18
19    fn open_reader(&self, item_identifier: &str) -> Result<Self::Reader> {
20        let f = File::open(item_identifier)?;
21        Ok(BufReader::new(f))
22    }
23
24    fn get_length(&self, item_identifier: &str) -> Result<u64> {
25        let metadata = fs::metadata(item_identifier)?;
26        Ok(metadata.len())
27    }
28
29    fn exists(&self, item_identifier: &str) -> bool {
30        fs::metadata(item_identifier).is_ok()
31    }
32}
33
34impl StorageWriteProvider for FileStorageProvider {
35    type Writer = BufWriter<File>;
36
37    fn open_writer(&self, item_identifier: &str) -> Result<Self::Writer> {
38        let f = OpenOptions::new().write(true).open(item_identifier)?;
39        Ok(BufWriter::new(f))
40    }
41
42    fn create_for_write(&self, item_identifier: &str) -> Result<Self::Writer> {
43        let f = OpenOptions::new()
44            .write(true)
45            .create(true)
46            .truncate(true)
47            .open(item_identifier)?;
48        Ok(BufWriter::new(f))
49    }
50
51    fn delete(&self, item_identifier: &str) -> Result<()> {
52        fs::remove_file(item_identifier)
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use std::io::{Read, Seek, SeekFrom, Write};
59
60    use tempfile::TempDir;
61
62    use super::*;
63
64    #[test]
65    fn test_file_reader() {
66        // Use TempDir for automatic deleting when going out of scope
67        let tmp_dir =
68            TempDir::with_prefix("test_file_reader").expect("Failed to create temporary directory");
69        let file_path = tmp_dir.path().join("test_file_reader.txt");
70        let file_name = file_path.to_str().unwrap();
71
72        let mut file = File::create(file_name).unwrap();
73        file.write_all(b"Hello, world!").unwrap();
74
75        let mut reader = FileStorageProvider.open_reader(file_name).unwrap();
76        let mut buffer = [0; 5];
77
78        reader.seek(SeekFrom::Start(0)).unwrap();
79        reader.read(&mut buffer).unwrap();
80        assert_eq!(&buffer, b"Hello");
81
82        reader.seek(SeekFrom::Start(5)).unwrap();
83        reader.read(&mut buffer).unwrap();
84        assert_eq!(&buffer, b", wor");
85    }
86
87    #[test]
88    fn test_file_create_write() {
89        let storage_provider = FileStorageProvider;
90
91        // Use TempDir for automatic deleting when going out of scope
92        let tmp_dir = TempDir::with_prefix("test_file_create_write")
93            .expect("Failed to create temporary directory");
94        let file_path = tmp_dir.path().join("test_file_create_write.txt");
95        let file_name = file_path.to_str().unwrap();
96
97        assert!(!storage_provider.exists(file_name));
98        {
99            let mut file = storage_provider.create_for_write(file_name).unwrap();
100            file.write_all(b"Hello, world! ").unwrap();
101            file.flush().unwrap();
102        }
103
104        assert!(storage_provider.exists(file_name));
105
106        {
107            let mut file = storage_provider.open_writer(file_name).unwrap();
108            file.seek(SeekFrom::End(0)).unwrap();
109            file.write_all(b"This is the second write! ").unwrap();
110            file.write_all(b"This is the third write!").unwrap();
111            file.flush().unwrap();
112        }
113
114        let expected = b"Hello, world! This is the second write! This is the third write!";
115
116        let mut reader = FileStorageProvider.open_reader(file_name).unwrap();
117        let mut file_data: Vec<u8> = Vec::new();
118        let read_size = reader.read_to_end(&mut file_data).unwrap();
119
120        assert_eq!(
121            expected.len(),
122            read_size,
123            "Did not read the expected number of bytes"
124        );
125        assert_eq!(expected, file_data.as_slice());
126    }
127
128    #[test]
129    fn test_file_storage_exists() {
130        // Use TempDir for automatic deleting when going out of scope
131        let tmp_dir = TempDir::with_prefix("test_file_storage_exists")
132            .expect("Failed to create temporary directory");
133        let file_path = tmp_dir.path().join("test_file_storage_exists.txt");
134        let file_name = file_path.to_str().unwrap();
135
136        assert!(!FileStorageProvider.exists(file_name));
137        File::create(file_name).unwrap();
138        assert!(FileStorageProvider.exists(file_name));
139    }
140
141    #[test]
142    fn test_file_storage_get_length() {
143        // Use TempDir for automatic deleting when going out of scope
144        let tmp_dir = TempDir::with_prefix("test_file_storage_get_length")
145            .expect("Failed to create temporary directory");
146        let file_path = tmp_dir.path().join("test_file_storage_get_length.txt");
147        let file_name = file_path.to_str().unwrap();
148
149        let mut file = File::create(file_name).unwrap();
150        file.write_all(b"Hello, world!").unwrap();
151
152        assert_eq!(FileStorageProvider.get_length(file_name).unwrap(), 13);
153        fs::remove_file(file_name).unwrap();
154    }
155}
diskann_providers/storage/file_storage_provider.rs

diskann_providers/storage/
file_storage_provider.rs
