agentmux_srv\backend\storage/
snapshot.rs1use std::fs;
29use std::path::{Path, PathBuf};
30use std::time::SystemTime;
31
32use chrono::{DateTime, Utc};
33use rusqlite::Connection;
34use tracing::{info, warn};
35
36const NAME_PRE_TOKEN: &str = "-pre-v";
41const NAME_SUFFIX: &str = ".bak";
42
43pub const MAX_SNAPSHOTS_PER_CHANNEL: usize = 5;
47
48const SNAPSHOT_DB_NAMES: &[&str] = &["objects.db", "sagas.db", "filestore.db"];
50
51#[derive(Debug, thiserror::Error)]
52pub enum SnapshotError {
53 #[error("io error during snapshot ({path}): {source}")]
54 Io {
55 path: PathBuf,
56 #[source]
57 source: std::io::Error,
58 },
59
60 #[error("sqlite error during snapshot: {0}")]
61 Sql(#[from] rusqlite::Error),
62
63 #[error("invalid channel name for snapshot: {0:?}")]
64 InvalidChannel(String),
65}
66
67impl SnapshotError {
68 fn io(path: impl Into<PathBuf>, source: std::io::Error) -> Self {
69 SnapshotError::Io {
70 path: path.into(),
71 source,
72 }
73 }
74}
75
76pub fn needs_snapshot(
82 objects_db: &Path,
83 current_schema_version: i64,
84) -> Result<Option<i64>, SnapshotError> {
85 if !objects_db.exists() {
86 return Ok(None);
87 }
88 let conn = Connection::open_with_flags(
92 objects_db,
93 rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY | rusqlite::OpenFlags::SQLITE_OPEN_URI,
94 )?;
95 let found: i64 = conn.query_row("PRAGMA user_version", [], |row| row.get(0))?;
96 drop(conn);
97 if found < current_schema_version {
98 Ok(Some(found))
99 } else {
100 Ok(None)
101 }
102}
103
104pub fn create_snapshot(
109 db_dir: &Path,
110 snapshots_dir: &Path,
111 channel: &str,
112 code_version: &str,
113) -> Result<PathBuf, SnapshotError> {
114 validate_channel(channel)?;
115 let snap_name = snapshot_dir_name(channel, code_version, SystemTime::now());
116 let snap_dir = snapshots_dir.join(&snap_name);
117 fs::create_dir_all(&snap_dir).map_err(|e| SnapshotError::io(&snap_dir, e))?;
118 for db_name in SNAPSHOT_DB_NAMES {
119 let src = db_dir.join(db_name);
120 if !src.exists() {
121 continue;
122 }
123 let dst = snap_dir.join(db_name);
124 let conn = Connection::open(&src)?;
125 let quoted = quote_sql_literal(&dst.to_string_lossy());
132 conn.execute_batch(&format!("VACUUM INTO {quoted};"))?;
133 }
134 info!(
135 snapshot = %snap_dir.display(),
136 channel,
137 code_version,
138 "created pre-migration snapshot",
139 );
140 Ok(snap_dir)
141}
142
143pub fn prune_snapshots(snapshots_dir: &Path, channel: &str) -> Result<usize, SnapshotError> {
148 if !snapshots_dir.exists() {
149 return Ok(0);
150 }
151 let prefix = format!("{channel}{NAME_PRE_TOKEN}");
152 let mut matches: Vec<(String, PathBuf)> = Vec::new();
157 let entries = fs::read_dir(snapshots_dir).map_err(|e| SnapshotError::io(snapshots_dir, e))?;
158 for entry in entries {
159 let entry = entry.map_err(|e| SnapshotError::io(snapshots_dir, e))?;
160 let path = entry.path();
161 if !path.is_dir() {
162 continue;
163 }
164 let fname = entry.file_name().to_string_lossy().into_owned();
165 if !fname.starts_with(&prefix) || !fname.ends_with(NAME_SUFFIX) {
166 continue;
167 }
168 let Some(ts) = parse_ts_from_name(&fname, &prefix) else {
169 warn!(snapshot = %fname, "snapshot name doesn't match expected layout — ignored for prune");
171 continue;
172 };
173 matches.push((ts, path));
174 }
175 if matches.len() <= MAX_SNAPSHOTS_PER_CHANNEL {
176 return Ok(0);
177 }
178 matches.sort_by(|a, b| a.0.cmp(&b.0));
181 let to_remove = matches.len() - MAX_SNAPSHOTS_PER_CHANNEL;
182 let mut removed = 0;
183 for (_, path) in matches.into_iter().take(to_remove) {
184 match fs::remove_dir_all(&path) {
185 Ok(()) => removed += 1,
186 Err(e) => warn!(snapshot = %path.display(), err = %e, "failed to prune snapshot"),
187 }
188 }
189 Ok(removed)
190}
191
192fn parse_ts_from_name(fname: &str, channel_prefix: &str) -> Option<String> {
197 let rest = fname.strip_prefix(channel_prefix)?;
198 let rest = rest.strip_suffix(NAME_SUFFIX)?;
199 const ISO_LEN: usize = 20; if rest.len() <= ISO_LEN + 1 {
206 return None;
207 }
208 let sep_idx = rest.len() - ISO_LEN - 1;
209 if rest.as_bytes().get(sep_idx) != Some(&b'-') {
210 return None;
211 }
212 Some(rest[sep_idx + 1..].to_string())
213}
214
215pub fn maybe_snapshot_pre_migration(
220 db_dir: &Path,
221 snapshots_dir: &Path,
222 channel: &str,
223 code_version: &str,
224 current_schema_version: i64,
225) -> Result<Option<PathBuf>, SnapshotError> {
226 let objects_db = db_dir.join("objects.db");
227 let Some(found_version) = needs_snapshot(&objects_db, current_schema_version)? else {
228 return Ok(None);
229 };
230 info!(
231 from_version = found_version,
232 to_version = current_schema_version,
233 channel,
234 "pre-migration snapshot needed",
235 );
236 fs::create_dir_all(snapshots_dir).map_err(|e| SnapshotError::io(snapshots_dir, e))?;
237 let snap = create_snapshot(db_dir, snapshots_dir, channel, code_version)?;
238 let pruned = prune_snapshots(snapshots_dir, channel)?;
239 if pruned > 0 {
240 info!(pruned, channel, "pruned old snapshots");
241 }
242 Ok(Some(snap))
243}
244
245fn validate_channel(channel: &str) -> Result<(), SnapshotError> {
252 if channel.is_empty()
253 || channel.contains('/')
254 || channel.contains('\\')
255 || channel.contains("..")
256 || channel.contains('\0')
257 {
258 return Err(SnapshotError::InvalidChannel(channel.to_string()));
259 }
260 Ok(())
261}
262
263fn snapshot_dir_name(channel: &str, code_version: &str, t: SystemTime) -> String {
264 let dt: DateTime<Utc> = t.into();
265 let ts = dt.format("%Y-%m-%dT%H-%M-%SZ").to_string();
268 format!("{channel}-pre-v{code_version}-{ts}.bak")
269}
270
271fn quote_sql_literal(s: &str) -> String {
272 let escaped = s.replace('\'', "''");
275 format!("'{escaped}'")
276}
277
278#[cfg(test)]
281mod tests {
282 use super::*;
283 use std::time::{Duration, UNIX_EPOCH};
284
285 fn make_db_with_version(path: &Path, version: i64) {
286 let conn = Connection::open(path).unwrap();
287 conn.execute_batch(&format!(
288 "CREATE TABLE IF NOT EXISTS t (x INTEGER); INSERT INTO t VALUES (42); PRAGMA user_version = {version};"
289 )).unwrap();
290 }
291
292 #[test]
293 fn needs_snapshot_returns_none_for_fresh_install() {
294 let tmp = tempfile::tempdir().unwrap();
295 let missing = tmp.path().join("objects.db");
296 assert_eq!(needs_snapshot(&missing, 4).unwrap(), None);
297 }
298
299 #[test]
300 fn needs_snapshot_returns_none_when_same_version() {
301 let tmp = tempfile::tempdir().unwrap();
302 let p = tmp.path().join("objects.db");
303 make_db_with_version(&p, 4);
304 assert_eq!(needs_snapshot(&p, 4).unwrap(), None);
305 }
306
307 #[test]
308 fn needs_snapshot_returns_some_when_db_is_behind() {
309 let tmp = tempfile::tempdir().unwrap();
310 let p = tmp.path().join("objects.db");
311 make_db_with_version(&p, 2);
312 assert_eq!(needs_snapshot(&p, 4).unwrap(), Some(2));
313 }
314
315 #[test]
316 fn needs_snapshot_returns_some_for_legacy_db_at_v0() {
317 let tmp = tempfile::tempdir().unwrap();
321 let p = tmp.path().join("objects.db");
322 let conn = Connection::open(&p).unwrap();
323 conn.execute_batch("CREATE TABLE t (x INTEGER); INSERT INTO t VALUES (1);")
324 .unwrap();
325 assert_eq!(needs_snapshot(&p, 4).unwrap(), Some(0));
326 }
327
328 #[test]
329 fn create_snapshot_copies_existing_dbs_only() {
330 let tmp = tempfile::tempdir().unwrap();
331 let db_dir = tmp.path().join("db");
332 let snaps_dir = tmp.path().join("snapshots");
333 fs::create_dir_all(&db_dir).unwrap();
334 make_db_with_version(&db_dir.join("objects.db"), 2);
335 make_db_with_version(&db_dir.join("sagas.db"), 1);
337
338 let snap = create_snapshot(&db_dir, &snaps_dir, "stable", "0.39.0").unwrap();
339 assert!(snap.is_dir());
340 assert!(snap.join("objects.db").exists());
341 assert!(snap.join("sagas.db").exists());
342 assert!(!snap.join("filestore.db").exists());
343
344 let conn = Connection::open(snap.join("objects.db")).unwrap();
346 let v: i64 = conn
347 .query_row("PRAGMA user_version", [], |r| r.get(0))
348 .unwrap();
349 assert_eq!(v, 2);
350 let x: i64 = conn.query_row("SELECT x FROM t LIMIT 1", [], |r| r.get(0)).unwrap();
352 assert_eq!(x, 42);
353 }
354
355 #[test]
356 fn snapshot_dir_name_uses_filename_safe_iso8601() {
357 let t = UNIX_EPOCH + Duration::from_secs(1_716_634_800); let name = snapshot_dir_name("stable", "0.39.0", t);
359 assert_eq!(name, "stable-pre-v0.39.0-2024-05-25T11-00-00Z.bak");
360 assert!(!name.contains(':'));
362 }
363
364 #[test]
365 fn validate_channel_rejects_traversal() {
366 assert!(validate_channel("").is_err());
367 assert!(validate_channel("..").is_err());
368 assert!(validate_channel("a/b").is_err());
369 assert!(validate_channel("a\\b").is_err());
370 assert!(validate_channel("ok").is_ok());
371 assert!(validate_channel("stable").is_ok());
372 }
373
374 #[test]
375 fn prune_keeps_newest_n_per_channel() {
376 let tmp = tempfile::tempdir().unwrap();
377 let snaps = tmp.path();
378 let extras = 3;
382 let mut created = Vec::new();
383 for i in 0..(MAX_SNAPSHOTS_PER_CHANNEL + extras) {
384 let t = UNIX_EPOCH + Duration::from_secs(1_700_000_000 + i as u64 * 60);
385 let name = snapshot_dir_name("stable", "0.39.0", t);
386 let p = snaps.join(&name);
387 fs::create_dir_all(&p).unwrap();
388 created.push((i, p));
389 }
390 let other = snaps.join("beta-pre-v0.39.0-2020-01-01T00-00-00Z.bak");
392 fs::create_dir_all(&other).unwrap();
393
394 let pruned = prune_snapshots(snaps, "stable").unwrap();
395 assert_eq!(pruned, extras);
396 for (i, p) in &created {
398 let should_exist = *i >= extras;
399 assert_eq!(p.exists(), should_exist, "i={i}");
400 }
401 assert!(other.exists(), "other-channel snapshot must survive");
402 }
403
404 #[test]
405 fn parse_ts_from_name_extracts_timestamp() {
406 let n = "stable-pre-v0.39.0-2024-05-25T11-00-00Z.bak";
407 let prefix = "stable-pre-v";
408 assert_eq!(
409 parse_ts_from_name(n, prefix).as_deref(),
410 Some("2024-05-25T11-00-00Z"),
411 );
412 }
413
414 #[test]
415 fn parse_ts_handles_semver_prerelease_in_code_version() {
416 let n = "stable-pre-v0.39.0-rc1-2024-05-25T11-00-00Z.bak";
417 let prefix = "stable-pre-v";
418 assert_eq!(
419 parse_ts_from_name(n, prefix).as_deref(),
420 Some("2024-05-25T11-00-00Z"),
421 );
422 }
423
424 #[test]
425 fn prune_noop_when_under_limit() {
426 let tmp = tempfile::tempdir().unwrap();
427 for i in 0..3 {
428 let t = UNIX_EPOCH + Duration::from_secs(1_700_000_000 + i as u64 * 60);
429 let p = tmp.path().join(snapshot_dir_name("stable", "0.39.0", t));
430 fs::create_dir_all(&p).unwrap();
431 }
432 assert_eq!(prune_snapshots(tmp.path(), "stable").unwrap(), 0);
433 }
434
435 #[test]
436 fn maybe_snapshot_returns_none_for_fresh_install() {
437 let tmp = tempfile::tempdir().unwrap();
438 let db_dir = tmp.path().join("db");
439 let snaps = tmp.path().join("snapshots");
440 fs::create_dir_all(&db_dir).unwrap();
441 let result =
442 maybe_snapshot_pre_migration(&db_dir, &snaps, "stable", "0.39.0", 4).unwrap();
443 assert!(result.is_none());
444 assert!(!snaps.exists());
446 }
447
448 #[test]
449 fn maybe_snapshot_runs_when_db_behind() {
450 let tmp = tempfile::tempdir().unwrap();
451 let db_dir = tmp.path().join("db");
452 let snaps = tmp.path().join("snapshots");
453 fs::create_dir_all(&db_dir).unwrap();
454 make_db_with_version(&db_dir.join("objects.db"), 2);
455 let result =
456 maybe_snapshot_pre_migration(&db_dir, &snaps, "stable", "0.39.0", 4).unwrap();
457 let snap = result.expect("snapshot should be created");
458 assert!(snap.is_dir());
459 assert!(snap.join("objects.db").exists());
460 }
461}