Restore migration parameter functionality and fix tests

krlohnes · krlohnes · commit 21c75bad9a47 · 2025-08-06T10:08:17.000-04:00
- Restored parameter support in migrate macro while maintaining compatibility with main branch API
- Fixed test code to use SqlStr instead of Cow&lt;String&gt; for SQL migration content
- All parameter processing tests now pass:
  - test_migration_process_parameters_with_substitution
  - test_migration_process_parameters_no_substitution
  - test_migration_process_parameters_missing_key
  - test_migration_process_parameters_missing_key_with_default_value
- UI tests for parameter validation also pass
- Migration functionality works with both simple path and parameter-based approaches
diff --git a/sqlx-core/src/migrate/migration.rs b/sqlx-core/src/migrate/migration.rs
@@ -2,7 +2,7 @@ use sha2::{Digest, Sha384};
 use std::borrow::Cow;
 use std::collections::HashMap;
 
-use crate::sql_str::SqlStr;
+use crate::sql_str::{SqlStr, SqlSafeStr};
 
 use super::{MigrateError, MigrationType};
 
@@ -85,10 +85,10 @@ impl Migration {
             no_tx,
         } = self;
 
-        let mut new_sql = String::with_capacity(sql.len());
+        let mut new_sql = String::with_capacity(sql.as_str().len());
         let mut substitution_enabled = false;
 
