From 9d1bc64cedd22115e60479b85c8db9b80086aba9 Mon Sep 17 00:00:00 2001
From: Austin Bonander <austin.bonander@gmail.com>
Date: Mon, 23 Sep 2024 02:06:46 -0700
Subject: [PATCH] feat: teach macros about `migrate.table-name`,
 `migrations-dir`

---
 sqlx-macros-core/src/migrate.rs   | 39 ++++++++++++++++++++++---------
 sqlx-macros-core/src/test_attr.rs | 14 +++++++----
 sqlx-macros/src/lib.rs            |  2 +-
 src/macros/mod.rs                 |  2 +-
 4 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/sqlx-macros-core/src/migrate.rs b/sqlx-macros-core/src/migrate.rs
index c9cf5b8e..56ac6140 100644
--- a/sqlx-macros-core/src/migrate.rs
+++ b/sqlx-macros-core/src/migrate.rs
@@ -3,12 +3,15 @@ extern crate proc_macro;
 
 use std::path::{Path, PathBuf};
 
-use proc_macro2::TokenStream;
+use proc_macro2::{Span, TokenStream};
 use quote::{quote, ToTokens, TokenStreamExt};
 use syn::LitStr;
-
+use syn::spanned::Spanned;
+use sqlx_core::config::Config;
 use sqlx_core::migrate::{Migration, MigrationType};
 
+pub const DEFAULT_PATH: &str = "./migrations";
+
 pub struct QuoteMigrationType(MigrationType);
 
 impl ToTokens for QuoteMigrationType {
@@ -81,20 +84,26 @@ impl ToTokens for QuoteMigration {
     }
 }
 
-pub fn expand_migrator_from_lit_dir(dir: LitStr) -> crate::Result<TokenStream> {
-    expand_migrator_from_dir(&dir.value(), dir.span())
+pub fn default_path(config: &Config) -> &str {
+    config.migrate.migrations_dir
+        .as_deref()
+        .unwrap_or(DEFAULT_PATH)
 }
 
-pub(crate) fn expand_migrator_from_dir(
-    dir: &str,
-    err_span: proc_macro2::Span,
-) -> crate::Result<TokenStream> {
-    let path = crate::common::resolve_path(dir, err_span)?;
+pub fn expand(path_arg: Option<LitStr>) -> crate::Result<TokenStream> {
+    let config = Config::from_crate();
 
-    expand_migrator(&path)
+   let path = match path_arg {
+       Some(path_arg) => crate::common::resolve_path(path_arg.value(), path_arg.span())?,
+       None => {
+           crate::common::resolve_path(default_path(config), Span::call_site())
+       }?
+   };
+
+    expand_with_path(config, &path)
 }
 
-pub(crate) fn expand_migrator(path: &Path) -> crate::Result<TokenStream> {
+pub fn expand_with_path(config: &Config, path: &Path) -> crate::Result<TokenStream> {
     let path = path.canonicalize().map_err(|e| {
         format!(
             "error canonicalizing migration directory {}: {e}",
@@ -119,11 +128,19 @@ pub(crate) fn expand_migrator(path: &Path) -> crate::Result<TokenStream> {
         proc_macro::tracked_path::path(path);
     }
 
+    let table_name = config.migrate.table_name
+        .as_deref()
+        .map_or_else(
+            || quote! {},
+            |name| quote! { table_name: Some(::std::borrow::Cow::Borrowed(#name)), }
+        );
+
     Ok(quote! {
         ::sqlx::migrate::Migrator {
             migrations: ::std::borrow::Cow::Borrowed(&[
                     #(#migrations),*
             ]),
+            #table_name
             ..::sqlx::migrate::Migrator::DEFAULT
         }
     })
diff --git a/sqlx-macros-core/src/test_attr.rs b/sqlx-macros-core/src/test_attr.rs
index d7c6eb04..403b6e7d 100644
--- a/sqlx-macros-core/src/test_attr.rs
+++ b/sqlx-macros-core/src/test_attr.rs
@@ -77,6 +77,8 @@ fn expand_simple(input: syn::ItemFn) -> TokenStream {
 
 #[cfg(feature = "migrate")]
 fn expand_advanced(args: AttributeArgs, input: syn::ItemFn) -> crate::Result<TokenStream> {
+    let config = sqlx_core::config::Config::from_crate();
+
     let ret = &input.sig.output;
     let name = &input.sig.ident;
     let inputs = &input.sig.inputs;
@@ -143,15 +145,17 @@ fn expand_advanced(args: AttributeArgs, input: syn::ItemFn) -> crate::Result<Tok
 
     let migrations = match args.migrations {
         MigrationsOpt::ExplicitPath(path) => {
-            let migrator = crate::migrate::expand_migrator_from_lit_dir(path)?;
+            let migrator = crate::migrate::expand(Some(path))?;
             quote! { args.migrator(&#migrator); }
         }
         MigrationsOpt::InferredPath if !inputs.is_empty() => {
-            let migrations_path =
-                crate::common::resolve_path("./migrations", proc_macro2::Span::call_site())?;
+            let path = crate::migrate::default_path(config);
 
-            if migrations_path.is_dir() {
-                let migrator = crate::migrate::expand_migrator(&migrations_path)?;
+            let resolved_path =
+                crate::common::resolve_path(path, proc_macro2::Span::call_site())?;
+
+            if resolved_path.is_dir() {
+                let migrator = crate::migrate::expand_with_path(config, &resolved_path)?;
                 quote! { args.migrator(&#migrator); }
             } else {
                 quote! {}
diff --git a/sqlx-macros/src/lib.rs b/sqlx-macros/src/lib.rs
index 987794ac..f527f5d2 100644
--- a/sqlx-macros/src/lib.rs
+++ b/sqlx-macros/src/lib.rs
@@ -69,7 +69,7 @@ pub fn migrate(input: TokenStream) -> TokenStream {
     use syn::LitStr;
 
     let input = syn::parse_macro_input!(input as LitStr);
-    match migrate::expand_migrator_from_lit_dir(input) {
+    match migrate::expand(input) {
         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
index 7f8ff747..c9602b55 100644
--- a/src/macros/mod.rs
+++ b/src/macros/mod.rs
@@ -814,6 +814,6 @@ macro_rules! migrate {
     }};
 
     () => {{
-        $crate::sqlx_macros::migrate!("./migrations")
+        $crate::sqlx_macros::migrate!()
     }};
 }