use async_trait::async_trait;
use sea_query::{
ColumnDef, Expr, Iden, Index, MysqlQueryBuilder, PostgresQueryBuilder, Query,
QueryStatementWriter, SchemaStatementBuilder, SqliteQueryBuilder, Table,
};
use sqlx::{AnyConnection, AnyPool};
use crate::error::{CoreError, CoreResult};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DbBackend {
Postgres,
Mysql,
Sqlite,
}
impl DbBackend {
pub fn from_url(url: &str) -> CoreResult<Self> {
if url.starts_with("postgres") {
Ok(Self::Postgres)
} else if url.starts_with("mysql") {
Ok(Self::Mysql)
} else if url.starts_with("sqlite") {
Ok(Self::Sqlite)
} else {
Err(CoreError::Config(format!(
"unrecognised database URL scheme: {url}"
)))
}
}
}
pub fn bool_col<T: sea_query::IntoIden>(name: T) -> ColumnDef {
ColumnDef::new(name).integer().to_owned()
}
pub const KEY_LEN: u32 = 255;
pub fn key_col<T: sea_query::IntoIden>(name: T) -> ColumnDef {
ColumnDef::new(name).string_len(KEY_LEN).to_owned()
}
fn schema_sql<S: SchemaStatementBuilder>(backend: DbBackend, stmt: &S) -> String {
match backend {
DbBackend::Postgres => stmt.build(PostgresQueryBuilder),
DbBackend::Mysql => stmt.build(MysqlQueryBuilder),
DbBackend::Sqlite => stmt.build(SqliteQueryBuilder),
}
}
fn query_sql<Q: QueryStatementWriter>(backend: DbBackend, stmt: &Q) -> String {
match backend {
DbBackend::Postgres => stmt.to_string(PostgresQueryBuilder),
DbBackend::Mysql => stmt.to_string(MysqlQueryBuilder),
DbBackend::Sqlite => stmt.to_string(SqliteQueryBuilder),
}
}
pub struct Schema<'c> {
conn: &'c mut AnyConnection,
backend: DbBackend,
}
impl Schema<'_> {
pub fn backend(&self) -> DbBackend {
self.backend
}
pub async fn exec<S: SchemaStatementBuilder>(&mut self, stmt: S) -> CoreResult<()> {
let sql = schema_sql(self.backend, &stmt);
drop(stmt);
sqlx::query(&sql).execute(&mut *self.conn).await?;
Ok(())
}
pub async fn raw(&mut self, sql: &str) -> CoreResult<()> {
sqlx::query(sql).execute(&mut *self.conn).await?;
Ok(())
}
}
#[async_trait(?Send)]
pub trait Migration: Send + Sync {
fn name(&self) -> &str;
async fn up(&self, schema: &mut Schema<'_>) -> CoreResult<()>;
async fn down(&self, _schema: &mut Schema<'_>) -> CoreResult<()> {
Err(CoreError::Irreversible {
module: String::new(),
name: self.name().to_string(),
})
}
}
pub struct SqlMigration {
name: String,
up: String,
down: Option<String>,
}
impl SqlMigration {
pub fn new(name: impl Into<String>, up: impl Into<String>) -> Self {
Self {
name: name.into(),
up: up.into(),
down: None,
}
}
pub fn reversible(mut self, down: impl Into<String>) -> Self {
self.down = Some(down.into());
self
}
}
#[async_trait(?Send)]
impl Migration for SqlMigration {
fn name(&self) -> &str {
&self.name
}
async fn up(&self, schema: &mut Schema<'_>) -> CoreResult<()> {
schema.raw(&self.up).await
}
async fn down(&self, schema: &mut Schema<'_>) -> CoreResult<()> {
match &self.down {
Some(sql) => schema.raw(sql).await,
None => Err(CoreError::Irreversible {
module: String::new(),
name: self.name.clone(),
}),
}
}
}
pub struct MigrationSet {
pub module_id: &'static str,
pub migrations: Vec<Box<dyn Migration>>,
}
impl MigrationSet {
pub fn new(module_id: &'static str, migrations: Vec<Box<dyn Migration>>) -> Self {
Self {
module_id,
migrations,
}
}
}
#[derive(Iden)]
enum LateriteMigrations {
Table,
ModuleId,
Name,
}
async fn ensure_tracking_table(pool: &AnyPool, backend: DbBackend) -> CoreResult<()> {
let stmt = Table::create()
.table(LateriteMigrations::Table)
.if_not_exists()
.col(
ColumnDef::new(LateriteMigrations::ModuleId)
.string_len(255)
.not_null(),
)
.col(
ColumnDef::new(LateriteMigrations::Name)
.string_len(255)
.not_null(),
)
.primary_key(
Index::create()
.col(LateriteMigrations::ModuleId)
.col(LateriteMigrations::Name),
)
.to_owned();
sqlx::query(&schema_sql(backend, &stmt))
.execute(pool)
.await?;
Ok(())
}
async fn is_applied(
pool: &AnyPool,
backend: DbBackend,
module_id: &str,
name: &str,
) -> CoreResult<bool> {
let stmt = Query::select()
.column(LateriteMigrations::Name)
.from(LateriteMigrations::Table)
.and_where(Expr::col(LateriteMigrations::ModuleId).eq(module_id))
.and_where(Expr::col(LateriteMigrations::Name).eq(name))
.limit(1)
.to_owned();
let found: Option<String> = sqlx::query_scalar(&query_sql(backend, &stmt))
.fetch_optional(pool)
.await?;
Ok(found.is_some())
}
pub async fn applied(
pool: &AnyPool,
backend: DbBackend,
module_id: &str,
) -> CoreResult<Vec<String>> {
ensure_tracking_table(pool, backend).await?;
let stmt = Query::select()
.column(LateriteMigrations::Name)
.from(LateriteMigrations::Table)
.and_where(Expr::col(LateriteMigrations::ModuleId).eq(module_id))
.order_by(LateriteMigrations::Name, sea_query::Order::Asc)
.to_owned();
let names: Vec<String> = sqlx::query_scalar(&query_sql(backend, &stmt))
.fetch_all(pool)
.await?;
Ok(names)
}
pub async fn run(pool: &AnyPool, backend: DbBackend, sets: &[MigrationSet]) -> CoreResult<()> {
ensure_tracking_table(pool, backend).await?;
for set in sets {
for migration in &set.migrations {
if is_applied(pool, backend, set.module_id, migration.name()).await? {
continue;
}
let mut tx = pool.begin().await?;
{
let mut schema = Schema {
conn: &mut tx,
backend,
};
migration.up(&mut schema).await?;
}
let insert = Query::insert()
.into_table(LateriteMigrations::Table)
.columns([LateriteMigrations::ModuleId, LateriteMigrations::Name])
.values_panic([set.module_id.into(), migration.name().into()])
.to_owned();
sqlx::query(&query_sql(backend, &insert))
.execute(&mut *tx)
.await?;
tx.commit().await?;
}
}
Ok(())
}
pub async fn rollback(
pool: &AnyPool,
backend: DbBackend,
set: &MigrationSet,
steps: usize,
) -> CoreResult<()> {
ensure_tracking_table(pool, backend).await?;
let mut done = 0;
for migration in set.migrations.iter().rev() {
if done >= steps {
break;
}
if !is_applied(pool, backend, set.module_id, migration.name()).await? {
continue;
}
let mut tx = pool.begin().await?;
{
let mut schema = Schema {
conn: &mut tx,
backend,
};
migration.down(&mut schema).await.map_err(|e| match e {
CoreError::Irreversible { name, .. } => CoreError::Irreversible {
module: set.module_id.to_string(),
name,
},
other => other,
})?;
}
let delete = Query::delete()
.from_table(LateriteMigrations::Table)
.and_where(Expr::col(LateriteMigrations::ModuleId).eq(set.module_id))
.and_where(Expr::col(LateriteMigrations::Name).eq(migration.name()))
.to_owned();
sqlx::query(&query_sql(backend, &delete))
.execute(&mut *tx)
.await?;
tx.commit().await?;
done += 1;
}
Ok(())
}
pub async fn reset(pool: &AnyPool, backend: DbBackend, set: &MigrationSet) -> CoreResult<()> {
rollback(pool, backend, set, set.migrations.len()).await
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Iden)]
enum Demo {
Table,
Id,
}
struct CreateDemo;
#[async_trait(?Send)]
impl Migration for CreateDemo {
fn name(&self) -> &str {
"0001_create_demo"
}
async fn up(&self, schema: &mut Schema<'_>) -> CoreResult<()> {
schema
.exec(
Table::create()
.table(Demo::Table)
.if_not_exists()
.col(ColumnDef::new(Demo::Id).integer().not_null())
.to_owned(),
)
.await
}
async fn down(&self, schema: &mut Schema<'_>) -> CoreResult<()> {
schema
.exec(Table::drop().table(Demo::Table).to_owned())
.await
}
}
async fn sqlite_pool() -> AnyPool {
sqlx::any::install_default_drivers();
sqlx::any::AnyPoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap()
}
#[tokio::test]
async fn applies_and_rolls_back_on_sqlite() {
let pool = sqlite_pool().await;
let backend = DbBackend::Sqlite;
let set = MigrationSet::new("test.demo", vec![Box::new(CreateDemo)]);
run(&pool, backend, std::slice::from_ref(&set))
.await
.unwrap();
run(&pool, backend, std::slice::from_ref(&set))
.await
.unwrap();
sqlx::query("insert into demo (id) values (1)")
.execute(&pool)
.await
.unwrap();
assert_eq!(applied(&pool, backend, "test.demo").await.unwrap().len(), 1);
reset(&pool, backend, &set).await.unwrap();
assert!(sqlx::query("select count(*) from demo")
.fetch_one(&pool)
.await
.is_err());
assert!(applied(&pool, backend, "test.demo")
.await
.unwrap()
.is_empty());
}
#[tokio::test]
async fn irreversible_migration_reports_module_and_name() {
let pool = sqlite_pool().await;
let backend = DbBackend::Sqlite;
let set = MigrationSet::new(
"test.oneway",
vec![Box::new(SqlMigration::new(
"0001_make_t",
"create table t (id integer not null)",
))],
);
run(&pool, backend, std::slice::from_ref(&set))
.await
.unwrap();
let err = rollback(&pool, backend, &set, 1).await.unwrap_err();
match err {
CoreError::Irreversible { module, name } => {
assert_eq!(module, "test.oneway");
assert_eq!(name, "0001_make_t");
}
other => panic!("expected Irreversible, got {other:?}"),
}
}
}