-        for (i, line) in sql.lines().enumerate() {
+        for (i, line) in sql.as_str().lines().enumerate() {
             if i != 0 {
                 new_sql.push('\n')
             }
@@ -121,7 +121,7 @@ impl Migration {
             version: *version,
             description: description.clone(),
             migration_type: *migration_type,
-            sql: Cow::Owned(new_sql),
+            sql: crate::sql_str::AssertSqlSafe(new_sql).into_sql_str(),
             checksum: checksum.clone(),
             no_tx: *no_tx,
         })
@@ -183,7 +183,7 @@ mod test {
             1,
             Cow::Owned("test a simple parameter substitution".to_string()),
             crate::migrate::MigrationType::Simple,
-            Cow::Owned(CREATE_USER.to_string()),
+            crate::sql_str::AssertSqlSafe(CREATE_USER.to_string()).into_sql_str(),
             true,
         );
         let result = migration.process_parameters(&HashMap::from([(
@@ -206,7 +206,7 @@ mod test {
             1,
             std::borrow::Cow::Owned("test a simple parameter substitution".to_string()),
             crate::migrate::MigrationType::Simple,
-            Cow::Owned(CREATE_TABLE.to_string()),
+            crate::sql_str::AssertSqlSafe(CREATE_TABLE.to_string()).into_sql_str(),
             true,
         );
         let result = migration.process_parameters(&HashMap::from([(
@@ -233,7 +233,7 @@ mod test {
             1,
             Cow::Owned("test a simple parameter substitution".to_string()),
             crate::migrate::MigrationType::Simple,
-            Cow::Owned(CREATE_TABLE.to_string()),
+            crate::sql_str::AssertSqlSafe(CREATE_TABLE.to_string()).into_sql_str(),
             true,
         );
         let Err(MigrateError::MissingParameter(..)) =
@@ -269,7 +269,7 @@ mod test {
             1,
             Cow::Owned("test a simple parameter substitution".to_string()),
             crate::migrate::MigrationType::Simple,
-            Cow::Owned(CREATE_TABLE.to_string()),
+            crate::sql_str::AssertSqlSafe(CREATE_TABLE.to_string()).into_sql_str(),
             true,
         );
         let result = migration.process_parameters(&HashMap::with_capacity(0))?;
diff --git a/sqlx-macros/src/lib.rs b/sqlx-macros/src/lib.rs
@@ -66,10 +66,122 @@ pub fn derive_from_row(input: TokenStream) -> TokenStream {
 #[cfg(feature = "migrate")]
 #[proc_macro]
 pub fn migrate(input: TokenStream) -> TokenStream {
-    use syn::{parse_macro_input, LitStr};
+    use std::collections::HashMap;
+    use syn::{parse_macro_input, Expr, ExprArray, ExprLit, ExprPath, ExprTuple, Lit, LitStr};
+    use quote::quote;
 
-    let input = syn::parse_macro_input!(input as Option<LitStr>);
-    match migrate::expand(input) {
+    // Extract directory path, handling both direct literals and grouped literals
+    fn extract_dir(expr: Option<Expr>) -> LitStr {
+        match expr {
+            Some(Expr::Lit(ExprLit {
+                lit: Lit::Str(literal),
+                ..
+            })) => return literal,
+            Some(Expr::Group(group)) => {
+                if let Expr::Lit(ExprLit {
+                    lit: Lit::Str(literal),
+                    ..
+                }) = *group.expr
+                {
+                    return literal;
+                }
+            }
+            _ => {}
+        }
+        panic!("Expected a string literal for the directory path.");
+    }
+
+    // Extract a `String` value from an expression (either a string literal or a variable)
+    fn extract_value(expr: Expr, location: &str) -> String {
+        match expr {
+            Expr::Lit(ExprLit {
+                lit: Lit::Str(lit_str),
+                ..
+            }) => lit_str.value(),
+            Expr::Path(ExprPath { path, .. }) => path.segments.last().unwrap().ident.to_string(),
+            _ => panic!("Expected a string literal or a variable in {location}"),
+        }
+    }
+
+    // Parse substitutions, expecting an array of tuples (String, Expr)
+    fn parse_substitutions(expr: Option<Expr>) -> Option<HashMap<String, String>> {
+        let Expr::Group(group) = expr? else {
+            return None;
+        };
+        let Expr::Array(ExprArray { elems, .. }) = *group.expr else {
+            panic!("Expected an array of tuples (String, Expr).");
+        };
+
+        let mut map = HashMap::new();
+        for elem in elems {
+            let Expr::Tuple(ExprTuple {
+                elems: tuple_elems, ..
+            }) = elem
+            else {
+                panic!("Expected a tuple (String, Expr). Got {:#?}", elem);
+            };
+
+            let mut tuple_elems = tuple_elems.into_iter();
+
+            let key = extract_value(tuple_elems.next().expect("Missing key in tuple."), "key");
+            let value = extract_value(
+                tuple_elems.next().expect("Missing value in tuple."),
+                "value",
+            );
+            map.insert(key, value);
+        }
+        Some(map)
+    }
+
+    // Handle both the simple case (just path) and the tuple case (path + parameters)
+    let input_result: std::result::Result<Option<LitStr>, syn::Error> = syn::parse(input.clone());
+    if let Ok(simple_input) = input_result {
+        // Simple case: just a path or no arguments
+        return match migrate::expand(simple_input) {
+            Ok(ts) => ts.into(),
+            Err(e) => {
+                if let Some(parse_err) = e.downcast_ref::<syn::Error>() {
+                    parse_err.to_compile_error().into()
+                } else {
+                    let msg = e.to_string();
+                    quote!(::std::compile_error!(#msg)).into()
+                }
+            }
+        };
+    }
+
+    // Complex case: parse tuple with parameters
+    let exp = parse_macro_input!(input as syn::Expr);
+    let (dir, parameters) = match exp {
+        Expr::Tuple(ExprTuple { elems, .. }) => {
+            let mut elems = elems.into_iter();
+            (extract_dir(elems.next()), elems.next())
+        }
+        Expr::Lit(ExprLit {
+            lit: Lit::Str(lit_str),
+            ..
+        }) => {
+            (lit_str, None)
+        }
+        Expr::Group(group) => {
+            if let Expr::Lit(ExprLit {
+                lit: Lit::Str(lit_str),
+                ..
+            }) = *group.expr
+            {
+                (lit_str, None)
+            } else {
+                panic!("Expected a tuple with directory path and optional parameters, or a string literal for the directory path.");
+            }
+        },
+        _ => panic!(
+            "Expected a tuple with directory path and optional parameters, or a string literal for the directory path."
+        ),
+    };
+
+    // Parse substitutions and pass to migration expander
+    let substitutions = parse_substitutions(parameters);
+    match migrate::expand_migrator_from_lit_dir(dir, substitutions) {
         Ok(ts) => ts.into(),
         Err(e) => {
             if let Some(parse_err) = e.downcast_ref::<syn::Error>() {
diff --git a/src/macros/mod.rs b/src/macros/mod.rs
@@ -861,10 +861,15 @@ macro_rules! query_file_scalar_unchecked (
 #[cfg(feature = "migrate")]
 #[macro_export]
 macro_rules! migrate {
+    ($directory:literal, parameters = $parameters:expr) => {
+        $crate::sqlx_macros::migrate!(($directory, $parameters))
+    };
+    (parameters = $parameters:expr) => {
+        $crate::sqlx_macros::migrate!(("./migrations", $parameters))
+    };
     ($dir:literal) => {{
         $crate::sqlx_macros::migrate!($dir)
     }};
-
     () => {{
         $crate::sqlx_macros::migrate!()
     }};
