1use std::fs;
2use std::path::{Path, PathBuf};
3use std::time::{Duration, SystemTime, UNIX_EPOCH};
4
5use anyhow::{Context, Result, bail};
6use sqlx::sqlite::SqliteConnectOptions;
7use sqlx::{Connection as _, SqliteConnection, SqlitePool};
8
9use super::MIGRATOR;
10
11const MIGRATION_BACKUP_KEEP: usize = 20;
12
13pub(super) async fn backup_before_pending_migrations(
14 path: &Path,
15 existed_before_open: bool,
16 pool: &SqlitePool,
17) -> Result<()> {
18 if !migration_backups_enabled() || !existed_before_open || !has_pending_migrations(pool).await?
19 {
20 return Ok(());
21 }
22 let backup_path = migration_backup_path(path)?;
23 let mut conn = pool.acquire().await?;
24 backup_database_with_connection(&mut conn, &backup_path).await?;
25 prune_migration_backups(path)?;
26 Ok(())
27}
28
29fn migration_backups_enabled() -> bool {
30 std::env::var_os("AVEN_DEV_MIGRATION_BACKUPS").is_some()
31}
32
33async fn has_pending_migrations(pool: &SqlitePool) -> Result<bool> {
34 let applied_versions =
35 match sqlx::query_scalar::<_, i64>("SELECT version FROM _sqlx_migrations")
36 .fetch_all(pool)
37 .await
38 {
39 Ok(versions) => versions,
40 Err(error) => {
41 let Some(db_error) = error.as_database_error() else {
42 return Err(error.into());
43 };
44 if db_error.code().as_deref() == Some("1") {
45 return Ok(MIGRATOR.iter().next().is_some());
46 }
47 return Err(error.into());
48 }
49 };
50 Ok(MIGRATOR
51 .iter()
52 .any(|migration| !applied_versions.contains(&migration.version)))
53}
54
55fn migration_backup_path(path: &Path) -> Result<PathBuf> {
56 default_sqlite_backup_path(path, "before-migrate")
57}
58
59pub fn default_backup_path(path: &Path, reason: &str) -> Result<PathBuf> {
60 backup_path_with_extension(path, reason, "aven-backup.tar.zst")
61}
62
63pub fn default_sqlite_backup_path(path: &Path, reason: &str) -> Result<PathBuf> {
64 backup_path_with_extension(path, reason, "sqlite")
65}
66
67fn backup_path_with_extension(path: &Path, reason: &str, extension: &str) -> Result<PathBuf> {
68 let parent = path.parent().unwrap_or_else(|| Path::new("."));
69 let backup_dir = parent.join("backups");
70 fs::create_dir_all(&backup_dir)
71 .with_context(|| format!("could not create {}", backup_dir.display()))?;
72 let stem = path
73 .file_name()
74 .and_then(|name| name.to_str())
75 .unwrap_or("db.sqlite");
76 Ok(backup_dir.join(format!(
77 "{stem}.{reason}-{}.{}",
78 backup_timestamp()?,
79 extension
80 )))
81}
82
83pub async fn backup_database(source: &Path, backup: &Path) -> Result<()> {
84 if !source.is_file() {
85 bail!("could not open source {}", source.display());
86 }
87 if let Some(parent) = backup.parent() {
88 fs::create_dir_all(parent)
89 .with_context(|| format!("could not create {}", parent.display()))?;
90 }
91 let mut conn = SqliteConnection::connect_with(
92 &SqliteConnectOptions::new()
93 .filename(source)
94 .read_only(true)
95 .foreign_keys(true)
96 .busy_timeout(Duration::from_secs(5)),
97 )
98 .await
99 .with_context(|| format!("could not open source {}", source.display()))?;
100 backup_database_with_connection(&mut conn, backup).await
101}
102
103pub fn wal_path(path: &Path) -> PathBuf {
104 PathBuf::from(format!("{}-wal", path.display()))
105}
106
107pub fn shm_path(path: &Path) -> PathBuf {
108 PathBuf::from(format!("{}-shm", path.display()))
109}
110
111pub async fn restore_database_file(target: &Path, source: &Path) -> Result<PathBuf> {
112 validate_sqlite_source(source).await?;
113 let safety = create_restore_safety_backup(target).await?;
114 let staging = target.with_extension("restore-staging");
115 if staging.exists() {
116 fs::remove_file(&staging)
117 .with_context(|| format!("could not remove {}", staging.display()))?;
118 }
119 fs::copy(source, &staging).with_context(|| {
120 format!(
121 "could not copy {} -> {}",
122 source.display(),
123 staging.display()
124 )
125 })?;
126 for sidecar in [wal_path(target), shm_path(target)] {
127 if sidecar.exists() {
128 fs::remove_file(&sidecar)
129 .with_context(|| format!("could not remove {}", sidecar.display()))?;
130 }
131 }
132 fs::rename(&staging, target)
133 .with_context(|| format!("could not replace {}", target.display()))?;
134 Ok(safety)
135}
136
137pub(crate) async fn create_restore_safety_backup(target: &Path) -> Result<PathBuf> {
138 let safety = default_sqlite_backup_path(target, "before-restore")?;
139 if target.exists() {
140 backup_database(target, &safety).await?;
141 } else {
142 SqliteConnection::connect_with(
143 &SqliteConnectOptions::new()
144 .filename(&safety)
145 .create_if_missing(true),
146 )
147 .await
148 .with_context(|| format!("could not create {}", safety.display()))?
149 .close()
150 .await?;
151 }
152 Ok(safety)
153}
154
155async fn validate_sqlite_source(source: &Path) -> Result<()> {
156 let mut conn = sqlx::SqliteConnection::connect_with(
157 &SqliteConnectOptions::new()
158 .filename(source)
159 .read_only(true)
160 .foreign_keys(true),
161 )
162 .await
163 .with_context(|| format!("could not open source {}", source.display()))?;
164 let quick_check: String = sqlx::query_scalar("PRAGMA quick_check")
165 .fetch_one(&mut conn)
166 .await?;
167 if quick_check != "ok" {
168 bail!("error backup-source-corrupt quick_check={quick_check}");
169 }
170 Ok(())
171}
172
173fn backup_timestamp() -> Result<u64> {
174 Ok(SystemTime::now()
175 .duration_since(UNIX_EPOCH)
176 .context("system clock is before unix epoch")?
177 .as_secs())
178}
179
180pub(crate) async fn backup_database_with_connection(
181 conn: &mut SqliteConnection,
182 backup: &Path,
183) -> Result<()> {
184 let parent = backup.parent().unwrap_or_else(|| Path::new("."));
185 fs::create_dir_all(parent).with_context(|| format!("could not create {}", parent.display()))?;
186 let staging_dir = tempfile::Builder::new()
187 .prefix(".aven-sqlite-backup-")
188 .tempdir_in(parent)
189 .with_context(|| format!("could not create backup staging in {}", parent.display()))?;
190 let staging = staging_dir.path().join("database.sqlite");
191 sqlx::query("VACUUM INTO ?")
192 .bind(staging.display().to_string())
193 .execute(&mut *conn)
194 .await
195 .with_context(|| format!("could not back up database to {}", backup.display()))?;
196 fs::rename(&staging, backup)
197 .with_context(|| format!("could not replace {}", backup.display()))?;
198 Ok(())
199}
200
201fn prune_migration_backups(path: &Path) -> Result<()> {
202 let Some(parent) = path.parent() else {
203 return Ok(());
204 };
205 let backup_dir = parent.join("backups");
206 let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
207 return Ok(());
208 };
209 let prefix = format!("{file_name}.before-migrate-");
210 let mut backups = fs::read_dir(&backup_dir)
211 .with_context(|| format!("could not read {}", backup_dir.display()))?
212 .filter_map(|entry| entry.ok())
213 .filter(|entry| {
214 entry
215 .file_name()
216 .to_str()
217 .is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".sqlite"))
218 })
219 .collect::<Vec<_>>();
220 backups.sort_by_key(|entry| entry.file_name());
221 let remove_count = backups.len().saturating_sub(MIGRATION_BACKUP_KEEP);
222 for entry in backups.into_iter().take(remove_count) {
223 let path = entry.path();
224 fs::remove_file(&path).with_context(|| format!("could not remove {}", path.display()))?;
225 }
226 Ok(())
227}
228
229#[cfg(test)]
230mod tests {
231 use super::super::{Database, begin_immediate, get_meta, set_meta};
232 use super::*;
233
234 #[tokio::test]
235 async fn in_process_backup_captures_wal_and_replaces_destination() {
236 let temp = tempfile::tempdir().unwrap();
237 let source = temp.path().join("source.sqlite");
238 let backup = temp.path().join("backup.sqlite");
239 let database = Database::open(&source).await.unwrap();
240 let mut writer = database.acquire_writer().await.unwrap();
241 let mut tx = begin_immediate(&mut writer).await.unwrap();
242 set_meta(&mut tx, "backup-test", "first").await.unwrap();
243 tx.commit().await.unwrap();
244 drop(writer);
245 assert!(wal_path(&source).exists());
246
247 fs::write(&backup, b"existing destination").unwrap();
248 backup_database(&source, &backup).await.unwrap();
249 let mut backup_conn = SqliteConnection::connect_with(
250 &SqliteConnectOptions::new()
251 .filename(&backup)
252 .read_only(true),
253 )
254 .await
255 .unwrap();
256 assert_eq!(
257 get_meta(&mut backup_conn, "backup-test")
258 .await
259 .unwrap()
260 .as_deref(),
261 Some("first")
262 );
263 drop(backup_conn);
264
265 let mut writer = database.acquire_writer().await.unwrap();
266 set_meta(&mut writer, "backup-test", "second")
267 .await
268 .unwrap();
269 drop(writer);
270 backup_database(&source, &backup).await.unwrap();
271 let mut backup_conn = SqliteConnection::connect_with(
272 &SqliteConnectOptions::new()
273 .filename(&backup)
274 .read_only(true),
275 )
276 .await
277 .unwrap();
278 assert_eq!(
279 get_meta(&mut backup_conn, "backup-test")
280 .await
281 .unwrap()
282 .as_deref(),
283 Some("second")
284 );
285 }
286
287 #[tokio::test]
288 async fn in_process_backup_rejects_missing_source() {
289 let temp = tempfile::tempdir().unwrap();
290 let backup = temp.path().join("backup.sqlite");
291 let error = backup_database(&temp.path().join("missing.sqlite"), &backup)
292 .await
293 .unwrap_err();
294 assert!(error.to_string().contains("could not open source"));
295 assert!(!backup.exists());
296 }
297
298 #[tokio::test]
299 async fn restore_replaces_sidecars_and_preserves_safety_copy() {
300 let temp = tempfile::tempdir().unwrap();
301 let target = temp.path().join("target.sqlite");
302 let source = temp.path().join("source.sqlite");
303 let target_database = Database::open(&target).await.unwrap();
304 let source_database = Database::open(&source).await.unwrap();
305 let mut target_writer = target_database.acquire_writer().await.unwrap();
306 set_meta(&mut target_writer, "restore-test", "target")
307 .await
308 .unwrap();
309 drop(target_writer);
310 let mut source_writer = source_database.acquire_writer().await.unwrap();
311 set_meta(&mut source_writer, "restore-test", "source")
312 .await
313 .unwrap();
314 drop(source_writer);
315 target_database.pool.close().await;
316 source_database.pool.close().await;
317 fs::write(wal_path(&target), b"stale wal").unwrap();
318 fs::write(shm_path(&target), b"stale shm").unwrap();
319
320 let safety = restore_database_file(&target, &source).await.unwrap();
321 assert!(!wal_path(&target).exists());
322 assert!(!shm_path(&target).exists());
323
324 let mut restored = SqliteConnection::connect_with(
325 &SqliteConnectOptions::new()
326 .filename(&target)
327 .read_only(true),
328 )
329 .await
330 .unwrap();
331 assert_eq!(
332 get_meta(&mut restored, "restore-test")
333 .await
334 .unwrap()
335 .as_deref(),
336 Some("source")
337 );
338 let mut preserved = SqliteConnection::connect_with(
339 &SqliteConnectOptions::new()
340 .filename(&safety)
341 .read_only(true),
342 )
343 .await
344 .unwrap();
345 assert_eq!(
346 get_meta(&mut preserved, "restore-test")
347 .await
348 .unwrap()
349 .as_deref(),
350 Some("target")
351 );
352 }
353}