1use crate::{DbError, DbOptions, Result};
17use rusqlite::Connection;
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)]
64pub struct MigrateOptions {
65 pub backup: bool,
67 pub backup_options: Option<DbOptions>,
69 pub forward_compatible: Option<fn(&Connection) -> bool>,
72 pub optimize_after: bool,
74}
75
76impl Default for MigrateOptions {
77 fn default() -> Self {
78 Self {
79 backup: true,
80 backup_options: None,
81 forward_compatible: None,
82 optimize_after: true,
83 }
84 }
85}
86
87#[derive(Clone, Debug, Default, PartialEq, Eq)]
88pub struct MigrationReport {
89 pub from: i64,
90 pub to: i64,
91 pub applied: Vec<i64>,
92 pub backup: Option<PathBuf>,
93 pub forward_compatible: bool,
95 pub interrupted_previous: Option<PendingMarker>,
97}
98
99#[derive(Clone, Debug, PartialEq, Eq)]
100pub struct PendingMarker {
101 pub from: i64,
102 pub to: i64,
103 pub pid: u32,
104}
105
106pub fn user_version(conn: &Connection) -> Result<i64> {
107 Ok(conn.query_row("PRAGMA user_version", [], |row| row.get(0))?)
108}
109
110fn set_user_version(conn: &Connection, version: i64) -> rusqlite::Result<()> {
111 conn.execute_batch(&format!("PRAGMA user_version = {version};"))
113}
114
115pub fn validate(migrations: &[Migration]) -> Result<()> {
116 for (index, migration) in migrations.iter().enumerate() {
117 let expected = index as i64 + 1;
118 if migration.version != expected {
119 return Err(DbError::InvalidMigrations(format!(
120 "expected version {expected} at position {index}, found {} ({})",
121 migration.version, migration.name
122 )));
123 }
124 }
125 Ok(())
126}
127
128fn marker_path(db: &Path) -> PathBuf {
129 let mut name = db.as_os_str().to_os_string();
130 name.push(".migrating");
131 PathBuf::from(name)
132}
133
134pub fn pending_marker(db: &Path) -> Option<PendingMarker> {
135 let text = std::fs::read_to_string(marker_path(db)).ok()?;
136 let mut marker = PendingMarker {
137 from: 0,
138 to: 0,
139 pid: 0,
140 };
141 for line in text.lines() {
142 match line.split_once('=')? {
143 ("from", v) => marker.from = v.parse().ok()?,
144 ("to", v) => marker.to = v.parse().ok()?,
145 ("pid", v) => marker.pid = v.parse().ok()?,
146 _ => {}
147 }
148 }
149 Some(marker)
150}
151
152pub fn migrate(
153 conn: &Connection,
154 db_path: Option<&Path>,
155 migrations: &[Migration],
156 options: &MigrateOptions,
157) -> Result<MigrationReport> {
158 validate(migrations)?;
159 let current = user_version(conn)?;
160 let max = migrations.last().map(|m| m.version).unwrap_or(0);
161 let mut report = MigrationReport {
162 from: current,
163 to: current,
164 ..Default::default()
165 };
166 report.interrupted_previous = db_path.and_then(pending_marker);
167
168 if current > max {
169 if options.forward_compatible.is_some_and(|probe| probe(conn)) {
170 report.forward_compatible = true;
171 return Ok(report);
172 }
173 return Err(DbError::NewerDatabase { current, max });
174 }
175 let pending: Vec<&Migration> = migrations.iter().filter(|m| m.version > current).collect();
176 if pending.is_empty() {
177 return Ok(report);
178 }
179
180 if current > 0 && options.backup {
181 if let Some(path) = db_path {
182 let name = path
183 .file_name()
184 .ok_or(DbError::InvalidDatabasePath)?
185 .to_string_lossy();
186 let backup = path.with_file_name(format!("{name}.pre-v{max}.bak"));
187 let backup_options = options
188 .backup_options
189 .clone()
190 .unwrap_or_else(|| DbOptions::new(path));
191 crate::online_backup_to(&backup_options, conn, &backup, None)?;
192 report.backup = Some(backup);
193 }
194 }
195
196 if let Some(path) = db_path {
197 let body = format!("from={current}\nto={max}\npid={}\n", std::process::id());
198 rightkit_fs::atomic_write(&marker_path(path), body.as_bytes())?;
199 }
200
201 for migration in pending {
202 conn.execute_batch("BEGIN IMMEDIATE;")?;
203 let applied = match &migration.up {
204 Step::Sql(sql) => conn.execute_batch(sql),
205 Step::Code(run) => run(conn),
206 }
207 .and_then(|()| set_user_version(conn, migration.version));
208 match applied {
209 Ok(()) => {
210 conn.execute_batch("COMMIT;")?;
211 report.applied.push(migration.version);
212 report.to = migration.version;
213 }
214 Err(source) => {
215 let _ = conn.execute_batch("ROLLBACK;");
216 return Err(DbError::MigrationFailed {
217 version: migration.version,
218 name: migration.name,
219 source,
220 });
221 }
222 }
223 }
224
225 if let Some(path) = db_path {
226 let _ = std::fs::remove_file(marker_path(path));
227 }
228 if options.optimize_after {
229 crate::optimize_boundary(conn)?;
230 }
231 Ok(report)
232}
233
234pub fn rollback_to(conn: &Connection, migrations: &[Migration], target: i64) -> Result<Vec<i64>> {
237 validate(migrations)?;
238 let current = user_version(conn)?;
239 if current <= target {
240 return Err(DbError::Rollback(format!(
241 "already at v{current}, at or below target v{target}"
242 )));
243 }
244 let mut plan = Vec::new();
245 for version in ((target + 1)..=current).rev() {
246 let migration = migrations
247 .iter()
248 .find(|m| m.version == version)
249 .ok_or_else(|| DbError::Rollback(format!("v{version} is not a known migration")))?;
250 let down = migration.down.ok_or_else(|| {
251 DbError::Rollback(format!(
252 "v{version} ({}) has no down migration",
253 migration.name
254 ))
255 })?;
256 plan.push((version, down));
257 }
258 let mut reverted = Vec::new();
259 for (version, down) in plan {
260 conn.execute_batch("BEGIN IMMEDIATE;")?;
261 let result = conn
262 .execute_batch(down)
263 .and_then(|()| set_user_version(conn, version - 1));
264 match result {
265 Ok(()) => {
266 conn.execute_batch("COMMIT;")?;
267 reverted.push(version);
268 }
269 Err(error) => {
270 let _ = conn.execute_batch("ROLLBACK;");
271 return Err(DbError::Rollback(format!(
272 "down v{version} failed and was rolled back: {error}"
273 )));
274 }
275 }
276 }
277 Ok(reverted)
278}
279
280pub fn column_exists(conn: &Connection, table: &str, column: &str) -> Result<bool> {
283 check_ident(table)?;
284 let mut statement = conn.prepare(&format!("PRAGMA table_xinfo({table})"))?;
285 let names = statement.query_map([], |row| row.get::<_, String>(1))?;
286 for name in names {
287 if name? == column {
288 return Ok(true);
289 }
290 }
291 Ok(false)
292}
293
294pub fn add_column_if_missing(
297 conn: &Connection,
298 table: &str,
299 column: &str,
300 decl: &str,
301) -> Result<()> {
302 check_ident(table)?;
303 check_ident(column)?;
304 if !column_exists(conn, table, column)? {
305 conn.execute_batch(&format!("ALTER TABLE {table} ADD COLUMN {column} {decl};"))?;
306 }
307 Ok(())
308}
309
310fn check_ident(ident: &str) -> Result<()> {
311 if !ident.is_empty()
312 && ident
313 .bytes()
314 .all(|b| b.is_ascii_alphanumeric() || b == b'_')
315 {
316 Ok(())
317 } else {
318 Err(DbError::InvalidMigrations(format!(
319 "unsafe SQL identifier {ident:?}"
320 )))
321 }
322}
323
324impl DbOptions {
325 pub fn open_writer_migrated(
328 &self,
329 migrations: &[Migration],
330 options: &MigrateOptions,
331 ) -> Result<(Connection, MigrationReport)> {
332 let conn = self.open_writer()?;
333 let mut options = options.clone();
334 options.backup_options.get_or_insert_with(|| self.clone());
335 let report = migrate(&conn, Some(&self.path), migrations, &options)?;
336 Ok((conn, report))
337 }
338}