sir_eel/

migrator.rs

use core::result;
use std::{
    fs::{self, DirEntry, File},
    path::PathBuf,
};

use serde::{Deserialize, Serialize};
use surrealdb::{
    Surreal,
    engine::any::{self, Any},
};

use crate::{Config, Result};

pub struct Migrator {
    conn: Surreal<Any>,
    config: Config,
}

// constructors
impl Migrator {
    pub async fn new(config: Config) -> Result<Self> {
        fs::create_dir_all(&config.migrations_dir)?;
        let db = any::connect(&config.db_url).await?;
        db.use_ns(&config.namespace)
            .use_db(&config.database)
            .await?;

        let query = r#"
            DEFINE TABLE IF NOT EXISTS schema_migrations SCHEMAFULL;
            DEFINE FIELD IF NOT EXISTS version ON schema_migrations TYPE string;
        "#;
        db.query(query).await?;

        Ok(Self { conn: db, config })
    }
}

// methods
impl Migrator {
    /// Generate a new migration file in the configured `migrations_dir` named
    /// by normalizing the given `input_name`.
    pub fn generate(&self, input_name: impl Into<String>) -> Result<PathBuf> {
        let now = chrono::Utc::now();
        let unixtime = format!("{}", now.timestamp_millis());
        let filename: String = input_name.into().replace(" ", "_").replace("-", "_");
        let pathbuf = self
            .config
            .migrations_dir
            .join(format!("{unixtime}_{filename}.surql"));

        File::create(&pathbuf)?;

        Ok(pathbuf)
    }

    pub async fn run(&self) -> Result<()> {
        for migration in self.pending_migrations().await? {
            let surql = fs::read_to_string(&migration.path)?;
            self.conn.query(surql).await?;

            let _: Option<Version> = self
                .conn
                .create("schema_migrations")
                .content(migration.version)
                .await?;
        }

        Ok(())
    }

    pub async fn pending_migrations(&self) -> Result<Vec<Migration>> {
        let versions = self.applied_versions().await?;
        let all_migration_files = self.migration_file_paths()?;
        let pending = all_migration_files
            .into_iter()
            .filter_map(|f| {
                f.file_name()
                    .to_str()
                    .and_then(|fname| match fname.split_once("_") {
                        Some((ts, rest)) => {
                            let ts = ts.to_string();

                            if versions.contains(&ts) {
                                return None;
                            }

                            if !rest.ends_with(".surql") {
                                return None;
                            }

                            ts.parse::<i64>()
                                .ok()
                                .and_then(chrono::DateTime::from_timestamp_millis)?;

                            Some(Migration {
                                version: Version { version: ts },
                                path: f.path(),
                            })
                        }
                        None => None, // ignore entries that don't have an "_"
                    })
            })
            .collect();

        Ok(pending)
    }

    fn migration_file_paths(&self) -> Result<Vec<DirEntry>> {
        let mut paths: Vec<_> = fs::read_dir(&self.config.migrations_dir)?
            .filter_map(result::Result::ok)
            .collect();
        paths.sort_by_key(|path| path.path());

        Ok(paths)
    }

    async fn applied_versions(&self) -> Result<Vec<String>> {
        let versions: Vec<Version> = self
            .conn
            .query("SELECT version FROM schema_migrations ORDER BY version ASC")
            .await?
            .take(0)?;
        Ok(versions.into_iter().map(|v| v.version).collect())
    }
}

#[derive(Serialize, Deserialize)]
pub struct Version {
    pub version: String,
}

pub struct Migration {
    pub version: Version,
    pub path: PathBuf,
}

#[cfg(test)]
mod tests {
    use std::{ffi::OsStr, time::Duration};
    type Result<T> = core::result::Result<T, Box<dyn core::error::Error>>;

    use super::*;

    fn build_config() -> Config {
        let path = tempfile::tempdir().unwrap().path().join("migrations");
        Config::builder()
            .migration_dir(path)
            .database("testdb")
            .namespace("testns")
            .db_url("memory")
            .build()
            .unwrap()
    }

