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}