Skip to main content

storekit/
migrate.rs

1use rusqlite::{Connection, OptionalExtension, params};
2
3const SCHEMA_META_SQL: &str =
4    "CREATE TABLE IF NOT EXISTS schema_meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);";
5
6#[derive(Debug, Clone, Copy)]
7pub struct Step {
8    pub to: i64,
9    pub sql: &'static str,
10    pub backfill: Option<fn(&Connection) -> rusqlite::Result<()>>,
11    foreign_keys_off: bool,
12}
13
14impl Step {
15    pub const fn new(
16        to: i64,
17        sql: &'static str,
18        backfill: Option<fn(&Connection) -> rusqlite::Result<()>>,
19    ) -> Self {
20        Self {
21            to,
22            sql,
23            backfill,
24            foreign_keys_off: false,
25        }
26    }
27
28    pub const fn with_foreign_keys_off(mut self) -> Self {
29        self.foreign_keys_off = true;
30        self
31    }
32}
33
34#[derive(Debug)]
35pub struct Migrations {
36    pub version: i64,
37    pub steps: &'static [Step],
38    pub before: Option<fn(&Connection) -> rusqlite::Result<()>>,
39}
40
41pub fn run(connection: &Connection, migrations: &Migrations) -> rusqlite::Result<()> {
42    connection.execute_batch(SCHEMA_META_SQL)?;
43    let current = stored_version(connection)?;
44
45    if current > 0
46        && let Some(before) = migrations.before
47    {
48        before(connection)?;
49    }
50
51    for step in migrations.steps.iter().filter(|step| step.to > current) {
52        run_step(connection, step)?;
53    }
54
55    if current != migrations.version {
56        connection.execute(
57            "INSERT INTO schema_meta(key, value) VALUES ('version', ?1)
58             ON CONFLICT(key) DO UPDATE SET value = excluded.value",
59            params![migrations.version.to_string()],
60        )?;
61    }
62    Ok(())
63}
64
65fn run_step(connection: &Connection, step: &Step) -> rusqlite::Result<()> {
66    if !step.foreign_keys_off {
67        connection.execute_batch(step.sql)?;
68    } else {
69        connection.pragma_update(None, "foreign_keys", "OFF")?;
70        let rebuilt = connection.execute_batch(step.sql);
71        let restored = connection.pragma_update(None, "foreign_keys", "ON");
72        rebuilt?;
73        restored?;
74    }
75    if let Some(backfill) = step.backfill {
76        backfill(connection)?;
77    }
78    Ok(())
79}
80
81fn stored_version(connection: &Connection) -> rusqlite::Result<i64> {
82    Ok(connection
83        .query_row(
84            "SELECT value FROM schema_meta WHERE key = 'version'",
85            [],
86            |row| row.get::<_, String>(0),
87        )
88        .optional()?
89        .and_then(|value| value.parse().ok())
90        .unwrap_or(0))
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    static OUT_OF_NUMBER_ORDER: &[Step] = &[
98        Step::new(
99            10,
100            "CREATE TABLE order_log(value INTEGER); INSERT INTO order_log VALUES (10);",
101            None,
102        ),
103        Step::new(8, "INSERT INTO order_log VALUES (8);", None),
104        Step::new(2, "INSERT INTO order_log VALUES (2);", None),
105    ];
106    static REBUILD_STEPS: &[Step] =
107        &[Step::new(1, "CREATE TABLE rebuilt(id INTEGER);", None).with_foreign_keys_off()];
108
109    #[test]
110    fn the_runner_honours_listed_order_and_records_the_version_once() {
111        let connection = Connection::open_in_memory().unwrap();
112        connection
113            .execute_batch(
114                "CREATE TABLE schema_meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
115                 INSERT INTO schema_meta VALUES ('version', '0');
116                 CREATE TABLE version_writes(count INTEGER NOT NULL);
117                 INSERT INTO version_writes VALUES (0);
118                 CREATE TRIGGER count_version_writes AFTER UPDATE ON schema_meta
119                 WHEN new.key = 'version'
120                 BEGIN UPDATE version_writes SET count = count + 1; END;",
121            )
122            .unwrap();
123        let migrations = Migrations {
124            version: 10,
125            steps: OUT_OF_NUMBER_ORDER,
126            before: None,
127        };
128
129        run(&connection, &migrations).unwrap();
130
131        let listed = connection
132            .prepare("SELECT value FROM order_log ORDER BY rowid")
133            .unwrap()
134            .query_map([], |row| row.get::<_, i64>(0))
135            .unwrap()
136            .collect::<rusqlite::Result<Vec<_>>>()
137            .unwrap();
138        assert_eq!(listed, vec![10, 8, 2]);
139        assert_eq!(stored_version(&connection).unwrap(), 10);
140        assert_eq!(
141            connection
142                .query_row("SELECT count FROM version_writes", [], |row| row
143                    .get::<_, i64>(0))
144                .unwrap(),
145            1,
146            "the target version is recorded only once after every listed step"
147        );
148    }
149
150    #[test]
151    fn a_rebuild_restores_foreign_keys() {
152        let connection = Connection::open_in_memory().unwrap();
153        connection
154            .pragma_update(None, "foreign_keys", "ON")
155            .unwrap();
156        let migrations = Migrations {
157            version: 1,
158            steps: REBUILD_STEPS,
159            before: None,
160        };
161
162        run(&connection, &migrations).unwrap();
163
164        assert!(
165            connection
166                .query_row("PRAGMA foreign_keys", [], |row| row.get::<_, bool>(0))
167                .unwrap()
168        );
169    }
170}