    #[tokio::test]
    async fn test_new_creates_schema_migrations_table() -> Result<()> {
        let mig = Migrator::new(build_config()).await?;

        #[derive(Deserialize)]
        struct Fields {
            pub version: String,
        }
        #[derive(Deserialize)]
        struct Info {
            pub fields: Fields,
        }

        let res: Option<Info> = mig
            .conn
            .query("info for table schema_migrations")
            .await?
            .take(0)?;

        assert!(res.is_some());

        let info = res.unwrap();
        assert_eq!(
            info.fields.version,
            "DEFINE FIELD version ON schema_migrations TYPE string PERMISSIONS FULL"
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_generate_ok() -> Result<()> {
        let mig = Migrator::new(build_config()).await?;
        let res = mig.generate("create users")?;

        assert!(res.is_file(), "res is not a file: {res:?}");

        let filename = res.file_name().and_then(OsStr::to_str).unwrap();
        let (ts, name) = filename.split_once("_").unwrap();
        let parsed = ts
            .parse::<i64>()
            .ok()
            .and_then(chrono::DateTime::from_timestamp_millis);

        assert!(
            parsed.is_some(),
            "migration file did not begin with a timestamp"
        );
        assert_eq!(name, "create_users.surql");

        Ok(())
    }

    #[tokio::test]
    async fn test_pending_migrations() -> Result<()> {
        let mig = Migrator::new(build_config()).await?;
        let pending = mig.pending_migrations().await?;

        assert!(pending.is_empty(), "expected no pending migrations");

        let path = mig.generate("create pokemon")?;

        let pending = mig.pending_migrations().await?;
        assert_eq!(pending.len(), 1, "expected one pending migration");

        let expected_filename = path.file_name();
        assert!(
            pending
                .iter()
                .any(|m| m.path.file_name() == expected_filename),
            "expected generated migration to be pending"
        );

        fs::write(path, "DEFINE TABLE pokemon;")?;
        mig.run().await?;

        let pending = mig.pending_migrations().await?;
        assert!(pending.is_empty(), "expected no pending migrations");

        Ok(())
    }

    #[tokio::test]
    async fn test_run() -> Result<()> {
        let mig = Migrator::new(build_config()).await?;
        let path = mig.generate("create users")?;
        let surql = r#"
        DEFINE TABLE users SCHEMAFULL;
        DEFINE FIELD name ON users TYPE string;
        "#;
        fs::write(path, surql)?;

        mig.run().await?;

        #[derive(Deserialize)]
        struct Fields {
            pub name: String,
        }
        #[derive(Deserialize)]
        struct Info {
            pub fields: Fields,
        }

        let res: Option<Info> = mig.conn.query("info for table users").await?.take(0)?;

        let info = res.unwrap();
        assert_eq!(
            info.fields.name,
            "DEFINE FIELD name ON users TYPE string PERMISSIONS FULL"
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_run_ignores_irrelevant_files() -> Result<()> {
        let mig = Migrator::new(build_config()).await?;
        fs::create_dir_all(&mig.config.migrations_dir)?;

        // files must have .surql extension
        let filename = format!("{}_foo.txt", chrono::Utc::now().timestamp_millis());
        let bad_file = mig.config.migrations_dir.join(filename);
        fs::write(bad_file, "Don't run me in the database")?;

        // file names must start with timestamp
        let bad_file = mig.config.migrations_dir.join("no_timestamp.surql");
        fs::write(bad_file, "Don't run me in the database")?;

        assert!(mig.run().await.is_ok());
        Ok(())
    }

    #[tokio::test]
    async fn test_run_runs_files_once() -> Result<()> {
        let mig = Migrator::new(build_config()).await?;
        let path_fruits = mig.generate("create fruits")?;
        fs::write(&path_fruits, "DEFINE TABLE fruits")?;

        tokio::time::sleep(Duration::from_millis(2)).await;

        let path_veg = mig.generate("create vegetables")?;
        fs::write(path_veg, "DEFINE TABLE vegetables")?;

        assert_eq!(
            mig.applied_versions().await?.len(),
            0,
            "schema_migrations should be empty"
        );

        #[derive(Serialize, Deserialize)]
        struct Mig {
            pub version: String,
        }

        let (ts, _) = path_fruits
            .file_name()
            .and_then(OsStr::to_str)
            .unwrap()
            .split_once("_")
            .unwrap();
        let res: Option<Mig> = mig
            .conn
            .create("schema_migrations")
            .content(Mig {
                version: ts.to_string(),
            })
            .await?;

        assert!(res.is_some());
        assert_eq!(
            mig.applied_versions().await?.len(),
            1,
            "schema_migrations should have one version"
        );

        // Run migrations
        mig.run().await?;

        #[derive(Deserialize)]
        struct Tables {
            pub fruits: Option<String>,
            pub vegetables: Option<String>,
        }

        #[derive(Deserialize)]
        struct DbInfo {
            pub tables: Tables,
        }

        let res: Option<DbInfo> = mig.conn.query("info for db").await?.take(0)?;
        let info = res.unwrap();

        assert!(
            info.tables.fruits.is_none(),
            "expected not to create fruits"
        );
        assert!(
            info.tables.vegetables.is_some(),
            "expected to create vegetables"
        );
        assert_eq!(
            mig.applied_versions().await?.len(),
            2,
            "schema_migrations should have two versions"
        );

        Ok(())
    }
}