1use rusqlite::{params, OptionalExtension};
15
16use super::core::FileStore;
17use super::{FileMeta, FileOpts};
18use crate::backend::storage::error::StoreError;
19
20pub struct SnapshotReader<'a> {
23 tx: &'a rusqlite::Transaction<'a>,
24}
25
26#[derive(Debug, Clone)]
28pub struct SnapshotFile {
29 pub size: i64,
30 pub gen: Option<String>,
32 pub meta: FileMeta,
33}
34
35impl SnapshotReader<'_> {
36 pub fn file(&self, zone_id: &str, name: &str) -> Result<Option<SnapshotFile>, StoreError> {
37 let Some(row) = super::counter::read_row(self.tx, zone_id, name)? else { return Ok(None) };
38 let meta: String = self.tx.query_row(
39 "SELECT meta FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
40 params![zone_id, name],
41 |r| r.get(0),
42 )?;
43 Ok(Some(SnapshotFile {
44 size: row.size,
45 gen: row.counter().map(|(gen, _)| gen),
46 meta: serde_json::from_str(&meta).unwrap_or_default(),
47 }))
48 }
49
50 pub fn bytes(&self, zone_id: &str, name: &str, offset: i64, len: i64) -> Result<Option<Vec<u8>>, StoreError> {
52 super::counter::read_bytes_exact(self.tx, zone_id, name, offset, len)
53 }
54}
55
56fn patch_range_in(tx: &rusqlite::Transaction<'_>, zone_id: &str, name: &str, offset: i64, data: &[u8]) -> Result<(), StoreError> {
61 if data.is_empty() {
62 return Ok(());
63 }
64 let pds = super::core::PART_DATA_SIZE as i64;
65 let end = offset + data.len() as i64;
66 for idx in (offset / pds)..=((end - 1) / pds) {
67 let part_start = idx * pds;
68 let mut part: Vec<u8> = tx
69 .query_row(
70 "SELECT data FROM db_file_data WHERE zoneid = ?1 AND name = ?2 AND partidx = ?3",
71 params![zone_id, name, idx as i32],
72 |r| r.get(0),
73 )
74 .optional()?
75 .unwrap_or_default();
76 let from = offset.max(part_start);
77 let to = end.min(part_start + pds);
78 let (lo, hi) = ((from - part_start) as usize, (to - part_start) as usize);
79 if part.len() < hi {
80 part.resize(hi, 0);
81 }
82 part[lo..hi].copy_from_slice(&data[(from - offset) as usize..(to - offset) as usize]);
83 tx.execute(
84 "REPLACE INTO db_file_data (zoneid, name, partidx, data) VALUES (?1, ?2, ?3, ?4)",
85 params![zone_id, name, idx as i32, part],
86 )?;
87 }
88 Ok(())
89}
90
91fn truncate_parts_in(tx: &rusqlite::Transaction<'_>, zone_id: &str, name: &str, new_size: i64) -> Result<(), StoreError> {
93 let pds = super::core::PART_DATA_SIZE as i64;
94 let keep_parts = (new_size + pds - 1) / pds;
95 tx.execute(
96 "DELETE FROM db_file_data WHERE zoneid = ?1 AND name = ?2 AND partidx >= ?3",
97 params![zone_id, name, keep_parts as i32],
98 )?;
99 let in_last = (new_size % pds) as usize;
100 if in_last > 0 {
101 let last = (keep_parts - 1) as i32;
102 let part: Option<Vec<u8>> = tx
103 .query_row(
104 "SELECT data FROM db_file_data WHERE zoneid = ?1 AND name = ?2 AND partidx = ?3",
105 params![zone_id, name, last],
106 |r| r.get(0),
107 )
108 .optional()?;
109 if let Some(mut part) = part.filter(|p| p.len() > in_last) {
110 part.truncate(in_last);
111 tx.execute(
112 "REPLACE INTO db_file_data (zoneid, name, partidx, data) VALUES (?1, ?2, ?3, ?4)",
113 params![zone_id, name, last, part],
114 )?;
115 }
116 }
117 Ok(())
118}
119
120#[derive(Debug, Clone)]
122pub struct DerivedSnapshot {
123 pub output_size: i64,
124 pub output_gen: Option<String>,
126 pub derived: Option<DerivedView>,
127}
128
129#[derive(Debug, Clone)]
130pub struct DerivedView {
131 pub size: i64,
132 pub bytes: Option<Vec<u8>>,
134 pub meta: FileMeta,
135}
136
137impl FileStore {
138 #[track_caller]
142 pub fn replace_file(&self, zone_id: &str, name: &str, data: &[u8], drop: &[&str]) -> Result<(), StoreError> {
143 self.replace_inner(zone_id, name, data, None, drop, None).map(|_| ())
144 }
145
146 #[track_caller]
151 #[cfg(test)]
152 pub fn put_file_with_meta(&self, zone_id: &str, name: &str, data: &[u8], meta: FileMeta) -> Result<(), StoreError> {
153 self.replace_inner(zone_id, name, data, Some(meta), &[], None).map(|_| ())
154 }
155
156 #[track_caller]
163 pub fn put_file_with_meta_if(
164 &self,
165 zone_id: &str,
166 name: &str,
167 data: &[u8],
168 meta: FileMeta,
169 guard: &str,
170 expected_gen: Option<&str>,
171 ) -> Result<bool, StoreError> {
172 self.replace_inner(zone_id, name, data, Some(meta), &[], Some((guard, expected_gen)))
173 }
174
175 #[track_caller]
184 #[allow(clippy::too_many_arguments)]
185 pub fn patch_derived_if(
186 &self,
187 zone_id: &str,
188 name: &str,
189 head: &[u8],
190 tail_at: i64,
191 tail: &[u8],
192 guard: &str,
193 expected_gen: &str,
194 expected_size: i64,
195 ) -> Result<bool, StoreError> {
196 if tail_at < head.len() as i64 {
197 return Err(StoreError::Other(format!("{zone_id}/{name}: patch tail at {tail_at} overlaps its head")));
198 }
199 let now = agentmux_common::time::now_ms();
200 let written = self.write_txn(|tx| {
201 let current = super::counter::read_row(tx, zone_id, guard)?.and_then(|r| r.counter()).map(|(gen, _)| gen);
202 if current.as_deref() != Some(expected_gen) {
203 return Ok(false);
204 }
205 let size: Option<i64> = tx
206 .query_row(
207 "SELECT size FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
208 params![zone_id, name],
209 |r| r.get(0),
210 )
211 .optional()?;
212 if size != Some(expected_size) || tail_at > expected_size {
213 return Ok(false);
214 }
215 patch_range_in(tx, zone_id, name, 0, head)?;
216 patch_range_in(tx, zone_id, name, tail_at, tail)?;
217 let new_size = tail_at + tail.len() as i64;
218 if new_size < expected_size {
219 truncate_parts_in(tx, zone_id, name, new_size)?;
220 }
221 tx.execute(
222 "UPDATE db_wave_file SET size = ?1, modts = ?2 WHERE zoneid = ?3 AND name = ?4",
223 params![new_size, now, zone_id, name],
224 )?;
225 Ok(true)
226 })?;
227 if written {
228 self.forget_cached(zone_id, &[name]);
229 }
230 Ok(written)
231 }
232
233 #[track_caller]
236 pub fn delete_files(&self, zone_id: &str, names: &[&str]) -> Result<(), StoreError> {
237 self.write_txn(|tx| {
238 for name in names {
239 tx.execute("DELETE FROM db_wave_file WHERE zoneid = ?1 AND name = ?2", params![zone_id, name])?;
240 tx.execute("DELETE FROM db_file_data WHERE zoneid = ?1 AND name = ?2", params![zone_id, name])?;
241 }
242 Ok(())
243 })?;
244 self.forget_cached(zone_id, names);
245 Ok(())
246 }
247
248 #[track_caller]
249 fn replace_inner(
250 &self,
251 zone_id: &str,
252 name: &str,
253 data: &[u8],
254 meta: Option<FileMeta>,
255 drop: &[&str],
256 guard: Option<(&str, Option<&str>)>,
257 ) -> Result<bool, StoreError> {
258 debug_assert!(!drop.contains(&name), "replace_file would drop the file it writes");
259 let now = agentmux_common::time::now_ms();
260 let opts_json = serde_json::to_string(&FileOpts::default())?;
261 let written = self.write_txn(|tx| {
262 if let Some((guard, Some(expected))) = guard {
263 let current = super::counter::read_row(tx, zone_id, guard)?.and_then(|r| r.counter()).map(|(gen, _)| gen);
264 if current.as_deref() != Some(expected) {
265 return Ok(false);
266 }
267 }
268 let current_meta: Option<String> = tx
269 .query_row(
270 "SELECT meta FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
271 params![zone_id, name],
272 |r| r.get(0),
273 )
274 .optional()?;
275 if current_meta.is_none() {
276 tx.execute(
277 "INSERT INTO db_wave_file (zoneid, name, size, createdts, modts, opts, meta)
278 VALUES (?1, ?2, 0, ?3, ?3, ?4, '{}')",
279 params![zone_id, name, now, opts_json],
280 )?;
281 }
282 Self::replace_content_in(tx, zone_id, name, data, now)?;
283 if let Some(meta) = meta {
284 let mut merged: FileMeta =
285 serde_json::from_str(current_meta.as_deref().unwrap_or("{}")).unwrap_or_default();
286 for (k, v) in meta {
287 if v.is_null() {
288 merged.remove(&k);
289 } else {
290 merged.insert(k, v);
291 }
292 }
293 tx.execute(
294 "UPDATE db_wave_file SET meta = ?1 WHERE zoneid = ?2 AND name = ?3",
295 params![serde_json::to_string(&merged)?, zone_id, name],
296 )?;
297 }
298 for d in drop {
299 tx.execute("DELETE FROM db_wave_file WHERE zoneid = ?1 AND name = ?2", params![zone_id, d])?;
300 tx.execute("DELETE FROM db_file_data WHERE zoneid = ?1 AND name = ?2", params![zone_id, d])?;
301 }
302 Ok(true)
303 })?;
304 if written {
305 let mut forget = vec![name];
306 forget.extend_from_slice(drop);
307 self.forget_cached(zone_id, &forget);
308 }
309 Ok(written)
310 }
311
312 #[track_caller]
319 pub fn read_bytes_db(&self, zone_id: &str, name: &str, offset: i64, len: i64) -> Result<Vec<u8>, StoreError> {
320 self.read_conn(|conn| super::counter::read_bytes_exact(conn, zone_id, name, offset, len))?.ok_or_else(|| {
321 StoreError::Other(format!("{zone_id}/{name}: bytes {offset}..{} not stored yet", offset + len))
322 })
323 }
324
325 #[track_caller]
330 pub fn read_snapshot<T>(&self, f: impl FnOnce(&SnapshotReader<'_>) -> Result<T, StoreError>) -> Result<T, StoreError> {
331 self.read_txn(|tx| f(&SnapshotReader { tx }))
332 }
333
334 #[track_caller]
339 pub fn derived_snapshot(
340 &self,
341 zone_id: &str,
342 output: &str,
343 derived: &str,
344 head: Option<i64>,
345 ) -> Result<Option<DerivedSnapshot>, StoreError> {
346 self.read_txn(|tx| {
347 let Some(row) = super::counter::read_row(tx, zone_id, output)? else { return Ok(None) };
348 let output_gen = row.counter().map(|(gen, _)| gen);
349 let derived_row: Option<(i64, String)> = tx
350 .query_row(
351 "SELECT size, meta FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
352 params![zone_id, derived],
353 |r| Ok((r.get(0)?, r.get(1)?)),
354 )
355 .optional()?;
356 let derived = match derived_row {
357 None => None,
358 Some((size, meta)) => {
359 let len = head.map_or(size, |h| h.min(size));
360 Some(DerivedView {
361 size,
362 bytes: super::counter::read_bytes_exact(tx, zone_id, derived, 0, len)?,
363 meta: serde_json::from_str(&meta).unwrap_or_default(),
364 })
365 }
366 };
367 Ok(Some(DerivedSnapshot { output_size: row.size, output_gen, derived }))
368 })
369 }
370
371 #[track_caller]
373 #[cfg(test)]
374 pub fn meta_db(&self, zone_id: &str, name: &str) -> Result<Option<FileMeta>, StoreError> {
375 let conn = self.lock_conn();
376 let meta: Option<String> = conn
377 .query_row(
378 "SELECT meta FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
379 params![zone_id, name],
380 |r| r.get(0),
381 )
382 .optional()?;
383 Ok(meta.map(|m| serde_json::from_str(&m).unwrap_or_default()))
384 }
385}