slave-framework 0.6.0

//! The migration engine: decides what to run and records what ran.

use std::collections::{BTreeMap, HashSet};
use std::time::{Duration, Instant};

use libsql::{Connection, TransactionBehavior, Value};

use super::{Registered, Schema};
use crate::{Error, Result};

/// Name of the table that records applied migrations.
pub const TABLE: &str = "slave_migrations";

/// A migration that completed successfully.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Outcome {
    pub name: String,
    pub duration: Duration,
}

/// A migration that failed. Its tracking record was not written (or removed).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Failure {
    pub name: String,
    pub error: String,
    /// Whether the migration's own changes were rolled back by its transaction.
    pub reverted: bool,
}

/// Result of [`Migrator::run`].
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RunReport {
    /// Batch number assigned to the migrations applied by this run.
    pub batch: i64,
    pub applied: Vec<Outcome>,
    /// Migrations that were already applied and therefore not run again.
    pub skipped: Vec<String>,
    /// The migration that stopped the run, if any. Later migrations were not attempted.
    pub failed: Option<Failure>,
}

/// Result of [`Migrator::rollback`].
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RollbackReport {
    /// The batch that was rolled back, or `None` if nothing had been applied.
    pub batch: Option<i64>,
    pub rolled_back: Vec<Outcome>,
    pub failed: Option<Failure>,
}

/// Where a migration stands relative to the database.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum State {
    Applied {
        batch: i64,
        executed_at: String,
    },
    Pending,
    /// Recorded as applied, but no longer registered in the application.
    Missing {
        batch: i64,
        executed_at: String,
    },
}

/// One row of [`Migrator::status`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Status {
    pub name: String,
    pub state: State,
}

struct Record {
    name: String,
    batch: i64,
    executed_at: String,
}

/// Runs, inspects and rolls back migrations on a single connection.
pub struct Migrator {
    conn: Connection,
    migrations: Vec<Registered>,
}

impl Migrator {
    /// Creates a migrator for `migrations`, which are run in name order.
    pub fn new(conn: Connection, mut migrations: Vec<Registered>) -> Result<Self> {
        migrations.sort_by_key(|m| m.name());
        if let Some(pair) = migrations.windows(2).find(|w| w[0].name() == w[1].name()) {
            return Err(Error::DuplicateMigration(pair[0].name().to_owned()));
        }
        Ok(Self { conn, migrations })
    }

    pub fn connection(&self) -> &Connection {
        &self.conn
    }

    /// Applies every pending migration, in order, as a new batch.
    ///
    /// Stops at the first failure; migrations applied before it stay applied.
    pub async fn run(&mut self) -> Result<RunReport> {
        let records = self.records().await?;
        let applied: HashSet<&str> = records.iter().map(|r| r.name.as_str()).collect();
        let batch = records.iter().map(|r| r.batch).max().unwrap_or(0) + 1;

        let mut report = RunReport {
            batch,
            ..RunReport::default()
        };
        let pending: Vec<Registered> = self
            .migrations
            .iter()
            .filter(|m| {
                let done = applied.contains(m.name());
                if done {
                    report.skipped.push(m.name().to_owned());
                }
                !done
            })
            .copied()
            .collect();

        for registered in pending {
            let migration = registered.migration();
            let started = Instant::now();
            let track = (
                format!("INSERT INTO {TABLE} (name, batch) VALUES (?1, ?2)"),
                vec![
                    Value::Text(registered.name().to_owned()),
                    Value::Integer(batch),
                ],
            );
            match self
                .apply(migration.transactional(), &migration.up(), track)
                .await
            {
                Ok(()) => report.applied.push(Outcome {
                    name: registered.name().to_owned(),
                    duration: started.elapsed(),
                }),
                Err(err) => {
                    report.failed = Some(Failure {
                        name: registered.name().to_owned(),
                        error: err.to_string(),
                        reverted: migration.transactional(),
                    });
                    break;
                }
            }
        }
        Ok(report)
    }

