escarmouche/src/migration.rs

use resource::resource_str;
use rusqlite::{Connection, Transaction};
use sha2::{Digest, Sha256};
use std::time::Instant;

#[derive(Debug)]
pub enum IntegrityError {
    MigrationCheckFailed {
        migration_id: usize,
        migration_definition: String,
        migration_sha256: String,
        existing_migration_sha256: String,
    },
}

#[derive(Debug)]
pub enum Error {
    Sql(rusqlite::Error),
    Integrity(IntegrityError),
}

const INITIAL_MIGRATION: &str = "
CREATE TABLE IF NOT EXISTS shitty_migrations (
    id                 INTEGER PRIMARY KEY,
    creation_timestamp DATETIME DEFAULT current_timestamp NOT NULL,
    duration           REAL DEFAULT 0 NOT NULL,
    sha256             BLOB(2) NOT NULL
);
";

fn get_id_of_new_migration(tx: &Transaction) -> Result<usize, rusqlite::Error> {
    let q = "SELECT COALESCE(MAX(id) + 1, 0) FROM shitty_migrations";
    return tx.query_row(q, (), |row| row.get::<usize, usize>(0));
}

pub fn migrate(conn: &mut Connection, migrations: &Vec<String>) -> Result<(), Error> {
    let tx = conn.transaction().map_err(|e| Error::Sql(e))?;
    tx.execute(INITIAL_MIGRATION, ())
        .map_err(|e| Error::Sql(e))?;
    let id = get_id_of_new_migration(&tx).map_err(|e| Error::Sql(e))?;

    for (i, migration) in migrations.iter().enumerate() {
        if i < id {
            check_migration(&tx, i, migration)?;
        } else {
            apply_migration(&tx, i, migration)?;
        }
    }
    tx.commit().map_err(|e| Error::Sql(e))?;
    return Ok(());
}

fn apply_migration(tx: &Transaction, id: usize, migration: &String) -> Result<(), Error> {
    let hash = compute_migration_sha256(migration);

    let q = "INSERT INTO shitty_migrations(id, sha256) VALUES (?1, ?2)";
    tx.execute(q, (&id, &hash)).map_err(|e| Error::Sql(e))?;

    let now = Instant::now();
    tx.execute_batch(migration).map_err(|e| Error::Sql(e))?;
    let elapsed_time = now.elapsed();

    let q = "UPDATE shitty_migrations SET duration = ?1 WHERE id = ?2";
    tx.execute(q, (elapsed_time.as_secs_f32(), &id))
        .map_err(|e| Error::Sql(e))?;

    return Ok(());
}

fn compute_migration_sha256(migration: &String) -> String {
    let mut hasher = Sha256::new();
    hasher.update(migration);
    let ret = hasher.finalize();
    return format!("{:x}", ret);
}

fn check_migration(tx: &Transaction, id: usize, migration: &String) -> Result<(), Error> {
    let q = "SELECT sha256 FROM shitty_migrations WHERE id = ?1";
    let res = tx
        .query_row(q, (&id,), |row| row.get::<usize, String>(0))
        .map_err(|e| Error::Sql(e))?;
    let expected = compute_migration_sha256(migration);
    if !res.eq(&expected) {
        return Err(Error::Integrity(IntegrityError::MigrationCheckFailed {
            migration_definition: migration.clone(),
            migration_id: id,
            migration_sha256: expected,
            existing_migration_sha256: res,
        }));
    }
    return Ok(());
}

fn get_migrations() -> Vec<String> {
    let mut ret = Vec::new();
    ret.push(
        resource_str!("resources/migrations/sql/create_tables_001.sql")
            .as_ref()
            .to_string(),
    );
    return ret;
}

pub fn apply_migrations(conn: &mut Connection) -> Result<(), Error> {
    let migrations = get_migrations();
    return migrate(conn, &migrations);
}

#[cfg(test)]
mod test_migration {
    use super::*;

    fn create_basic_migrations() -> Vec<String> {
        let migration1 = "CREATE TABLE people(name TEXT, age INTEGER);".to_string();
        let migration2 = "INSERT INTO people(name, age) VALUES (\"didier\", 37);".to_string();
        return vec![migration1, migration2];
    }

    fn create_basic_migrations_2() -> Vec<String> {
        let migration1 = "CREATE TABLE people(name TEXT, age INTEGER);".to_string();
        let migration2 = "INSERT INTO people(name, age) VALUES (\"stephanie\", 38);".to_string();
        return vec![migration1, migration2];
    }

    #[test]
    fn basic() {
        let mut conn = Connection::open_in_memory().unwrap();
        let ret = migrate(&mut conn, &create_basic_migrations());
        assert!(ret.is_ok());
        let tx = conn.transaction().unwrap();
        assert_eq!(get_id_of_new_migration(&tx).unwrap(), 2);

        let q = "SELECT COUNT(*) FROM people;";
        let c = tx.query_row(q, (), |x| x.get::<usize, usize>(0)).unwrap();
        assert_eq!(c, c);
    }

    #[test]
    fn idempotency() -> Result<(), Error> {
        let mut conn = Connection::open_in_memory().unwrap();

        migrate(&mut conn, &create_basic_migrations())?;
        migrate(&mut conn, &create_basic_migrations())?;

        let tx = conn.transaction().unwrap();

        assert_eq!(get_id_of_new_migration(&tx).unwrap(), 2);
        let q = "SELECT COUNT(*) FROM people;";
        let c = tx.query_row(q, (), |x| x.get::<usize, usize>(0)).unwrap();
        assert_eq!(c, c);
        return Ok(());
    }

    #[test]
    fn broken_migration() -> Result<(), Error> {
        let mut conn = Connection::open_in_memory().unwrap();

        migrate(&mut conn, &create_basic_migrations())?;
        let ret = migrate(&mut conn, &create_basic_migrations_2());
        assert!(ret.is_err());
        return Ok(());
    }
}