Skip to main content

smugmug_cli/cache/
hash_store.rs

1use anyhow::{Context, Result};
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use std::sync::Arc;
5
6#[derive(Debug, Serialize, Deserialize, Clone)]
7pub struct UploadedFile {
8    pub smugmug_uri: String,
9    pub album_key: String,
10    pub image_key: String,
11    pub uploaded_at: DateTime<Utc>,
12    pub file_size: u64,
13    pub original_path: String,
14}
15
16/// Thread-safe hash store using sled embedded database
17pub struct HashStore {
18    db: Arc<sled::Db>,
19}
20
21impl HashStore {
22    /// Create a new HashStore at the specified path
23    pub fn new(path: &str) -> Result<Self> {
24        let db = sled::open(path).context("Failed to open sled database")?;
25
26        Ok(HashStore { db: Arc::new(db) })
27    }
28
29    /// Get an uploaded file by its hash
30    pub fn get(&self, hash: &str) -> Result<Option<UploadedFile>> {
31        let value = self
32            .db
33            .get(hash.as_bytes())
34            .context("Failed to read from database")?;
35
36        match value {
37            Some(bytes) => {
38                let file: UploadedFile =
39                    serde_json::from_slice(&bytes).context("Failed to deserialize UploadedFile")?;
40                Ok(Some(file))
41            }
42            None => Ok(None),
43        }
44    }
45
46    /// Insert a new uploaded file record
47    pub fn insert(&self, hash: &str, file: UploadedFile) -> Result<()> {
48        let serialized = serde_json::to_vec(&file).context("Failed to serialize UploadedFile")?;
49
50        self.db
51            .insert(hash.as_bytes(), serialized)
52            .context("Failed to insert into database")?;
53
54        self.db.flush().context("Failed to flush database")?;
55
56        Ok(())
57    }
58
59    /// Clear all entries from the cache
60    pub fn clear(&self) -> Result<()> {
61        self.db.clear().context("Failed to clear database")?;
62
63        self.db.flush().context("Failed to flush database")?;
64
65        Ok(())
66    }
67
68    /// Get cache statistics
69    pub fn stats(&self) -> Result<CacheStats> {
70        let mut total_entries = 0;
71        let mut total_size = 0u64;
72        let mut oldest: Option<DateTime<Utc>> = None;
73        let mut newest: Option<DateTime<Utc>> = None;
74
75        for item in self.db.iter() {
76            let (_key, value) = item.context("Failed to read database entry")?;
77            total_entries += 1;
78
79            // Deserialize to get file size and timestamp
80            if let Ok(file) = serde_json::from_slice::<UploadedFile>(&value) {
81                total_size += file.file_size;
82                let timestamp = file.uploaded_at;
83
84                oldest = Some(oldest.map_or(timestamp, |old| old.min(timestamp)));
85                newest = Some(newest.map_or(timestamp, |new| new.max(timestamp)));
86            }
87        }
88
89        Ok(CacheStats {
90            total_entries,
91            total_size,
92            oldest_entry: oldest,
93            newest_entry: newest,
94        })
95    }
96
97    /// Get the total number of entries in the cache
98    #[allow(dead_code)]
99    pub fn count(&self) -> Result<usize> {
100        let mut count = 0;
101        for item in self.db.iter() {
102            item.context("Failed to read database entry")?;
103            count += 1;
104        }
105        Ok(count)
106    }
107
108    /// Get the total size of cached data in bytes (sum of file_size fields)
109    #[allow(dead_code)]
110    pub fn size(&self) -> Result<u64> {
111        let mut total_size = 0u64;
112        for item in self.db.iter() {
113            let (_key, value) = item.context("Failed to read database entry")?;
114            if let Ok(file) = serde_json::from_slice::<UploadedFile>(&value) {
115                total_size += file.file_size;
116            }
117        }
118        Ok(total_size)
119    }
120}
121
122// Implement Clone for HashStore since it uses Arc internally
123impl Clone for HashStore {
124    fn clone(&self) -> Self {
125        HashStore {
126            db: Arc::clone(&self.db),
127        }
128    }
129}
130
131#[derive(Debug, Clone)]
132pub struct CacheStats {
133    pub total_entries: usize,
134    pub total_size: u64,
135    pub oldest_entry: Option<DateTime<Utc>>,
136    pub newest_entry: Option<DateTime<Utc>>,
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use chrono::TimeZone;
143
144    fn temp_dir() -> String {
145        let dir = tempfile::tempdir().unwrap();
146        dir.path().to_str().unwrap().to_string()
147    }
148
149    fn create_test_file(timestamp: DateTime<Utc>) -> UploadedFile {
150        UploadedFile {
151            smugmug_uri: "https://api.smugmug.com/api/v2/image/test".to_string(),
152            album_key: "album123".to_string(),
153            image_key: "img456".to_string(),
154            uploaded_at: timestamp,
155            file_size: 1024,
156            original_path: "/test/path.jpg".to_string(),
157        }
158    }
159
160    #[test]
161    fn test_new_store() {
162        let dir = temp_dir();
163        let store = HashStore::new(&dir);
164        assert!(store.is_ok());
165    }
166
167    #[test]
168    fn test_insert_and_get() {
169        let dir = temp_dir();
170        let store = HashStore::new(&dir).unwrap();
171
172        let file = create_test_file(Utc::now());
173        let hash = "abc123";
174
175        // Insert
176        store.insert(hash, file.clone()).unwrap();
177
178        // Retrieve
179        let retrieved = store.get(hash).unwrap();
180        assert!(retrieved.is_some());
181        let retrieved = retrieved.unwrap();
182        assert_eq!(retrieved.smugmug_uri, file.smugmug_uri);
183        assert_eq!(retrieved.album_key, file.album_key);
184        assert_eq!(retrieved.image_key, file.image_key);
185    }
186
187    #[test]
188    fn test_get_nonexistent() {
189        let dir = temp_dir();
190        let store = HashStore::new(&dir).unwrap();
191
192        let result = store.get("nonexistent").unwrap();
193        assert!(result.is_none());
194    }
195
196    #[test]
197    fn test_clear() {
198        let dir = temp_dir();
199        let store = HashStore::new(&dir).unwrap();
200
201        // Insert multiple entries
202        for i in 0..5 {
203            let hash = format!("hash{}", i);
204            let file = create_test_file(Utc::now());
205            store.insert(&hash, file).unwrap();
206        }
207
208        // Verify entries exist
209        let count_before = store.count().unwrap();
210        assert_eq!(count_before, 5);
211
212        // Clear
213        store.clear().unwrap();
214
215        // Verify cleared
216        let count_after = store.count().unwrap();
217        assert_eq!(count_after, 0);
218    }
219
220    #[test]
221    fn test_count() {
222        let dir = temp_dir();
223        let store = HashStore::new(&dir).unwrap();
224
225        // Empty store
226        assert_eq!(store.count().unwrap(), 0);
227
228        // Add entries
229        for i in 0..10 {
230            let hash = format!("hash{}", i);
231            let file = create_test_file(Utc::now());
232            store.insert(&hash, file).unwrap();
233        }
234
235        assert_eq!(store.count().unwrap(), 10);
236    }
237
238    #[test]
239    fn test_size() {
240        let dir = temp_dir();
241        let store = HashStore::new(&dir).unwrap();
242
243        // Empty store
244        assert_eq!(store.size().unwrap(), 0);
245
246        // Add entries
247        let file = create_test_file(Utc::now());
248        store.insert("hash1", file).unwrap();
249
250        // Size should be greater than 0
251        let size = store.size().unwrap();
252        assert!(size > 0);
253    }
254
255    #[test]
256    fn test_stats_empty() {
257        let dir = temp_dir();
258        let store = HashStore::new(&dir).unwrap();
259
260        let stats = store.stats().unwrap();
261        assert_eq!(stats.total_entries, 0);
262        assert_eq!(stats.total_size, 0);
263        assert!(stats.oldest_entry.is_none());
264        assert!(stats.newest_entry.is_none());
265    }
266
267    #[test]
268    fn test_stats_single_entry() {
269        let dir = temp_dir();
270        let store = HashStore::new(&dir).unwrap();
271
272        let timestamp = Utc.with_ymd_and_hms(2024, 3, 15, 12, 0, 0).unwrap();
273        let file = create_test_file(timestamp);
274        store.insert("hash1", file).unwrap();
275
276        let stats = store.stats().unwrap();
277        assert_eq!(stats.total_entries, 1);
278        assert!(stats.total_size > 0);
279        assert_eq!(stats.oldest_entry, Some(timestamp));
280        assert_eq!(stats.newest_entry, Some(timestamp));
281    }
282
283    #[test]
284    fn test_stats_multiple_entries() {
285        let dir = temp_dir();
286        let store = HashStore::new(&dir).unwrap();
287
288        // Create entries with different timestamps
289        let oldest = Utc.with_ymd_and_hms(2024, 1, 15, 12, 0, 0).unwrap();
290        let middle = Utc.with_ymd_and_hms(2024, 2, 15, 12, 0, 0).unwrap();
291        let newest = Utc.with_ymd_and_hms(2024, 3, 15, 12, 0, 0).unwrap();
292
293        store.insert("hash1", create_test_file(middle)).unwrap();
294        store.insert("hash2", create_test_file(oldest)).unwrap();
295        store.insert("hash3", create_test_file(newest)).unwrap();
296
297        let stats = store.stats().unwrap();
298        assert_eq!(stats.total_entries, 3);
299        assert!(stats.total_size > 0);
300        assert_eq!(stats.oldest_entry, Some(oldest));
301        assert_eq!(stats.newest_entry, Some(newest));
302    }
303
304    #[test]
305    fn test_stats_size_calculation() {
306        let dir = temp_dir();
307        let store = HashStore::new(&dir).unwrap();
308
309        // Insert one entry
310        let file1 = create_test_file(Utc::now());
311        store.insert("hash1", file1).unwrap();
312        let size1 = store.stats().unwrap().total_size;
313
314        // Insert another entry
315        let file2 = create_test_file(Utc::now());
316        store.insert("hash2", file2).unwrap();
317        let size2 = store.stats().unwrap().total_size;
318
319        // Size should have increased
320        assert!(size2 > size1);
321    }
322
323    #[test]
324    fn test_clone() {
325        let dir = temp_dir();
326        let store1 = HashStore::new(&dir).unwrap();
327
328        let file = create_test_file(Utc::now());
329        store1.insert("hash1", file).unwrap();
330
331        // Clone the store
332        let store2 = store1.clone();
333
334        // Both should be able to read the same data
335        assert!(store1.get("hash1").unwrap().is_some());
336        assert!(store2.get("hash1").unwrap().is_some());
337        assert_eq!(store1.count().unwrap(), store2.count().unwrap());
338    }
339
340    #[test]
341    fn test_thread_safety() {
342        use std::thread;
343
344        let dir = temp_dir();
345        let store = HashStore::new(&dir).unwrap();
346
347        let store1 = store.clone();
348        let store2 = store.clone();
349
350        let handle1 = thread::spawn(move || {
351            for i in 0..50 {
352                let hash = format!("thread1_hash{}", i);
353                let file = create_test_file(Utc::now());
354                store1.insert(&hash, file).unwrap();
355            }
356        });
357
358        let handle2 = thread::spawn(move || {
359            for i in 0..50 {
360                let hash = format!("thread2_hash{}", i);
361                let file = create_test_file(Utc::now());
362                store2.insert(&hash, file).unwrap();
363            }
364        });
365
366        handle1.join().unwrap();
367        handle2.join().unwrap();
368
369        // Should have 100 entries total
370        assert_eq!(store.count().unwrap(), 100);
371    }
372
373    #[test]
374    fn test_update_existing_entry() {
375        let dir = temp_dir();
376        let store = HashStore::new(&dir).unwrap();
377
378        let hash = "hash1";
379        let file1 = create_test_file(Utc.with_ymd_and_hms(2024, 1, 1, 12, 0, 0).unwrap());
380        let file2 = create_test_file(Utc.with_ymd_and_hms(2024, 2, 1, 12, 0, 0).unwrap());
381
382        // Insert first version
383        store.insert(hash, file1).unwrap();
384        let retrieved1 = store.get(hash).unwrap().unwrap();
385        assert_eq!(
386            retrieved1.uploaded_at,
387            Utc.with_ymd_and_hms(2024, 1, 1, 12, 0, 0).unwrap()
388        );
389
390        // Update with second version
391        store.insert(hash, file2).unwrap();
392        let retrieved2 = store.get(hash).unwrap().unwrap();
393        assert_eq!(
394            retrieved2.uploaded_at,
395            Utc.with_ymd_and_hms(2024, 2, 1, 12, 0, 0).unwrap()
396        );
397
398        // Should still have only 1 entry
399        assert_eq!(store.count().unwrap(), 1);
400    }
401}