use crate::{repo, schema, Entity, Executor, Result};
use core::future::Future;
use core::marker::PhantomData;
use core::pin::Pin;
pub type StepFuture<'a> = Pin<Box<dyn Future<Output = Result<()>> + 'a>>;
#[derive(Debug, Clone, Entity)]
#[orm(table = "_orm_migrations")]
struct Applied {
#[orm(pk)]
name: String,
}
pub trait Migration<X: Executor> {
fn name(&self) -> &'static str;
fn up<'a>(&'a self, db: &'a X) -> StepFuture<'a>;
}
pub struct CreateTable<E> {
name: &'static str,
_entity: PhantomData<E>,
}
impl<E> CreateTable<E> {
pub fn new(name: &'static str) -> Self {
Self { name, _entity: PhantomData }
}
}
impl<E: Entity, X: Executor> Migration<X> for CreateTable<E> {
fn name(&self) -> &'static str {
self.name
}
fn up<'a>(&'a self, db: &'a X) -> StepFuture<'a> {
Box::pin(async move {
for ddl in schema::schema_ddl::<E>(db.dialect()) {
db.execute(&ddl, vec![]).await?;
}
Ok(())
})
}
}
pub struct Migrator<X> {
steps: Vec<Box<dyn Migration<X>>>,
}
impl<X: Executor> Default for Migrator<X> {
fn default() -> Self {
Self::new()
}
}
impl<X: Executor> Migrator<X> {
pub fn new() -> Self {
Self { steps: Vec::new() }
}
#[allow(clippy::should_implement_trait)] pub fn add(mut self, m: impl Migration<X> + 'static) -> Self {
self.steps.push(Box::new(m));
self
}
pub fn create_table<E: Entity + 'static>(self, name: &'static str) -> Self {
self.add(CreateTable::<E>::new(name))
}
pub fn names(&self) -> Vec<&'static str> {
self.steps.iter().map(|s| s.name()).collect()
}
pub async fn run(&self, db: &X) -> Result<Vec<&'static str>> {
if db.is_sharded() {
crate::sharding::ensure_control_tables(db).await?;
}
db.execute(&schema::create_table_ddl::<Applied>(db.dialect()), vec![]).await?;
let done: Vec<String> =
repo::all::<Applied, _>(db).await?.into_iter().map(|a| a.name).collect();
let mut applied = Vec::new();
for step in &self.steps {
if done.iter().any(|n| n == step.name()) {
continue;
}
step.up(db).await?;
repo::insert(db, &Applied { name: step.name().to_string() }).await?;
applied.push(step.name());
}
Ok(applied)
}
pub async fn run_on_each(&self, dbs: &[&X]) -> Result<()> {
for db in dbs {
self.run(db).await?;
}
Ok(())
}
}