1use orbok_core::{OrbokResult, StorageCategory};
9use orbok_db::Catalog;
10use orbok_db::repo::StorageAccountingRepository;
11use std::path::Path;
12
13pub fn update_storage_accounting(
19 catalog: &Catalog,
20 cache_db_path: &Path,
21) -> OrbokResult<Vec<(StorageCategory, u64, u64)>> {
22 let storage = StorageAccountingRepository::new(catalog);
23 let mut rows = Vec::new();
24
25 macro_rules! measure {
26 ($cat:expr, $size:expr, $count:expr) => {{
27 storage.upsert($cat, $size, $count)?;
28 rows.push(($cat, $size, $count));
29 }};
30 }
31
32 let conn = catalog.lock();
33
34 let catalog_path = catalog.path();
37 let catalog_bytes = if catalog_path.to_str() == Some(":memory:") {
38 let pages: i64 = conn
40 .query_row("PRAGMA page_count", [], |r| r.get(0))
41 .unwrap_or(0);
42 let page_size: i64 = conn
43 .query_row("PRAGMA page_size", [], |r| r.get(0))
44 .unwrap_or(4096);
45 (pages * page_size) as u64
46 } else {
47 std::fs::metadata(catalog_path).map(|m| m.len()).unwrap_or(0)
48 };
49 let source_count: i64 = conn
51 .query_row("SELECT COUNT(*) FROM sources WHERE status != 'removed'", [], |r| r.get(0))
52 .unwrap_or(0);
53 drop(conn); measure!(StorageCategory::PersistentCatalog, catalog_bytes, source_count as u64);
55
56 let conn = catalog.lock();
58 let kw_count: i64 = conn
59 .query_row("SELECT COUNT(*) FROM keyword_index_records WHERE status='active'", [], |r| r.get(0))
60 .unwrap_or(0);
61 let kw_bytes = kw_count as u64 * 256;
63 drop(conn);
64 measure!(StorageCategory::KeywordIndex, kw_bytes, kw_count as u64);
65
66 let conn = catalog.lock();
68 let (emb_count, emb_bytes): (i64, i64) = conn
69 .query_row(
70 "SELECT COUNT(*), COALESCE(SUM(LENGTH(vector_blob)), 0) FROM embeddings WHERE status='active'",
71 [],
72 |r| Ok((r.get(0)?, r.get(1)?)),
73 )
74 .unwrap_or((0, 0));
75 drop(conn);
76 measure!(StorageCategory::VectorIndex, emb_bytes as u64, emb_count as u64);
77
78 let conn = catalog.lock();
80 let (snip_count, snip_bytes): (i64, i64) = conn
81 .query_row(
82 "SELECT COUNT(*), COALESCE(SUM(size_bytes), 0) FROM snippet_cache",
83 [],
84 |r| Ok((r.get(0)?, r.get(1)?)),
85 )
86 .unwrap_or((0, 0));
87 drop(conn);
88 measure!(StorageCategory::SnippetCache, snip_bytes as u64, snip_count as u64);
89
90 let conn = catalog.lock();
92 let sr_count: i64 = conn
93 .query_row("SELECT COUNT(*) FROM search_result_cache", [], |r| r.get(0))
94 .unwrap_or(0);
95 drop(conn);
96 measure!(StorageCategory::SearchCache, sr_count as u64 * 512, sr_count as u64);
97
98 let cache_bytes = std::fs::metadata(cache_db_path)
100 .map(|m| m.len())
101 .unwrap_or(0);
102 let conn = catalog.lock();
103 let extract_count: i64 = conn
104 .query_row("SELECT COUNT(*) FROM extraction_records WHERE status='succeeded'", [], |r| r.get(0))
105 .unwrap_or(0);
106 drop(conn);
107 measure!(StorageCategory::TemporaryExtraction, cache_bytes, extract_count as u64);
108
109 let conn = catalog.lock();
111 let evt_count: i64 = conn
112 .query_row("SELECT COUNT(*) FROM app_events", [], |r| r.get(0))
113 .unwrap_or(0);
114 drop(conn);
115 measure!(StorageCategory::Logs, evt_count as u64 * 256, evt_count as u64);
116
117 measure!(StorageCategory::ModelFiles, 0, 0);
119
120 Ok(rows)
121}