1#[cfg(feature = "sqlite")]
7use std::collections::BTreeSet;
8use std::collections::{HashMap, HashSet};
9use std::future::Future;
10use std::path::Path;
11#[cfg(feature = "sqlite")]
12use std::sync::{Mutex, PoisonError};
13use std::time::Duration;
14#[cfg(any(feature = "postgres", feature = "sqlite"))]
15use std::time::Instant;
16
17use crate::backend;
18use crate::database::Pool;
19use crate::schema::Statement;
20use crate::transaction::{Target, Transaction};
21use crate::{Database, Driver, Error, Value};
22
23#[derive(Debug, Clone)]
26pub struct Migration {
27 pub name: String,
29 pub up: String,
32 pub down: String,
34 schema: Option<SchemaBody>,
36}
37
38#[derive(Debug, Clone)]
39struct SchemaBody {
40 up: Vec<Statement>,
41 down: Vec<Statement>,
42}
43
44impl Migration {
45 pub fn new(name: impl Into<String>, up: impl Into<String>, down: impl Into<String>) -> Self {
47 Self {
48 name: name.into(),
49 up: up.into(),
50 down: down.into(),
51 schema: None,
52 }
53 }
54
55 pub fn schema(name: impl Into<String>, up: Vec<Statement>, down: Vec<Statement>) -> Self {
60 Self {
61 name: name.into(),
62 up: String::new(),
63 down: String::new(),
64 schema: Some(SchemaBody { up, down }),
65 }
66 }
67
68 fn up_scripts(&self, driver: Driver) -> Result<Vec<String>, Error> {
70 match &self.schema {
71 Some(body) => render(&body.up, driver),
72 None => Ok(vec![self.up.clone()]),
73 }
74 }
75
76 fn down_scripts(&self, driver: Driver) -> Result<Vec<String>, Error> {
78 match &self.schema {
79 Some(body) => render(&body.down, driver),
80 None => Ok(vec![self.down.clone()]),
81 }
82 }
83}
84
85fn render(statements: &[Statement], driver: Driver) -> Result<Vec<String>, Error> {
86 let mut sql = Vec::new();
87 for statement in statements {
88 sql.extend(statement.try_to_sql(driver)?);
89 }
90 Ok(sql)
91}
92
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
95pub enum State {
96 Applied { batch: i64 },
97 Pending,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct Status {
103 pub name: String,
104 pub state: State,
105}
106
107#[derive(Debug, Clone)]
110pub struct Migrator {
111 migrations: Vec<Migration>,
112 lock_timeout: Duration,
113}
114
115const LOCK_PREFIX: &str = "orbit_migrations";
118
119const DEFAULT_LOCK_TIMEOUT: Duration = Duration::from_secs(10);
121
122#[cfg(any(feature = "postgres", feature = "sqlite"))]
124const LOCK_POLL: Duration = Duration::from_millis(50);
125
126const TABLE: &str = "__diesel_schema_migrations";
128
129#[cfg(feature = "sqlite")]
132static SQLITE_HELD: Mutex<BTreeSet<String>> = Mutex::new(BTreeSet::new());
133
134impl Migrator {
135 pub fn new(mut migrations: Vec<Migration>) -> Result<Self, Error> {
139 for migration in &migrations {
140 if !has_timestamp(&migration.name) {
141 return Err(Error::InvalidMigrationName {
142 name: migration.name.clone(),
143 });
144 }
145 }
146 migrations.sort_by(|a, b| a.name.cmp(&b.name));
147 for pair in migrations.windows(2) {
148 if pair[0].name == pair[1].name {
149 return Err(Error::DuplicateMigration {
150 name: pair[0].name.clone(),
151 });
152 }
153 if timestamp(&pair[0].name) == timestamp(&pair[1].name) {
154 return Err(Error::DuplicateTimestamp {
155 first: pair[0].name.clone(),
156 second: pair[1].name.clone(),
157 });
158 }
159 }
160 Ok(Self {
161 migrations,
162 lock_timeout: DEFAULT_LOCK_TIMEOUT,
163 })
164 }
165
166 pub fn with_lock_timeout(mut self, timeout: Duration) -> Self {
168 self.lock_timeout = timeout;
169 self
170 }
171
172 pub fn migrations(&self) -> &[Migration] {
174 &self.migrations
175 }
176
177 pub fn from_dir(dir: impl AsRef<Path>) -> Result<Self, Error> {
180 let dir = dir.as_ref();
181 let entries = std::fs::read_dir(dir).map_err(|source| Error::Read {
182 path: dir.to_path_buf(),
183 source,
184 })?;
185 let mut migrations = Vec::new();
186 for entry in entries {
187 let entry = entry.map_err(|source| Error::Read {
188 path: dir.to_path_buf(),
189 source,
190 })?;
191 let name = entry.file_name().to_string_lossy().into_owned();
192 let path = entry.path();
193 if name.starts_with('.') || !path.is_dir() {
194 continue;
195 }
196 let up = read_file(&path.join("up.sql"))?;
197 let down = read_file(&path.join("down.sql"))?;
198 migrations.push(Migration::new(name, up, down));
199 }
200 Self::new(migrations)
201 }
202
203 pub async fn run(&self, db: &Database) -> Result<Vec<String>, Error> {
215 self.with_lock(db, self.run_unlocked(db)).await
216 }
217
218 async fn run_unlocked(&self, db: &Database) -> Result<Vec<String>, Error> {
219 ensure_table(db).await?;
220 let applied = applied(db).await?;
221 let pending: Vec<&Migration> = self
222 .migrations
223 .iter()
224 .filter(|m| !applied.contains_key(timestamp(&m.name)))
225 .collect();
226 if pending.is_empty() {
227 return Ok(Vec::new());
228 }
229
230 let max: Option<i32> = backend::scalar(
231 Target::Pool(&db.pool),
232 &format!("SELECT MAX(batch) FROM {TABLE}"),
233 &[],
234 )
235 .await?;
236 let batch = i64::from(max.unwrap_or(0)) + 1;
237 let mut names = Vec::with_capacity(pending.len());
238 for migration in pending {
239 let scripts = migration.up_scripts(db.driver())?;
240 let mut tx = Transaction::begin(&db.pool).await?;
241 for script in &scripts {
242 backend::run_script(Target::Tx(&mut tx.inner), script)
243 .await
244 .map_err(|source| Error::MigrationFailed {
245 name: migration.name.clone(),
246 source,
247 })?;
248 }
249 let (sql, values) = prepared(
250 db.driver(),
251 &format!("INSERT INTO {TABLE} (version, batch) VALUES (?, ?)"),
252 vec![
253 Value::Text(timestamp(&migration.name).to_owned()),
254 Value::Int(batch),
255 ],
256 );
257 backend::execute(Target::Tx(&mut tx.inner), &sql, &values)
258 .await
259 .map_err(|source| Error::MigrationFailed {
260 name: migration.name.clone(),
261 source,
262 })?;
263 tx.commit().await?;
264 names.push(migration.name.clone());
265 }
266 Ok(names)
267 }
268
269 pub async fn status(&self, db: &Database) -> Result<Vec<Status>, Error> {
272 ensure_table(db).await?;
273 let applied = applied(db).await?;
274 Ok(self
275 .migrations
276 .iter()
277 .map(|m| Status {
278 name: m.name.clone(),
279 state: match applied.get(timestamp(&m.name)) {
280 Some(&batch) => State::Applied { batch },
281 None => State::Pending,
282 },
283 })
284 .collect())
285 }
286
287 pub async fn rollback(&self, db: &Database) -> Result<Vec<String>, Error> {
296 self.with_lock(db, self.rollback_unlocked(db)).await
297 }
298
299 async fn rollback_unlocked(&self, db: &Database) -> Result<Vec<String>, Error> {
300 ensure_table(db).await?;
301 let last: Option<i32> = backend::scalar(
302 Target::Pool(&db.pool),
303 &format!("SELECT MAX(batch) FROM {TABLE}"),
304 &[],
305 )
306 .await?;
307 let Some(batch) = last else {
308 return Ok(Vec::new());
309 };
310 let (sql, values) = prepared(
311 db.driver(),
312 &format!("SELECT version FROM {TABLE} WHERE batch = ? ORDER BY version DESC"),
313 vec![Value::Int(i64::from(batch))],
314 );
315 let versions: Vec<(String,)> =
316 backend::fetch_all(Target::Pool(&db.pool), &sql, &values).await?;
317
318 let mut to_revert = Vec::with_capacity(versions.len());
319 for (version,) in &versions {
320 let migration = self
321 .migrations
322 .iter()
323 .find(|m| timestamp(&m.name) == version.as_str())
324 .ok_or_else(|| Error::UnknownMigration {
325 name: version.clone(),
326 })?;
327 to_revert.push(migration);
328 }
329 let names = to_revert.iter().map(|m| m.name.clone()).collect();
330 for migration in to_revert {
331 let scripts = migration.down_scripts(db.driver())?;
332 let mut tx = Transaction::begin(&db.pool).await?;
333 for script in &scripts {
334 backend::run_script(Target::Tx(&mut tx.inner), script)
335 .await
336 .map_err(|source| Error::RollbackFailed {
337 name: migration.name.clone(),
338 source,
339 })?;
340 }
341 let (sql, values) = prepared(
342 db.driver(),
343 &format!("DELETE FROM {TABLE} WHERE version = ?"),
344 vec![Value::Text(timestamp(&migration.name).to_owned())],
345 );
346 backend::execute(Target::Tx(&mut tx.inner), &sql, &values)
347 .await
348 .map_err(|source| Error::RollbackFailed {
349 name: migration.name.clone(),
350 source,
351 })?;
352 tx.commit().await?;
353 }
354 Ok(names)
355 }
356}
357
358impl Migrator {
359 #[deprecated(
365 since = "0.1.1",
366 note = "not needed any more: run, status and rollback upgrade the migrations table \
367 themselves. This returns an empty list and will be removed in a later release"
368 )]
369 pub async fn import_diesel_history(&self, db: &Database) -> Result<Vec<String>, Error> {
370 self.with_lock(db, import_unlocked(db)).await
371 }
372
373 pub async fn unmatched_diesel_versions(&self, db: &Database) -> Result<Vec<String>, Error> {
377 let Some(versions) = recorded_versions(db).await? else {
378 return Ok(Vec::new());
379 };
380 let mut unmatched: Vec<String> = versions
381 .into_iter()
382 .filter(|version| {
383 !self
384 .migrations
385 .iter()
386 .any(|m| timestamp(&m.name) == version.as_str())
387 })
388 .collect();
389 unmatched.sort();
390 Ok(unmatched)
391 }
392}
393
394#[macro_export]
402macro_rules! migrations {
403 ($dir:literal, $($name:literal),* $(,)?) => {
404 $crate::Migrator::new(vec![
405 $(
406 $crate::Migration::new(
407 $name,
408 include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/", $dir, "/", $name, "/up.sql")),
409 include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/", $dir, "/", $name, "/down.sql")),
410 ),
411 )*
412 ])
413 };
414}
415
416impl Migrator {
417 async fn with_lock<T>(
420 &self,
421 db: &Database,
422 body: impl Future<Output = Result<T, Error>>,
423 ) -> Result<T, Error> {
424 match &db.pool {
425 #[cfg(feature = "mysql")]
426 Pool::MySql(pool) => mysql_lock(pool, self.lock_timeout, body).await,
427 #[cfg(feature = "postgres")]
428 Pool::Postgres(pool) => postgres_lock(pool, self.lock_timeout, body).await,
429 #[cfg(feature = "sqlite")]
430 Pool::Sqlite(pool) => sqlite_lock(pool, self.lock_timeout, body).await,
431 }
432 }
433}
434
435#[cfg(feature = "mysql")]
438async fn mysql_lock<T>(
439 pool: &sqlx::MySqlPool,
440 timeout: Duration,
441 body: impl Future<Output = Result<T, Error>>,
442) -> Result<T, Error> {
443 let name = format!(
444 "{LOCK_PREFIX}:{}",
445 pool.connect_options().get_database().unwrap_or("")
446 );
447 let mut conn = pool.acquire().await?;
448 let got: Option<i64> = sqlx::query_scalar("SELECT GET_LOCK(?, ?)")
449 .bind(&name)
450 .bind(timeout.as_secs() as i64)
451 .fetch_one(&mut *conn)
452 .await?;
453 if got != Some(1) {
454 return Err(Error::MigrationsLocked {
455 seconds: timeout.as_secs(),
456 });
457 }
458
459 let result = body.await;
460 let released = sqlx::query("SELECT RELEASE_LOCK(?)")
461 .bind(&name)
462 .execute(&mut *conn)
463 .await;
464 if released.is_err() {
465 drop(conn.detach());
467 }
468 let value = result?;
469 released?;
470 Ok(value)
471}
472
473#[cfg(feature = "postgres")]
476async fn postgres_lock<T>(
477 pool: &sqlx::PgPool,
478 timeout: Duration,
479 body: impl Future<Output = Result<T, Error>>,
480) -> Result<T, Error> {
481 let name = format!(
482 "{LOCK_PREFIX}:{}",
483 pool.connect_options().get_database().unwrap_or("")
484 );
485 let key = advisory_key(&name);
486 let mut conn = pool.acquire().await?;
487 let deadline = Instant::now() + timeout;
488 loop {
489 let got: bool = sqlx::query_scalar("SELECT pg_try_advisory_lock($1)")
490 .bind(key)
491 .fetch_one(&mut *conn)
492 .await?;
493 if got {
494 break;
495 }
496 if Instant::now() >= deadline {
497 return Err(Error::MigrationsLocked {
498 seconds: timeout.as_secs(),
499 });
500 }
501 tokio::time::sleep(LOCK_POLL).await;
502 }
503
504 let result = body.await;
505 let released = sqlx::query_scalar::<_, bool>("SELECT pg_advisory_unlock($1)")
506 .bind(key)
507 .fetch_one(&mut *conn)
508 .await;
509 if released.is_err() {
510 drop(conn.detach());
512 }
513 let value = result?;
514 released?;
515 Ok(value)
516}
517
518#[cfg(feature = "sqlite")]
524async fn sqlite_lock<T>(
525 pool: &sqlx::SqlitePool,
526 timeout: Duration,
527 body: impl Future<Output = Result<T, Error>>,
528) -> Result<T, Error> {
529 let name = format!(
530 "{LOCK_PREFIX}:{}",
531 pool.connect_options().get_filename().display()
532 );
533 let deadline = Instant::now() + timeout;
534 loop {
535 let taken = SQLITE_HELD
536 .lock()
537 .unwrap_or_else(PoisonError::into_inner)
538 .insert(name.clone());
539 if taken {
540 break;
541 }
542 if Instant::now() >= deadline {
543 return Err(Error::MigrationsLocked {
544 seconds: timeout.as_secs(),
545 });
546 }
547 tokio::time::sleep(LOCK_POLL).await;
548 }
549
550 let result = body.await;
551 SQLITE_HELD
552 .lock()
553 .unwrap_or_else(PoisonError::into_inner)
554 .remove(&name);
555 result
556}
557
558#[cfg(feature = "postgres")]
561fn advisory_key(name: &str) -> i64 {
562 let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
563 for byte in name.bytes() {
564 hash ^= u64::from(byte);
565 hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
566 }
567 hash as i64
568}
569
570async fn import_unlocked(db: &Database) -> Result<Vec<String>, Error> {
572 ensure_table(db).await?;
573 Ok(Vec::new())
574}
575
576fn has_timestamp(name: &str) -> bool {
578 let bytes = name.as_bytes();
579 bytes.len() > 15 && bytes[..14].iter().all(u8::is_ascii_digit) && bytes[14] == b'_'
580}
581
582fn timestamp(name: &str) -> &str {
584 &name[..14]
585}
586
587fn prepared(driver: Driver, sql: &str, values: Vec<Value>) -> (String, Vec<Value>) {
589 driver.prepare(sql, &values)
590}
591
592#[cfg_attr(
594 not(all(feature = "mysql", feature = "postgres", feature = "sqlite")),
595 allow(unused_variables)
596)]
597fn pick<'a>(driver: Driver, mysql: &'a str, postgres: &'a str, sqlite: &'a str) -> &'a str {
598 match driver {
599 #[cfg(feature = "mysql")]
600 Driver::MySql => mysql,
601 #[cfg(feature = "postgres")]
602 Driver::Postgres => postgres,
603 #[cfg(feature = "sqlite")]
604 Driver::Sqlite => sqlite,
605 }
606}
607
608async fn ensure_table(db: &Database) -> Result<(), Error> {
613 let driver = db.driver();
614 let create = pick(
615 driver,
616 "CREATE TABLE IF NOT EXISTS __diesel_schema_migrations (\
617 version VARCHAR(50) NOT NULL PRIMARY KEY, \
618 run_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, \
619 batch INT NOT NULL DEFAULT 0)",
620 "CREATE TABLE IF NOT EXISTS __diesel_schema_migrations (\
621 version VARCHAR(50) NOT NULL PRIMARY KEY, \
622 run_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, \
623 batch INT NOT NULL DEFAULT 0)",
624 "CREATE TABLE IF NOT EXISTS __diesel_schema_migrations (\
625 version TEXT NOT NULL PRIMARY KEY, \
626 run_on TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, \
627 batch INTEGER NOT NULL DEFAULT 0)",
628 );
629 backend::execute(Target::Pool(&db.pool), create, &[]).await?;
630 if has_batch_column(db).await? {
631 return Ok(());
632 }
633 let added = backend::execute(
634 Target::Pool(&db.pool),
635 "ALTER TABLE __diesel_schema_migrations ADD COLUMN batch INT NOT NULL DEFAULT 0",
636 &[],
637 )
638 .await;
639 match added {
640 Ok(_) => Ok(()),
641 Err(err) if is_duplicate_column(&err) => Ok(()),
643 Err(err) => Err(err.into()),
644 }
645}
646
647async fn has_batch_column(db: &Database) -> Result<bool, Error> {
649 let sql = pick(
650 db.driver(),
651 "SELECT COUNT(*) FROM information_schema.columns \
652 WHERE table_schema = DATABASE() AND table_name = '__diesel_schema_migrations' \
653 AND column_name = 'batch'",
654 "SELECT COUNT(*) FROM information_schema.columns \
655 WHERE table_schema = current_schema() AND table_name = '__diesel_schema_migrations' \
656 AND column_name = 'batch'",
657 "SELECT COUNT(*) FROM pragma_table_info('__diesel_schema_migrations') \
658 WHERE name = 'batch'",
659 );
660 let found: i64 = backend::scalar(Target::Pool(&db.pool), sql, &[]).await?;
661 Ok(found > 0)
662}
663
664fn is_duplicate_column(err: &sqlx::Error) -> bool {
667 err.as_database_error().is_some_and(|db_err| {
668 matches!(db_err.code().as_deref(), Some("1060" | "42701"))
669 || db_err.message().contains("duplicate column name")
670 })
671}
672
673async fn applied(db: &Database) -> Result<HashMap<String, i64>, Error> {
675 let rows: Vec<(String, i32)> = backend::fetch_all(
676 Target::Pool(&db.pool),
677 &format!("SELECT version, batch FROM {TABLE}"),
678 &[],
679 )
680 .await?;
681 Ok(rows
682 .into_iter()
683 .map(|(version, batch)| (version, i64::from(batch)))
684 .collect())
685}
686
687async fn recorded_versions(db: &Database) -> Result<Option<HashSet<String>>, Error> {
690 let sql = pick(
691 db.driver(),
692 "SELECT COUNT(*) FROM information_schema.tables \
693 WHERE table_schema = DATABASE() AND table_name = '__diesel_schema_migrations'",
694 "SELECT COUNT(*) FROM information_schema.tables \
695 WHERE table_schema = current_schema() AND table_name = '__diesel_schema_migrations'",
696 "SELECT COUNT(*) FROM sqlite_master \
697 WHERE type = 'table' AND name = '__diesel_schema_migrations'",
698 );
699 let found: i64 = backend::scalar(Target::Pool(&db.pool), sql, &[]).await?;
700 if found == 0 {
701 return Ok(None);
702 }
703 let versions: Vec<(String,)> = backend::fetch_all(
704 Target::Pool(&db.pool),
705 &format!("SELECT version FROM {TABLE}"),
706 &[],
707 )
708 .await?;
709 Ok(Some(versions.into_iter().map(|(v,)| v).collect()))
710}
711
712pub(crate) fn read_file(path: &Path) -> Result<String, Error> {
714 std::fs::read_to_string(path).map_err(|source| Error::Read {
715 path: path.to_path_buf(),
716 source,
717 })
718}