Skip to main content

rightkit_db/
migrate.rs

1//! Ordered, transactional, crash-aware schema migrations.
2//!
3//! The library owns the *mechanism*; apps own the schemas. Rules, taken from the strongest
4//! donors (CodeRight `storage::migrations`, MailRight `db::migrations`, WorkRight `local-store`):
5//!
6//! - versions are contiguous from 1 and default to `PRAGMA user_version`;
7//! - each step runs in its own `BEGIN IMMEDIATE` transaction together with its version bump, so
8//!   a failing step rolls back to the last good version;
9//! - before the first pending step on a non-empty database an online backup is taken to
10//!   `<db>.pre-v<target>.bak` (WAL-inclusive, atomic, owner-only);
11//! - a `<db>.migrating` marker brackets the batch so a crash mid-batch is reported on next open;
12//! - a database newer than this build is refused unless the app's probe says the subset it
13//!   needs is still present (additive forward compatibility);
14//! - rollback replays `down` steps and refuses when any step in the range lacks one.
15
16use crate::{DbError, DbOptions, Result};
17use rusqlite::{Connection, OptionalExtension};
18use std::path::{Path, PathBuf};
19
20pub enum Step {
21    Sql(&'static str),
22    /// Rust step (data backfills, table rebuilds); runs inside the step's transaction.
23    Code(fn(&Connection) -> rusqlite::Result<()>),
24}
25
26pub struct Migration {
27    pub version: i64,
28    pub name: &'static str,
29    pub up: Step,
30    /// Reverse SQL; `None` makes this step unrevertable (rollback refuses to guess).
31    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/// Where migration versions are stored. Table storage uses Blueprint-compatible
64/// `key TEXT PRIMARY KEY, value TEXT NOT NULL` rows, leaving `user_version` untouched.
65#[derive(Clone, Debug, Default, PartialEq, Eq)]
66pub enum VersionStorage {
67    #[default]
68    UserVersion,
69    /// Table names must contain only ASCII letters, digits, or underscores.
70    /// Keys are bound parameters and may contain arbitrary text.
71    Table { table: String, key: String },
72}
73
74#[derive(Clone)]
75pub struct MigrateOptions {
76    /// Version location shared by reads and transactional migration writes.
77    pub version_storage: VersionStorage,
78    /// Take a pre-migration backup when migrating an existing (version > 0) on-disk database.
79    pub backup: bool,
80    /// Key source for the backup (SQLCipher); `DbOptions::open_writer_migrated` fills this in.
81    pub backup_options: Option<DbOptions>,
82    /// For a database newer than this build: `true` means the subset this build reads/writes is
83    /// intact, so open it read-write without migrating.
84    pub forward_compatible: Option<fn(&Connection) -> bool>,
85    /// Refresh planner statistics after migrating (`PRAGMA optimize`).
86    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    /// Opened a newer-but-compatible database without migrating.
108    pub forward_compatible: bool,
109    /// A previous batch started but never finished (crash or failed step).
110    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    // `user_version` cannot be bound as a parameter; the value is an i64, never text.
126    conn.execute_batch(&format!("PRAGMA user_version = {version};"))
127}
128
129/// Read the selected migration version without creating a table or row.
130/// Missing table/key means zero; malformed or negative versions are refused.
131pub 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
161// Called only inside the migration/down transaction, after identifier validation.
162fn 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
300/// Revert to `target` by replaying each step's `down`, highest first, one transaction per step.
301/// Refuses before touching anything unless every step in the range has a `down`.
302pub fn rollback_to(conn: &Connection, migrations: &[Migration], target: i64) -> Result<Vec<i64>> {
303    rollback_to_with_storage(conn, migrations, target, &VersionStorage::UserVersion)
304}
305
306/// Revert using the same version storage selected for [`migrate`].
307pub 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
361/// True if `table` has `column`, including STORED/VIRTUAL generated columns (`table_xinfo`;
362/// `table_info` omits them, which made MailRight re-run an ALTER on every open).
363pub 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
375/// Idempotent `ALTER TABLE .. ADD COLUMN` (SQLite has no `IF NOT EXISTS` for columns).
376/// `table`, `column` and `decl` must be literals from the app, never user input.
377pub 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    /// Open the writer and bring the schema up to date (see [`migrate`]); the pre-migration
407    /// backup is keyed like the database itself.
408    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}