Migration enums (#312)

* generate EmbeddedMigrations enum * create enum from migration * consolidate regex and remove lazy_static * fmt, cleanup * remove unused trait * use From for migration enum * remove unneeded usings * add feature flag, update example project * fmt * Update refinery_core/src/util.rs Co-authored-by: João Oliveira <[email protected]> * Update refinery_core/src/util.rs Co-authored-by: João Oliveira <[email protected]> * Update refinery_core/src/util.rs Co-authored-by: João Oliveira <[email protected]> * Update refinery_core/src/util.rs Co-authored-by: João Oliveira <[email protected]> --------- Co-authored-by: João Oliveira <[email protected]>
rust-db · Mar 28, 2024 · 135acc8 · 135acc8
1 parent 9672719
commit 135acc8
Show file tree

Hide file tree

Showing 11 changed files with 183 additions and 51 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -61,7 +61,7 @@ jobs:
       - run: rustup self update
       - run: cd refinery_core && cargo test --all-features -- --test-threads 1
       - run: cd refinery && cargo build --all-features
-      - run: cd refinery_macros && cargo test
+      - run: cd refinery_macros && cargo test --features=enums
       - run: cd refinery_cli && cargo test
 
   test-sqlite:

diff --git a/Cargo.toml b/Cargo.toml
@@ -3,5 +3,6 @@ members = [
     "refinery",
     "refinery_cli",
     "refinery_core",
-    "refinery_macros"
+    "refinery_macros",
+    "examples"
 ]
diff --git a/examples/Cargo.toml b/examples/Cargo.toml
@@ -0,0 +1,19 @@
+[package]
+name = "refinery-examples"
+version = "0.8.12"
+authors = ["Katharina Fey <[email protected]>", "João Oliveira <[email protected]>"]
+description = "Minimal Refinery usage example"
+license = "MIT OR Apache-2.0"
+documentation = "https://docs.rs/refinery/"
+repository = "https://github.com/rust-db/refinery"
+edition = "2021"
+
+[features]
+enums = ["refinery/enums"]
+
+[dependencies]
+refinery = { path = "../refinery", features = ["rusqlite"] }
+rusqlite = "0.29"
+barrel = { version = "0.7", features = ["sqlite3"] }
+log = "0.4"
+env_logger = "0.11"
diff --git a/examples/main.rs b/examples/main.rs
diff --git a/examples/src/main.rs b/examples/src/main.rs
@@ -0,0 +1,42 @@
+use barrel::backend::Sqlite as Sql;
+use log::info;
+use refinery::Migration;
+use rusqlite::Connection;
+
+refinery::embed_migrations!("migrations");
+
+fn main() {
+    env_logger::init();
+
+    let mut conn = Connection::open_in_memory().unwrap();
+
+    let use_iteration = std::env::args().any(|a| a.to_lowercase().eq("--iterate"));
+
+    if use_iteration {
+        // create an iterator over migrations as they run
+        for migration in migrations::runner().run_iter(&mut conn) {
+            process_migration(migration.expect("Migration failed!"));
+        }
+    } else {
+        // or run all migrations in one go
+        migrations::runner().run(&mut conn).unwrap();
+    }
+}
+
+fn process_migration(migration: Migration) {
+    #[cfg(not(feature = "enums"))]
+    {
+        // run something after each migration
+        info!("Post-processing a migration: {}", migration)
+    }
+
+    #[cfg(feature = "enums")]
+    {
+        // or with the `enums` feature enabled, match against migrations to run specific post-migration steps
+        use migrations::EmbeddedMigration;
+        match migration.into() {
+            EmbeddedMigration::Initial(m) => info!("V{}: Initialized the database!", m.version()),
+            m => info!("Got a migration: {:?}", m),
+        }
+    }
+}
diff --git a/refinery/Cargo.toml b/refinery/Cargo.toml
@@ -24,6 +24,7 @@ tiberius = ["refinery-core/tiberius"]
 tiberius-config = ["refinery-core/tiberius", "refinery-core/tiberius-config"]
 serde = ["refinery-core/serde"]
 toml = ["refinery-core/toml"]
+enums = ["refinery-macros/enums"]
 
 [dependencies]
 refinery-core = { version = "0.8.12", path = "../refinery_core" }

diff --git a/refinery_core/src/lib.rs b/refinery_core/src/lib.rs
@@ -9,7 +9,9 @@ pub use crate::error::Error;
 pub use crate::runner::{Migration, Report, Runner, Target};
 pub use crate::traits::r#async::AsyncMigrate;
 pub use crate::traits::sync::Migrate;
-pub use crate::util::{find_migration_files, load_sql_migrations, MigrationType};
+pub use crate::util::{
+    find_migration_files, load_sql_migrations, parse_migration_name, MigrationType,
+};
 
 #[cfg(feature = "rusqlite")]
 pub use rusqlite;

diff --git a/refinery_core/src/runner.rs b/refinery_core/src/runner.rs
@@ -1,4 +1,3 @@
-use regex::Regex;
 use siphasher::sip::SipHasher13;
 use time::OffsetDateTime;
 
@@ -7,19 +6,12 @@ use std::cmp::Ordering;
 use std::collections::VecDeque;
 use std::fmt;
 use std::hash::{Hash, Hasher};
-use std::sync::OnceLock;
 
-use crate::error::Kind;
 use crate::traits::{sync::migrate as sync_migrate, DEFAULT_MIGRATION_TABLE_NAME};
+use crate::util::parse_migration_name;
 use crate::{AsyncMigrate, Error, Migrate};
 use std::fmt::Formatter;
 
-// regex used to match file names
-pub fn file_match_re() -> &'static Regex {
-    static RE: OnceLock<regex::Regex> = OnceLock::new();
-    RE.get_or_init(|| Regex::new(r"^([U|V])(\d+(?:\.\d+)?)__(\w+)").unwrap())
-}
-
 /// An enum set that represents the type of the Migration
 #[derive(Clone, PartialEq)]
 pub enum Type {
@@ -84,20 +76,7 @@ impl Migration {
     /// Create an unapplied migration, name and version are parsed from the input_name,
     /// which must be named in the format (U|V){1}__{2}.rs where {1} represents the migration version and {2} the name.
     pub fn unapplied(input_name: &str, sql: &str) -> Result<Migration, Error> {
-        let captures = file_match_re()
-            .captures(input_name)
-            .filter(|caps| caps.len() == 4)
-            .ok_or_else(|| Error::new(Kind::InvalidName, None))?;
-        let version: i32 = captures[2]
-            .parse()
-            .map_err(|_| Error::new(Kind::InvalidVersion, None))?;
-
-        let name: String = (&captures[3]).into();
-        let prefix = match &captures[1] {
-            "V" => Type::Versioned,
-            "U" => Type::Unversioned,
-            _ => unreachable!(),
-        };
+        let (prefix, version, name) = parse_migration_name(input_name)?;
 
         // Previously, `std::collections::hash_map::DefaultHasher` was used
         // to calculate the checksum and the implementation at that time

diff --git a/refinery_core/src/util.rs b/refinery_core/src/util.rs
@@ -1,10 +1,32 @@
 use crate::error::{Error, Kind};
+use crate::runner::Type;
 use crate::Migration;
 use regex::Regex;
 use std::ffi::OsStr;
 use std::path::{Path, PathBuf};
+use std::sync::OnceLock;
 use walkdir::{DirEntry, WalkDir};
 
+const STEM_RE: &'static str = r"^([U|V])(\d+(?:\.\d+)?)__(\w+)";
+
+/// Matches the stem of a migration file.
+fn file_stem_re() -> &'static Regex {
+    static RE: OnceLock<Regex> = OnceLock::new();
+    RE.get_or_init(|| Regex::new(STEM_RE).unwrap())
+}
+
+/// Matches the stem + extension of a SQL migration file.
+fn file_re_sql() -> &'static Regex {
+    static RE: OnceLock<Regex> = OnceLock::new();
+    RE.get_or_init(|| Regex::new([STEM_RE, r"\.sql$"].concat().as_str()).unwrap())
+}
+
+/// Matches the stem + extension of any migration file.
+fn file_re_all() -> &'static Regex {
+    static RE: OnceLock<Regex> = OnceLock::new();
+    RE.get_or_init(|| Regex::new([STEM_RE, r"\.(rs|sql)$"].concat().as_str()).unwrap())
+}
+
 /// enum containing the migration types used to search for migrations
 /// either just .sql files or both .sql and .rs
 pub enum MigrationType {
@@ -13,16 +35,34 @@ pub enum MigrationType {
 }
 
 impl MigrationType {
-    fn file_match_re(&self) -> Regex {
-        let ext = match self {
-            MigrationType::All => "(rs|sql)",
-            MigrationType::Sql => "sql",
-        };
-        let re_str = format!(r"^(U|V)(\d+(?:\.\d+)?)__(\w+)\.{}$", ext);
-        Regex::new(re_str.as_str()).unwrap()
+    fn file_match_re(&self) -> &'static Regex {
+        match self {
+            MigrationType::All => file_re_all(),
+            MigrationType::Sql => file_re_sql(),
+        }
     }
 }
 
+/// Parse a migration filename stem into a prefix, version, and name.
+pub fn parse_migration_name(name: &str) -> Result<(Type, i32, String), Error> {
+    let captures = file_stem_re()
+        .captures(name)
+        .filter(|caps| caps.len() == 4)
+        .ok_or_else(|| Error::new(Kind::InvalidName, None))?;
+    let version: i32 = captures[2]
+        .parse()
+        .map_err(|_| Error::new(Kind::InvalidVersion, None))?;
+
+    let name: String = (&captures[3]).into();
+    let prefix = match &captures[1] {
+        "V" => Type::Versioned,
+        "U" => Type::Unversioned,
+        _ => unreachable!(),
+    };
+
+    Ok((prefix, version, name))
+}
+
 /// find migrations on file system recursively across directories given a location and [MigrationType]
 pub fn find_migration_files(
     location: impl AsRef<Path>,

diff --git a/refinery_macros/Cargo.toml b/refinery_macros/Cargo.toml
@@ -8,6 +8,9 @@ documentation = "https://docs.rs/refinery/"
 repository = "https://github.com/rust-db/refinery"
 edition = "2018"
 
+[features]
+enums = []
+
 [lib]
 proc-macro = true
 
@@ -17,6 +20,7 @@ quote = "1"
 syn = "2"
 proc-macro2 = "1"
 regex = "1"
+heck = "0.4"
 
 [dev-dependencies]
 tempfile = "3"
diff --git a/refinery_macros/src/lib.rs b/refinery_macros/src/lib.rs
@@ -1,6 +1,7 @@
 //! Contains Refinery macros that are used to import and embed migration files.
 #![recursion_limit = "128"]
 
+use heck::ToUpperCamelCase;
 use proc_macro::TokenStream;
 use proc_macro2::{Span as Span2, TokenStream as TokenStream2};
 use quote::quote;
@@ -31,6 +32,42 @@ fn migration_fn_quoted<T: ToTokens>(_migrations: Vec<T>) -> TokenStream2 {
     result
 }
 
+fn migration_enum_quoted(migration_names: &[impl AsRef<str>]) -> TokenStream2 {
+    if cfg!(feature = "enums") {
+        let mut variants = Vec::new();
+        let mut discriminants = Vec::new();
+
+        for m in migration_names {
+            let m = m.as_ref();
+            let (_, version, name) = refinery_core::parse_migration_name(m)
+                .unwrap_or_else(|e| panic!("Couldn't parse migration filename '{}': {:?}", m, e));
+            let variant = Ident::new(name.to_upper_camel_case().as_str(), Span2::call_site());
+            variants.push(quote! { #variant(Migration) = #version });
+            discriminants.push(quote! { #version => Self::#variant(migration) });
+        }
+        discriminants.push(quote! { v => panic!("Invalid migration version '{}'", v) });
+
+        let result = quote! {
+            #[repr(i32)]
+            #[derive(Debug)]
+            pub enum EmbeddedMigration {
+                #(#variants),*
+            }
+
+            impl From<Migration> for EmbeddedMigration {
+                fn from(migration: Migration) -> Self {
+                    match migration.version() as i32 {
+                        #(#discriminants),*
+                    }
+                }
+            }
+        };
+        result
+    } else {
+        quote!()
+    }
+}
+
 /// Interpret Rust or SQL migrations and inserts a function called runner that when called returns a [`Runner`] instance with the collected migration modules.
 ///
 /// When called without arguments `embed_migrations` searches for migration files on a directory called `migrations` at the root level of your crate.
@@ -56,6 +93,7 @@ pub fn embed_migrations(input: TokenStream) -> TokenStream {
 
     let mut migrations_mods = Vec::new();
     let mut _migrations = Vec::new();
+    let mut migration_filenames = Vec::new();
 
     for migration in migration_files {
         // safe to call unwrap as find_migration_filenames returns canonical paths
@@ -65,6 +103,7 @@ pub fn embed_migrations(input: TokenStream) -> TokenStream {
             .unwrap();
         let path = migration.display().to_string();
         let extension = migration.extension().unwrap();
+        migration_filenames.push(filename.clone());
 
         if extension == "sql" {
             _migrations.push(quote! {(#filename, include_str!(#path).to_string())});
@@ -85,10 +124,12 @@ pub fn embed_migrations(input: TokenStream) -> TokenStream {
     }
 
     let fnq = migration_fn_quoted(_migrations);
+    let enums = migration_enum_quoted(migration_filenames.as_slice());
     (quote! {
         pub mod migrations {
             #(#migrations_mods)*
             #fnq
+            #enums
         }
     })
     .into()
@@ -98,6 +139,27 @@ pub fn embed_migrations(input: TokenStream) -> TokenStream {
 mod tests {
     use super::{migration_fn_quoted, quote};
 
+    #[test]
+    #[cfg(feature = "enums")]
+    fn test_enum_fn() {
+        let expected = concat! {
+            "# [repr (i32)] # [derive (Debug)] ",
+            "pub enum EmbeddedMigration { ",
+            "Foo (Migration) = 1i32 , ",
+            "BarBaz (Migration) = 3i32 ",
+            "} ",
+            "impl From < Migration > for EmbeddedMigration { ",
+            "fn from (migration : Migration) -> Self { ",
+            "match migration . version () as i32 { ",
+            "1i32 => Self :: Foo (migration) , ",
+            "3i32 => Self :: BarBaz (migration) , ",
+            "v => panic ! (\"Invalid migration version '{}'\" , v) ",
+            "} } }"
+        };
+        let enums = super::migration_enum_quoted(&["V1__foo", "U3__barBAZ"]).to_string();
+        assert_eq!(expected, enums);
+    }
+
     #[test]
     fn test_quote_fn() {
         let migs = vec![quote!("V1__first", "valid_sql_file")];