    /// Reverts the most recent batch, newest migration first.
    ///
    /// Refuses (changing nothing) if any migration in the batch is irreversible
    /// or no longer registered.
    pub async fn rollback(&mut self) -> Result<RollbackReport> {
        let records = self.records().await?;
        let Some(batch) = records.iter().map(|r| r.batch).max() else {
            return Ok(RollbackReport::default());
        };

        let mut targets: Vec<&Record> = records.iter().filter(|r| r.batch == batch).collect();
        targets.sort_by(|a, b| b.name.cmp(&a.name));

        let mut plan = Vec::with_capacity(targets.len());
        let (mut missing, mut irreversible) = (Vec::new(), Vec::new());
        for record in targets {
            match self.find(&record.name) {
                None => missing.push(record.name.clone()),
                Some(registered) => match registered.migration().down() {
                    None => irreversible.push(record.name.clone()),
                    Some(down) => plan.push((registered, down)),
                },
            }
        }
        if !missing.is_empty() {
            return Err(Error::RollbackRefused {
                batch,
                reason: format!("not registered in this application: {}", missing.join(", ")),
            });
        }
        if !irreversible.is_empty() {
            return Err(Error::RollbackRefused {
                batch,
                reason: format!(
                    "irreversible (down() returns None): {}",
                    irreversible.join(", ")
                ),
            });
        }

        let mut report = RollbackReport {
            batch: Some(batch),
            ..RollbackReport::default()
        };
        for (registered, down) in plan {
            let transactional = registered.migration().transactional();
            let started = Instant::now();
            let track = (
                format!("DELETE FROM {TABLE} WHERE name = ?1"),
                vec![Value::Text(registered.name().to_owned())],
            );
            match self.apply(transactional, &down, track).await {
                Ok(()) => report.rolled_back.push(Outcome {
                    name: registered.name().to_owned(),
                    duration: started.elapsed(),
                }),
                Err(err) => {
                    report.failed = Some(Failure {
                        name: registered.name().to_owned(),
                        error: err.to_string(),
                        reverted: transactional,
                    });
                    break;
                }
            }
        }
        Ok(report)
    }

    /// Lists every registered or recorded migration in name order.
    pub async fn status(&mut self) -> Result<Vec<Status>> {
        let mut states: BTreeMap<String, State> = self
            .migrations
            .iter()
            .map(|m| (m.name().to_owned(), State::Pending))
            .collect();
        for record in self.records().await? {
            let registered = states.contains_key(&record.name);
            let (batch, executed_at) = (record.batch, record.executed_at);
            let state = if registered {
                State::Applied { batch, executed_at }
            } else {
                State::Missing { batch, executed_at }
            };
            states.insert(record.name, state);
        }
        Ok(states
            .into_iter()
            .map(|(name, state)| Status { name, state })
            .collect())
    }

    fn find(&self, name: &str) -> Option<Registered> {
        self.migrations.iter().find(|m| m.name() == name).copied()
    }

    /// Creates the tracking table if needed and returns its rows.
    async fn records(&self) -> Result<Vec<Record>> {
        self.conn
            .execute(
                &format!(
                    "CREATE TABLE IF NOT EXISTS {TABLE} (
                        id INTEGER PRIMARY KEY,
                        name TEXT NOT NULL UNIQUE,
                        batch INTEGER NOT NULL,
                        executed_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
                    )"
                ),
                (),
            )
            .await?;

        let mut rows = self
            .conn
            .query(
                &format!("SELECT name, batch, executed_at FROM {TABLE} ORDER BY name"),
                (),
            )
            .await?;
        let mut records = Vec::new();
        while let Some(row) = rows.next().await? {
            records.push(Record {
                name: row.get(0)?,
                batch: row.get(1)?,
                executed_at: row.get(2)?,
            });
        }
        Ok(records)
    }

    /// Executes `schema` followed by the tracking statement, atomically when `transactional`.
    async fn apply(
        &mut self,
        transactional: bool,
        schema: &Schema,
        track: (String, Vec<Value>),
    ) -> Result<()> {
        if !transactional {
            return execute(&self.conn, schema, track).await;
        }

        let tx = self
            .conn
            .transaction_with_behavior(TransactionBehavior::Immediate)
            .await?;
        match execute(&tx, schema, track).await {
            Ok(()) => Ok(tx.commit().await?),
            Err(err) => {
                // Report the original failure; a rollback error would only obscure it.
                let _ = tx.rollback().await;
                Err(err)
            }
        }
    }
}

async fn execute(
    conn: &Connection,
    schema: &Schema,
    (sql, params): (String, Vec<Value>),
) -> Result<()> {
    for (index, statement) in schema.statements().iter().enumerate() {
        conn.execute_batch(statement)
            .await
            .map_err(|source| Error::Statement {
                index: index + 1,
                source,
            })?;
    }
    conn.execute(&sql, params).await?;
    Ok(())
}