ruvector_snapshot/
manager.rs

1use crate::error::{Result, SnapshotError};
2use crate::snapshot::{Snapshot, SnapshotData};
3use crate::storage::SnapshotStorage;
4
5/// Manages snapshot operations for collections
6pub struct SnapshotManager {
7    storage: Box<dyn SnapshotStorage>,
8}
9
10impl SnapshotManager {
11    /// Create a new snapshot manager with the given storage backend
12    pub fn new(storage: Box<dyn SnapshotStorage>) -> Self {
13        Self { storage }
14    }
15
16    /// Create a snapshot of a collection
17    ///
18    /// # Arguments
19    /// * `snapshot_data` - The complete snapshot data including vectors and configuration
20    ///
21    /// # Returns
22    /// * `Snapshot` - Metadata about the created snapshot
23    pub async fn create_snapshot(&self, snapshot_data: SnapshotData) -> Result<Snapshot> {
24        // Validate snapshot data
25        if snapshot_data.vectors.is_empty() {
26            return Err(SnapshotError::storage(
27                "Cannot create snapshot of empty collection",
28            ));
29        }
30
31        // Verify all vectors have the same dimension
32        let expected_dim = snapshot_data.config.dimension;
33        for (idx, vector) in snapshot_data.vectors.iter().enumerate() {
34            if vector.vector.len() != expected_dim {
35                return Err(SnapshotError::storage(format!(
36                    "Vector {} has dimension {} but expected {}",
37                    idx,
38                    vector.vector.len(),
39                    expected_dim
40                )));
41            }
42        }
43
44        // Save the snapshot
45        self.storage.save(&snapshot_data).await
46    }
47
48    /// Restore a snapshot by ID
49    ///
50    /// # Arguments
51    /// * `id` - The unique snapshot identifier
52    ///
53    /// # Returns
54    /// * `SnapshotData` - The complete snapshot data including vectors and configuration
55    pub async fn restore_snapshot(&self, id: &str) -> Result<SnapshotData> {
56        if id.is_empty() {
57            return Err(SnapshotError::storage("Snapshot ID cannot be empty"));
58        }
59
60        self.storage.load(id).await
61    }
62
63    /// List all available snapshots
64    ///
65    /// # Returns
66    /// * `Vec<Snapshot>` - List of all snapshot metadata, sorted by creation date (newest first)
67    pub async fn list_snapshots(&self) -> Result<Vec<Snapshot>> {
68        self.storage.list().await
69    }
70
71    /// List snapshots for a specific collection
72    ///
73    /// # Arguments
74    /// * `collection_name` - Name of the collection to filter by
75    ///
76    /// # Returns
77    /// * `Vec<Snapshot>` - List of snapshots for the specified collection
78    pub async fn list_snapshots_for_collection(
79        &self,
80        collection_name: &str,
81    ) -> Result<Vec<Snapshot>> {
82        let all_snapshots = self.storage.list().await?;
83        Ok(all_snapshots
84            .into_iter()
85            .filter(|s| s.collection_name == collection_name)
86            .collect())
87    }
88
89    /// Delete a snapshot by ID
90    ///
91    /// # Arguments
92    /// * `id` - The unique snapshot identifier
93    pub async fn delete_snapshot(&self, id: &str) -> Result<()> {
94        if id.is_empty() {
95            return Err(SnapshotError::storage("Snapshot ID cannot be empty"));
96        }
97
98        self.storage.delete(id).await
99    }
100
101    /// Get snapshot metadata by ID
102    ///
103    /// # Arguments
104    /// * `id` - The unique snapshot identifier
105    ///
106    /// # Returns
107    /// * `Snapshot` - Metadata about the snapshot
108    pub async fn get_snapshot_info(&self, id: &str) -> Result<Snapshot> {
109        let snapshots = self.storage.list().await?;
110        snapshots
111            .into_iter()
112            .find(|s| s.id == id)
113            .ok_or_else(|| SnapshotError::SnapshotNotFound(id.to_string()))
114    }
115
116    /// Delete old snapshots, keeping only the N most recent
117    ///
118    /// # Arguments
119    /// * `collection_name` - Name of the collection
120    /// * `keep_count` - Number of recent snapshots to keep
121    ///
122    /// # Returns
123    /// * `usize` - Number of snapshots deleted
124    pub async fn cleanup_old_snapshots(
125        &self,
126        collection_name: &str,
127        keep_count: usize,
128    ) -> Result<usize> {
129        let snapshots = self.list_snapshots_for_collection(collection_name).await?;
130
131        if snapshots.len() <= keep_count {
132            return Ok(0);
133        }
134
135        let to_delete = &snapshots[keep_count..];
136        let mut deleted = 0;
137
138        for snapshot in to_delete {
139            if self.storage.delete(&snapshot.id).await.is_ok() {
140                deleted += 1;
141            }
142        }
143
144        Ok(deleted)
145    }
146
147    /// Get the total size of all snapshots in bytes
148    pub async fn total_size(&self) -> Result<u64> {
149        let snapshots = self.storage.list().await?;
150        Ok(snapshots.iter().map(|s| s.size_bytes).sum())
151    }
152
153    /// Get the total size of snapshots for a specific collection
154    pub async fn collection_size(&self, collection_name: &str) -> Result<u64> {
155        let snapshots = self.list_snapshots_for_collection(collection_name).await?;
156        Ok(snapshots.iter().map(|s| s.size_bytes).sum())
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163    use crate::snapshot::{CollectionConfig, DistanceMetric, VectorRecord};
164    use crate::storage::LocalStorage;
165    use std::path::PathBuf;
166
167    fn create_test_snapshot_data(name: &str, vector_count: usize) -> SnapshotData {
168        let config = CollectionConfig {
169            dimension: 3,
170            metric: DistanceMetric::Cosine,
171            hnsw_config: None,
172        };
173
174        let vectors = (0..vector_count)
175            .map(|i| {
176                VectorRecord::new(
177                    format!("v{}", i),
178                    vec![i as f32, (i + 1) as f32, (i + 2) as f32],
179                    None,
180                )
181            })
182            .collect();
183
184        SnapshotData::new(name.to_string(), config, vectors)
185    }
186
187    #[tokio::test]
188    async fn test_create_and_restore_snapshot() {
189        let temp_dir = std::env::temp_dir().join("ruvector-manager-test");
190        let storage = Box::new(LocalStorage::new(temp_dir.clone()));
191        let manager = SnapshotManager::new(storage);
192
193        let snapshot_data = create_test_snapshot_data("test-collection", 5);
194        let id = snapshot_data.id().to_string();
195
196        // Create snapshot
197        let snapshot = manager.create_snapshot(snapshot_data).await.unwrap();
198        assert_eq!(snapshot.id, id);
199        assert_eq!(snapshot.vectors_count, 5);
200
201        // Restore snapshot
202        let restored = manager.restore_snapshot(&id).await.unwrap();
203        assert_eq!(restored.id(), id);
204        assert_eq!(restored.vectors_count(), 5);
205
206        // Cleanup
207        let _ = manager.delete_snapshot(&id).await;
208        let _ = std::fs::remove_dir_all(temp_dir);
209    }
210
211    #[tokio::test]
212    async fn test_list_snapshots() {
213        let temp_dir = std::env::temp_dir().join("ruvector-list-test");
214        let storage = Box::new(LocalStorage::new(temp_dir.clone()));
215        let manager = SnapshotManager::new(storage);
216
217        // Create multiple snapshots
218        let snapshot1 = create_test_snapshot_data("collection-1", 3);
219        let snapshot2 = create_test_snapshot_data("collection-2", 5);
220
221        let id1 = snapshot1.id().to_string();
222        let id2 = snapshot2.id().to_string();
223
224        manager.create_snapshot(snapshot1).await.unwrap();
225        manager.create_snapshot(snapshot2).await.unwrap();
226
227        // List all
228        let all_snapshots = manager.list_snapshots().await.unwrap();
229        assert!(all_snapshots.len() >= 2);
230
231        // List by collection
232        let collection1_snapshots = manager
233            .list_snapshots_for_collection("collection-1")
234            .await
235            .unwrap();
236        assert_eq!(collection1_snapshots.len(), 1);
237
238        // Cleanup
239        let _ = manager.delete_snapshot(&id1).await;
240        let _ = manager.delete_snapshot(&id2).await;
241        let _ = std::fs::remove_dir_all(temp_dir);
242    }
243
244    #[tokio::test]
245    async fn test_cleanup_old_snapshots() {
246        let temp_dir = std::env::temp_dir().join("ruvector-cleanup-test");
247        let storage = Box::new(LocalStorage::new(temp_dir.clone()));
248        let manager = SnapshotManager::new(storage);
249
250        // Create multiple snapshots for the same collection
251        for i in 0..5 {
252            let snapshot_data = create_test_snapshot_data("test-collection", i + 1);
253            manager.create_snapshot(snapshot_data).await.unwrap();
254            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
255        }
256
257        // Cleanup, keeping only 2 most recent
258        let deleted = manager
259            .cleanup_old_snapshots("test-collection", 2)
260            .await
261            .unwrap();
262        assert_eq!(deleted, 3);
263
264        // Verify only 2 remain
265        let remaining = manager
266            .list_snapshots_for_collection("test-collection")
267            .await
268            .unwrap();
269        assert_eq!(remaining.len(), 2);
270
271        // Cleanup
272        let _ = std::fs::remove_dir_all(temp_dir);
273    }
274
275    #[tokio::test]
276    async fn test_snapshot_validation() {
277        let temp_dir = std::env::temp_dir().join("ruvector-validation-test");
278        let storage = Box::new(LocalStorage::new(temp_dir.clone()));
279        let manager = SnapshotManager::new(storage);
280
281        // Test empty collection
282        let config = CollectionConfig {
283            dimension: 3,
284            metric: DistanceMetric::Cosine,
285            hnsw_config: None,
286        };
287        let empty_data = SnapshotData::new("empty".to_string(), config, vec![]);
288        let result = manager.create_snapshot(empty_data).await;
289        assert!(result.is_err());
290
291        // Cleanup
292        let _ = std::fs::remove_dir_all(temp_dir);
293    }
294}
ruvector_snapshot/manager.rs

ruvector_snapshot/
manager.rs
