1use crate::{DbError, DbOptions, Result};
17use rusqlite::{Connection, OptionalExtension};
18use std::path::{Path, PathBuf};
19
20pub enum Step {
21 Sql(&'static str),
22 Code(fn(&Connection) -> rusqlite::Result<()>),
24}
25
26pub struct Migration {
27 pub version: i64,
28 pub name: &'static str,
29 pub up: Step,
30 pub down: Option<&'static str>,
32}
33
34impl Migration {
35 pub const fn sql(version: i64, name: &'static str, up: &'static str) -> Self {
36 Self {
37 version,
38 name,
39 up: Step::Sql(up),
40 down: None,
41 }
42 }
43
44 pub const fn with_down(mut self, down: &'static str) -> Self {
45 self.down = Some(down);
46 self
47 }
48
49 pub const fn code(
50 version: i64,
51 name: &'static str,
52 up: fn(&Connection) -> rusqlite::Result<()>,
53 ) -> Self {
54 Self {
55 version,
56 name,
57 up: Step::Code(up),
58 down: None,
59 }
60 }
61}
62
63#[derive(Clone, Debug, Default, PartialEq, Eq)]
66pub enum VersionStorage {
67 #[default]
68 UserVersion,
69 Table { table: String, key: String },
72}
73
74#[derive(Clone)]
75pub struct MigrateOptions {
76 pub version_storage: VersionStorage,
78 pub backup: bool,
80 pub backup_options: Option<DbOptions>,
82 pub forward_compatible: Option<fn(&Connection) -> bool>,
85 pub optimize_after: bool,
87}
88
89impl Default for MigrateOptions {
90 fn default() -> Self {
91 Self {
92 version_storage: VersionStorage::default(),
93 backup: true,
94 backup_options: None,
95 forward_compatible: None,
96 optimize_after: true,
97 }
98 }
99}
100
101#[derive(Clone, Debug, Default, PartialEq, Eq)]
102pub struct MigrationReport {
103 pub from: i64,
104 pub to: i64,
105 pub applied: Vec<i64>,
106 pub backup: Option<PathBuf>,
107 pub forward_compatible: bool,
109 pub interrupted_previous: Option<PendingMarker>,
111}
112
113#[derive(Clone, Debug, PartialEq, Eq)]
114pub struct PendingMarker {
115 pub from: i64,
116 pub to: i64,
117 pub pid: u32,
118}
119
120pub fn user_version(conn: &Connection) -> Result<i64> {
121 Ok(conn.query_row("PRAGMA user_version", [], |row| row.get(0))?)
122}
123
124fn set_user_version(conn: &Connection, version: i64) -> rusqlite::Result<()> {
125 conn.execute_batch(&format!("PRAGMA user_version = {version};"))
127}
128
129pub fn migration_version(conn: &Connection, storage: &VersionStorage) -> Result<i64> {
132 let VersionStorage::Table { table, key } = storage else {
133 return user_version(conn);
134 };
135 check_ident(table)?;
136 let exists: bool = conn.query_row(
137 "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name=?1 COLLATE NOCASE)",
138 [table],
139 |row| row.get(0),
140 )?;
141 if !exists {
142 return Ok(0);
143 }
144 let value: Option<String> = conn
145 .query_row(
146 &format!("SELECT CAST(\"value\" AS TEXT) FROM \"{table}\" WHERE \"key\"=?1"),
147 [key],
148 |row| row.get(0),
149 )
150 .optional()?;
151 match value {
152 None => Ok(0),
153 Some(value) => value
154 .parse::<i64>()
155 .ok()
156 .filter(|version| *version >= 0)
157 .ok_or_else(|| DbError::InvalidMigrations(format!("invalid stored version {value:?}"))),
158 }
159}
160
161fn set_version(conn: &Connection, storage: &VersionStorage, version: i64) -> rusqlite::Result<()> {
163 match storage {
164 VersionStorage::UserVersion => set_user_version(conn, version),
165 VersionStorage::Table { table, key } => {
166 conn.execute_batch(&format!(
167 "CREATE TABLE IF NOT EXISTS \"{table}\" (\"key\" TEXT PRIMARY KEY, \"value\" TEXT NOT NULL);"
168 ))?;
169 conn.execute(
170 &format!(
171 "INSERT INTO \"{table}\"(\"key\",\"value\") VALUES(?1,?2) \
172 ON CONFLICT(\"key\") DO UPDATE SET \"value\"=excluded.\"value\""
173 ),
174 rusqlite::params![key, version.to_string()],
175 )?;
176 Ok(())
177 }
178 }
179}
180
181pub fn validate(migrations: &[Migration]) -> Result<()> {
182 for (index, migration) in migrations.iter().enumerate() {
183 let expected = index as i64 + 1;
184 if migration.version != expected {
185 return Err(DbError::InvalidMigrations(format!(
186 "expected version {expected} at position {index}, found {} ({})",
187 migration.version, migration.name
188 )));
189 }
190 }
191 Ok(())
192}
193
194fn marker_path(db: &Path) -> PathBuf {
195 let mut name = db.as_os_str().to_os_string();
196 name.push(".migrating");
197 PathBuf::from(name)
198}
199
200pub fn pending_marker(db: &Path) -> Option<PendingMarker> {
201 let text = std::fs::read_to_string(marker_path(db)).ok()?;
202 let mut marker = PendingMarker {
203 from: 0,
204 to: 0,
205 pid: 0,
206 };
207 for line in text.lines() {
208 match line.split_once('=')? {
209 ("from", v) => marker.from = v.parse().ok()?,
210 ("to", v) => marker.to = v.parse().ok()?,
211 ("pid", v) => marker.pid = v.parse().ok()?,
212 _ => {}
213 }
214 }
215 Some(marker)
216}
217
218pub fn migrate(
219 conn: &Connection,
220 db_path: Option<&Path>,
221 migrations: &[Migration],
222 options: &MigrateOptions,
223) -> Result<MigrationReport> {
224 validate(migrations)?;
225 let current = migration_version(conn, &options.version_storage)?;
226 let max = migrations.last().map(|m| m.version).unwrap_or(0);
227 let mut report = MigrationReport {
228 from: current,
229 to: current,
230 ..Default::default()
231 };
232 report.interrupted_previous = db_path.and_then(pending_marker);
233
234 if current > max {
235 if options.forward_compatible.is_some_and(|probe| probe(conn)) {
236 report.forward_compatible = true;
237 return Ok(report);
238 }
239 return Err(DbError::NewerDatabase { current, max });
240 }
241 let pending: Vec<&Migration> = migrations.iter().filter(|m| m.version > current).collect();
242 if pending.is_empty() {
243 return Ok(report);
244 }
245
246 if current > 0 && options.backup {
247 if let Some(path) = db_path {
248 let name = path
249 .file_name()
250 .ok_or(DbError::InvalidDatabasePath)?
251 .to_string_lossy();
252 let backup = path.with_file_name(format!("{name}.pre-v{max}.bak"));
253 let backup_options = options
254 .backup_options
255 .clone()
256 .unwrap_or_else(|| DbOptions::new(path));
257 crate::online_backup_to(&backup_options, conn, &backup, None)?;
258 report.backup = Some(backup);
259 }
260 }
261
262 if let Some(path) = db_path {
263 let body = format!("from={current}\nto={max}\npid={}\n", std::process::id());
264 rightkit_fs::atomic_write(&marker_path(path), body.as_bytes())?;
265 }
266
267 for migration in pending {
268 conn.execute_batch("BEGIN IMMEDIATE;")?;
269 let applied = match &migration.up {
270 Step::Sql(sql) => conn.execute_batch(sql),
271 Step::Code(run) => run(conn),
272 }
273 .and_then(|()| set_version(conn, &options.version_storage, migration.version));
274 match applied {
275 Ok(()) => {
276 conn.execute_batch("COMMIT;")?;
277 report.applied.push(migration.version);
278 report.to = migration.version;
279 }
280 Err(source) => {
281 let _ = conn.execute_batch("ROLLBACK;");
282 return Err(DbError::MigrationFailed {
283 version: migration.version,
284 name: migration.name,
285 source,
286 });
287 }
288 }
289 }
290
291 if let Some(path) = db_path {
292 let _ = std::fs::remove_file(marker_path(path));
293 }
294 if options.optimize_after {
295 crate::optimize_boundary(conn)?;
296 }
297 Ok(report)
298}
299
300pub fn rollback_to(conn: &Connection, migrations: &[Migration], target: i64) -> Result<Vec<i64>> {
303 rollback_to_with_storage(conn, migrations, target, &VersionStorage::UserVersion)
304}
305
306pub fn rollback_to_with_storage(
308 conn: &Connection,
309 migrations: &[Migration],
310 target: i64,
311 storage: &VersionStorage,
312) -> Result<Vec<i64>> {
313 validate(migrations)?;
314 if target < 0 {
315 return Err(DbError::Rollback(
316 "target version must be nonnegative".into(),
317 ));
318 }
319 let current = migration_version(conn, storage)?;
320 if current <= target {
321 return Err(DbError::Rollback(format!(
322 "already at v{current}, at or below target v{target}"
323 )));
324 }
325 let mut plan = Vec::new();
326 for version in ((target + 1)..=current).rev() {
327 let migration = migrations
328 .iter()
329 .find(|m| m.version == version)
330 .ok_or_else(|| DbError::Rollback(format!("v{version} is not a known migration")))?;
331 let down = migration.down.ok_or_else(|| {
332 DbError::Rollback(format!(
333 "v{version} ({}) has no down migration",
334 migration.name
335 ))
336 })?;
337 plan.push((version, down));
338 }
339 let mut reverted = Vec::new();
340 for (version, down) in plan {
341 conn.execute_batch("BEGIN IMMEDIATE;")?;
342 let result = conn
343 .execute_batch(down)
344 .and_then(|()| set_version(conn, storage, version - 1));
345 match result {
346 Ok(()) => {
347 conn.execute_batch("COMMIT;")?;
348 reverted.push(version);
349 }
350 Err(error) => {
351 let _ = conn.execute_batch("ROLLBACK;");
352 return Err(DbError::Rollback(format!(
353 "down v{version} failed and was rolled back: {error}"
354 )));
355 }
356 }
357 }
358 Ok(reverted)
359}
360
361pub fn column_exists(conn: &Connection, table: &str, column: &str) -> Result<bool> {
364 check_ident(table)?;
365 let mut statement = conn.prepare(&format!("PRAGMA table_xinfo({table})"))?;
366 let names = statement.query_map([], |row| row.get::<_, String>(1))?;
367 for name in names {
368 if name? == column {
369 return Ok(true);
370 }
371 }
372 Ok(false)
373}
374
375pub fn add_column_if_missing(
378 conn: &Connection,
379 table: &str,
380 column: &str,
381 decl: &str,
382) -> Result<()> {
383 check_ident(table)?;
384 check_ident(column)?;
385 if !column_exists(conn, table, column)? {
386 conn.execute_batch(&format!("ALTER TABLE {table} ADD COLUMN {column} {decl};"))?;
387 }
388 Ok(())
389}
390
391fn check_ident(ident: &str) -> Result<()> {
392 if !ident.is_empty()
393 && ident
394 .bytes()
395 .all(|b| b.is_ascii_alphanumeric() || b == b'_')
396 {
397 Ok(())
398 } else {
399 Err(DbError::InvalidMigrations(format!(
400 "unsafe SQL identifier {ident:?}"
401 )))
402 }
403}
404
405impl DbOptions {
406 pub fn open_writer_migrated(
409 &self,
410 migrations: &[Migration],
411 options: &MigrateOptions,
412 ) -> Result<(Connection, MigrationReport)> {
413 let conn = self.open_writer()?;
414 let mut options = options.clone();
415 options.backup_options.get_or_insert_with(|| self.clone());
416 let report = migrate(&conn, Some(&self.path), migrations, &options)?;
417 Ok((conn, report))
418 }
419}