#[cfg(feature = "sqlite")]
use std::collections::BTreeSet;
use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::path::Path;
#[cfg(feature = "sqlite")]
use std::sync::{Mutex, PoisonError};
use std::time::Duration;
#[cfg(any(feature = "postgres", feature = "sqlite"))]
use std::time::Instant;
use crate::backend;
use crate::database::Pool;
use crate::schema::Statement;
use crate::transaction::{Target, Transaction};
use crate::{Database, Driver, Error, Value};
#[derive(Debug, Clone)]
pub struct Migration {
pub name: String,
pub up: String,
pub down: String,
schema: Option<SchemaBody>,
}
#[derive(Debug, Clone)]
struct SchemaBody {
up: Vec<Statement>,
down: Vec<Statement>,
}
impl Migration {
pub fn new(name: impl Into<String>, up: impl Into<String>, down: impl Into<String>) -> Self {
Self {
name: name.into(),
up: up.into(),
down: down.into(),
schema: None,
}
}
pub fn schema(name: impl Into<String>, up: Vec<Statement>, down: Vec<Statement>) -> Self {
Self {
name: name.into(),
up: String::new(),
down: String::new(),
schema: Some(SchemaBody { up, down }),
}
}
fn up_scripts(&self, driver: Driver) -> Result<Vec<String>, Error> {
match &self.schema {
Some(body) => render(&body.up, driver),
None => Ok(vec![self.up.clone()]),
}
}
fn down_scripts(&self, driver: Driver) -> Result<Vec<String>, Error> {
match &self.schema {
Some(body) => render(&body.down, driver),
None => Ok(vec![self.down.clone()]),
}
}
}
fn render(statements: &[Statement], driver: Driver) -> Result<Vec<String>, Error> {
let mut sql = Vec::new();
for statement in statements {
sql.extend(statement.try_to_sql(driver)?);
}
Ok(sql)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum State {
Applied { batch: i64 },
Pending,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Status {
pub name: String,
pub state: State,
}
#[derive(Debug, Clone)]
pub struct Migrator {
migrations: Vec<Migration>,
lock_timeout: Duration,
}
const LOCK_PREFIX: &str = "orbit_migrations";
const DEFAULT_LOCK_TIMEOUT: Duration = Duration::from_secs(10);
#[cfg(any(feature = "postgres", feature = "sqlite"))]
const LOCK_POLL: Duration = Duration::from_millis(50);
const TABLE: &str = "__diesel_schema_migrations";
#[cfg(feature = "sqlite")]
static SQLITE_HELD: Mutex<BTreeSet<String>> = Mutex::new(BTreeSet::new());
impl Migrator {
pub fn new(mut migrations: Vec<Migration>) -> Result<Self, Error> {
for migration in &migrations {
if !has_timestamp(&migration.name) {
return Err(Error::InvalidMigrationName {
name: migration.name.clone(),
});
}
}
migrations.sort_by(|a, b| a.name.cmp(&b.name));
for pair in migrations.windows(2) {
if pair[0].name == pair[1].name {
return Err(Error::DuplicateMigration {
name: pair[0].name.clone(),
});
}
if timestamp(&pair[0].name) == timestamp(&pair[1].name) {
return Err(Error::DuplicateTimestamp {
first: pair[0].name.clone(),
second: pair[1].name.clone(),
});
}
}
Ok(Self {
migrations,
lock_timeout: DEFAULT_LOCK_TIMEOUT,
})
}
pub fn with_lock_timeout(mut self, timeout: Duration) -> Self {
self.lock_timeout = timeout;
self
}
pub fn migrations(&self) -> &[Migration] {
&self.migrations
}
pub fn from_dir(dir: impl AsRef<Path>) -> Result<Self, Error> {
let dir = dir.as_ref();
let entries = std::fs::read_dir(dir).map_err(|source| Error::Read {
path: dir.to_path_buf(),
source,
})?;
let mut migrations = Vec::new();
for entry in entries {
let entry = entry.map_err(|source| Error::Read {
path: dir.to_path_buf(),
source,
})?;
let name = entry.file_name().to_string_lossy().into_owned();
let path = entry.path();
if name.starts_with('.') || !path.is_dir() {
continue;
}
let up = read_file(&path.join("up.sql"))?;
let down = read_file(&path.join("down.sql"))?;
migrations.push(Migration::new(name, up, down));
}
Self::new(migrations)
}
pub async fn run(&self, db: &Database) -> Result<Vec<String>, Error> {
self.with_lock(db, self.run_unlocked(db)).await
}
async fn run_unlocked(&self, db: &Database) -> Result<Vec<String>, Error> {
ensure_table(db).await?;
let applied = applied(db).await?;
let pending: Vec<&Migration> = self
.migrations
.iter()
.filter(|m| !applied.contains_key(timestamp(&m.name)))
.collect();
if pending.is_empty() {
return Ok(Vec::new());
}
let max: Option<i32> = backend::scalar(
Target::Pool(&db.pool),
&format!("SELECT MAX(batch) FROM {TABLE}"),
&[],
)
.await?;
let batch = i64::from(max.unwrap_or(0)) + 1;
let mut names = Vec::with_capacity(pending.len());
for migration in pending {
let scripts = migration.up_scripts(db.driver())?;
let mut tx = Transaction::begin(&db.pool).await?;
for script in &scripts {
backend::run_script(Target::Tx(&mut tx.inner), script)
.await
.map_err(|source| Error::MigrationFailed {
name: migration.name.clone(),
source,
})?;
}
let (sql, values) = prepared(
db.driver(),
&format!("INSERT INTO {TABLE} (version, batch) VALUES (?, ?)"),
vec![
Value::Text(timestamp(&migration.name).to_owned()),
Value::Int(batch),
],
);
backend::execute(Target::Tx(&mut tx.inner), &sql, &values)
.await
.map_err(|source| Error::MigrationFailed {
name: migration.name.clone(),
source,
})?;
tx.commit().await?;
names.push(migration.name.clone());
}
Ok(names)
}
pub async fn status(&self, db: &Database) -> Result<Vec<Status>, Error> {
ensure_table(db).await?;
let applied = applied(db).await?;
Ok(self
.migrations
.iter()
.map(|m| Status {
name: m.name.clone(),
state: match applied.get(timestamp(&m.name)) {
Some(&batch) => State::Applied { batch },
None => State::Pending,
},
})
.collect())
}
pub async fn rollback(&self, db: &Database) -> Result<Vec<String>, Error> {
self.with_lock(db, self.rollback_unlocked(db)).await
}
async fn rollback_unlocked(&self, db: &Database) -> Result<Vec<String>, Error> {
ensure_table(db).await?;
let last: Option<i32> = backend::scalar(
Target::Pool(&db.pool),
&format!("SELECT MAX(batch) FROM {TABLE}"),
&[],
)
.await?;
let Some(batch) = last else {
return Ok(Vec::new());
};
let (sql, values) = prepared(
db.driver(),
&format!("SELECT version FROM {TABLE} WHERE batch = ? ORDER BY version DESC"),
vec![Value::Int(i64::from(batch))],
);
let versions: Vec<(String,)> =
backend::fetch_all(Target::Pool(&db.pool), &sql, &values).await?;
let mut to_revert = Vec::with_capacity(versions.len());
for (version,) in &versions {
let migration = self
.migrations
.iter()
.find(|m| timestamp(&m.name) == version.as_str())
.ok_or_else(|| Error::UnknownMigration {
name: version.clone(),
})?;
to_revert.push(migration);
}
let names = to_revert.iter().map(|m| m.name.clone()).collect();
for migration in to_revert {
let scripts = migration.down_scripts(db.driver())?;
let mut tx = Transaction::begin(&db.pool).await?;
for script in &scripts {
backend::run_script(Target::Tx(&mut tx.inner), script)
.await
.map_err(|source| Error::RollbackFailed {
name: migration.name.clone(),
source,
})?;
}
let (sql, values) = prepared(
db.driver(),
&format!("DELETE FROM {TABLE} WHERE version = ?"),
vec![Value::Text(timestamp(&migration.name).to_owned())],
);
backend::execute(Target::Tx(&mut tx.inner), &sql, &values)
.await
.map_err(|source| Error::RollbackFailed {
name: migration.name.clone(),
source,
})?;
tx.commit().await?;
}
Ok(names)
}
}
impl Migrator {
#[deprecated(
since = "0.1.1",
note = "not needed any more: run, status and rollback upgrade the migrations table \
themselves. This returns an empty list and will be removed in a later release"
)]
pub async fn import_diesel_history(&self, db: &Database) -> Result<Vec<String>, Error> {
self.with_lock(db, import_unlocked(db)).await
}
pub async fn unmatched_diesel_versions(&self, db: &Database) -> Result<Vec<String>, Error> {
let Some(versions) = recorded_versions(db).await? else {
return Ok(Vec::new());
};
let mut unmatched: Vec<String> = versions
.into_iter()
.filter(|version| {
!self
.migrations
.iter()
.any(|m| timestamp(&m.name) == version.as_str())
})
.collect();
unmatched.sort();
Ok(unmatched)
}
}
#[macro_export]
macro_rules! migrations {
($dir:literal, $($name:literal),* $(,)?) => {
$crate::Migrator::new(vec![
$(
$crate::Migration::new(
$name,
include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/", $dir, "/", $name, "/up.sql")),
include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/", $dir, "/", $name, "/down.sql")),
),
)*
])
};
}
impl Migrator {
async fn with_lock<T>(
&self,
db: &Database,
body: impl Future<Output = Result<T, Error>>,
) -> Result<T, Error> {
match &db.pool {
#[cfg(feature = "mysql")]
Pool::MySql(pool) => mysql_lock(pool, self.lock_timeout, body).await,
#[cfg(feature = "postgres")]
Pool::Postgres(pool) => postgres_lock(pool, self.lock_timeout, body).await,
#[cfg(feature = "sqlite")]
Pool::Sqlite(pool) => sqlite_lock(pool, self.lock_timeout, body).await,
}
}
}
#[cfg(feature = "mysql")]
async fn mysql_lock<T>(
pool: &sqlx::MySqlPool,
timeout: Duration,
body: impl Future<Output = Result<T, Error>>,
) -> Result<T, Error> {
let name = format!(
"{LOCK_PREFIX}:{}",
pool.connect_options().get_database().unwrap_or("")
);
let mut conn = pool.acquire().await?;
let got: Option<i64> = sqlx::query_scalar("SELECT GET_LOCK(?, ?)")
.bind(&name)
.bind(timeout.as_secs() as i64)
.fetch_one(&mut *conn)
.await?;
if got != Some(1) {
return Err(Error::MigrationsLocked {
seconds: timeout.as_secs(),
});
}
let result = body.await;
let released = sqlx::query("SELECT RELEASE_LOCK(?)")
.bind(&name)
.execute(&mut *conn)
.await;
if released.is_err() {
drop(conn.detach());
}
let value = result?;
released?;
Ok(value)
}
#[cfg(feature = "postgres")]
async fn postgres_lock<T>(
pool: &sqlx::PgPool,
timeout: Duration,
body: impl Future<Output = Result<T, Error>>,
) -> Result<T, Error> {
let name = format!(
"{LOCK_PREFIX}:{}",
pool.connect_options().get_database().unwrap_or("")
);
let key = advisory_key(&name);
let mut conn = pool.acquire().await?;
let deadline = Instant::now() + timeout;
loop {
let got: bool = sqlx::query_scalar("SELECT pg_try_advisory_lock($1)")
.bind(key)
.fetch_one(&mut *conn)
.await?;
if got {
break;
}
if Instant::now() >= deadline {
return Err(Error::MigrationsLocked {
seconds: timeout.as_secs(),
});
}
tokio::time::sleep(LOCK_POLL).await;
}
let result = body.await;
let released = sqlx::query_scalar::<_, bool>("SELECT pg_advisory_unlock($1)")
.bind(key)
.fetch_one(&mut *conn)
.await;
if released.is_err() {
drop(conn.detach());
}
let value = result?;
released?;
Ok(value)
}
#[cfg(feature = "sqlite")]
async fn sqlite_lock<T>(
pool: &sqlx::SqlitePool,
timeout: Duration,
body: impl Future<Output = Result<T, Error>>,
) -> Result<T, Error> {
let name = format!(
"{LOCK_PREFIX}:{}",
pool.connect_options().get_filename().display()
);
let deadline = Instant::now() + timeout;
loop {
let taken = SQLITE_HELD
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(name.clone());
if taken {
break;
}
if Instant::now() >= deadline {
return Err(Error::MigrationsLocked {
seconds: timeout.as_secs(),
});
}
tokio::time::sleep(LOCK_POLL).await;
}
let result = body.await;
SQLITE_HELD
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&name);
result
}
#[cfg(feature = "postgres")]
fn advisory_key(name: &str) -> i64 {
let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
for byte in name.bytes() {
hash ^= u64::from(byte);
hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
}
hash as i64
}
async fn import_unlocked(db: &Database) -> Result<Vec<String>, Error> {
ensure_table(db).await?;
Ok(Vec::new())
}
fn has_timestamp(name: &str) -> bool {
let bytes = name.as_bytes();
bytes.len() > 15 && bytes[..14].iter().all(u8::is_ascii_digit) && bytes[14] == b'_'
}
fn timestamp(name: &str) -> &str {
&name[..14]
}
fn prepared(driver: Driver, sql: &str, values: Vec<Value>) -> (String, Vec<Value>) {
driver.prepare(sql, &values)
}
#[cfg_attr(
not(all(feature = "mysql", feature = "postgres", feature = "sqlite")),
allow(unused_variables)
)]
fn pick<'a>(driver: Driver, mysql: &'a str, postgres: &'a str, sqlite: &'a str) -> &'a str {
match driver {
#[cfg(feature = "mysql")]
Driver::MySql => mysql,
#[cfg(feature = "postgres")]
Driver::Postgres => postgres,
#[cfg(feature = "sqlite")]
Driver::Sqlite => sqlite,
}
}
async fn ensure_table(db: &Database) -> Result<(), Error> {
let driver = db.driver();
let create = pick(
driver,
"CREATE TABLE IF NOT EXISTS __diesel_schema_migrations (\
version VARCHAR(50) NOT NULL PRIMARY KEY, \
run_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, \
batch INT NOT NULL DEFAULT 0)",
"CREATE TABLE IF NOT EXISTS __diesel_schema_migrations (\
version VARCHAR(50) NOT NULL PRIMARY KEY, \
run_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, \
batch INT NOT NULL DEFAULT 0)",
"CREATE TABLE IF NOT EXISTS __diesel_schema_migrations (\
version TEXT NOT NULL PRIMARY KEY, \
run_on TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, \
batch INTEGER NOT NULL DEFAULT 0)",
);
backend::execute(Target::Pool(&db.pool), create, &[]).await?;
if has_batch_column(db).await? {
return Ok(());
}
let added = backend::execute(
Target::Pool(&db.pool),
"ALTER TABLE __diesel_schema_migrations ADD COLUMN batch INT NOT NULL DEFAULT 0",
&[],
)
.await;
match added {
Ok(_) => Ok(()),
Err(err) if is_duplicate_column(&err) => Ok(()),
Err(err) => Err(err.into()),
}
}
async fn has_batch_column(db: &Database) -> Result<bool, Error> {
let sql = pick(
db.driver(),
"SELECT COUNT(*) FROM information_schema.columns \
WHERE table_schema = DATABASE() AND table_name = '__diesel_schema_migrations' \
AND column_name = 'batch'",
"SELECT COUNT(*) FROM information_schema.columns \
WHERE table_schema = current_schema() AND table_name = '__diesel_schema_migrations' \
AND column_name = 'batch'",
"SELECT COUNT(*) FROM pragma_table_info('__diesel_schema_migrations') \
WHERE name = 'batch'",
);
let found: i64 = backend::scalar(Target::Pool(&db.pool), sql, &[]).await?;
Ok(found > 0)
}
fn is_duplicate_column(err: &sqlx::Error) -> bool {
err.as_database_error().is_some_and(|db_err| {
matches!(db_err.code().as_deref(), Some("1060" | "42701"))
|| db_err.message().contains("duplicate column name")
})
}
async fn applied(db: &Database) -> Result<HashMap<String, i64>, Error> {
let rows: Vec<(String, i32)> = backend::fetch_all(
Target::Pool(&db.pool),
&format!("SELECT version, batch FROM {TABLE}"),
&[],
)
.await?;
Ok(rows
.into_iter()
.map(|(version, batch)| (version, i64::from(batch)))
.collect())
}
async fn recorded_versions(db: &Database) -> Result<Option<HashSet<String>>, Error> {
let sql = pick(
db.driver(),
"SELECT COUNT(*) FROM information_schema.tables \
WHERE table_schema = DATABASE() AND table_name = '__diesel_schema_migrations'",
"SELECT COUNT(*) FROM information_schema.tables \
WHERE table_schema = current_schema() AND table_name = '__diesel_schema_migrations'",
"SELECT COUNT(*) FROM sqlite_master \
WHERE type = 'table' AND name = '__diesel_schema_migrations'",
);
let found: i64 = backend::scalar(Target::Pool(&db.pool), sql, &[]).await?;
if found == 0 {
return Ok(None);
}
let versions: Vec<(String,)> = backend::fetch_all(
Target::Pool(&db.pool),
&format!("SELECT version FROM {TABLE}"),
&[],
)
.await?;
Ok(Some(versions.into_iter().map(|(v,)| v).collect()))
}
pub(crate) fn read_file(path: &Path) -> Result<String, Error> {
std::fs::read_to_string(path).map_err(|source| Error::Read {
path: path.to_path_buf(),
source,
})
}