use std::collections::HashMap;
use worker::D1Database;
use crate::{error::MigrationError, migration::Migration, tracker};
pub struct MigrationStatus {
pub id: String,
pub name: String,
pub applied: bool,
pub applied_at: Option<String>,
}
pub struct MigrationRunner<'db> {
db: &'db D1Database,
migrations: Vec<Migration>,
}
impl<'db> MigrationRunner<'db> {
pub fn new(db: &'db D1Database) -> Self {
Self { db, migrations: vec![] }
}
pub fn register(mut self, m: Migration) -> Self {
self.migrations.push(m);
self
}
pub async fn run(&self) -> Result<usize, MigrationError> {
tracker::ensure_table(self.db).await?;
let mut applied = 0;
for m in &self.migrations {
if !tracker::is_applied(self.db, m.id).await? {
self.db.prepare(m.up)
.run().await
.map_err(|e| MigrationError::Sql(format!("{}: {}", m.id, e)))?;
tracker::record(self.db, m.id, m.name).await?;
applied += 1;
}
}
Ok(applied)
}
pub async fn rollback(&self, steps: usize) -> Result<usize, MigrationError> {
tracker::ensure_table(self.db).await?;
let entries = tracker::applied_entries(self.db).await?;
let to_roll: Vec<String> = entries.into_iter().rev().take(steps).map(|(id, _)| id).collect();
let mut count = 0;
for id in &to_roll {
let m = self.migrations.iter().find(|m| m.id == id.as_str())
.ok_or_else(|| MigrationError::NotFound(id.clone()))?;
let down = m.down.ok_or_else(|| MigrationError::MissingDown(id.clone()))?;
self.db.prepare(down)
.run().await
.map_err(|e| MigrationError::Sql(format!("{}: {}", id, e)))?;
tracker::remove(self.db, id).await?;
count += 1;
}
Ok(count)
}
pub async fn status(&self) -> Result<Vec<MigrationStatus>, MigrationError> {
tracker::ensure_table(self.db).await?;
let entries = tracker::applied_entries(self.db).await?;
let map: HashMap<String, String> = entries.into_iter().collect();
Ok(self.migrations.iter().map(|m| {
let at = map.get(m.id).cloned();
MigrationStatus {
id: m.id.to_string(),
name: m.name.to_string(),
applied: at.is_some(),
applied_at: at,
}
}).collect())
}
}