From 5db8cb0a3d3e7e7184c00adf8ae0ff4219b9c083 Mon Sep 17 00:00:00 2001
From: Matthias Seiffert <matthias-seiffert@hotmail.de>
Date: Thu, 9 Jan 2025 20:34:32 +0100
Subject: [PATCH] Allow modifying the connection in migrations (#2397)

* Add function run_cli_with_connection

* Delegate run_cli to run_cli_with_connection

* Update sea-orm-migration/src/cli.rs

---------

Co-authored-by: Chris Tsang <chris.2y3@outlook.com>
---
 sea-orm-migration/src/cli.rs | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/sea-orm-migration/src/cli.rs b/sea-orm-migration/src/cli.rs
index aa0d3668..6716fb01 100644
--- a/sea-orm-migration/src/cli.rs
+++ b/sea-orm-migration/src/cli.rs
@@ -1,9 +1,11 @@
+use std::future::Future;
+
 use clap::Parser;
 use dotenvy::dotenv;
 use std::{error::Error, fmt::Display, process::exit};
 use tracing_subscriber::{prelude::*, EnvFilter};
 
-use sea_orm::{ConnectOptions, Database, DbConn};
+use sea_orm::{ConnectOptions, Database, DbConn, DbErr};
 use sea_orm_cli::{run_migrate_generate, run_migrate_init, MigrateSubcommands};
 
 use super::MigratorTrait;
@@ -13,6 +15,20 @@ const MIGRATION_DIR: &str = "./";
 pub async fn run_cli<M>(migrator: M)
 where
     M: MigratorTrait,
+{
+    run_cli_with_connection(migrator, Database::connect).await;
+}
+
+/// Same as [`run_cli`] where you provide the function to create the [`DbConn`].
+///
+/// This allows configuring the database connection as you see fit.
+/// E.g. you can change settings in [`ConnectOptions`] or you can load sqlite
+/// extensions.
+pub async fn run_cli_with_connection<M, F, Fut>(migrator: M, make_connection: F)
+where
+    M: MigratorTrait,
+    F: FnOnce(ConnectOptions) -> Fut,
+    Fut: Future<Output = Result<DbConn, DbErr>>,
 {
     dotenv().ok();
     let cli = Cli::parse();
@@ -25,11 +41,12 @@ where
     let connect_options = ConnectOptions::new(url)
         .set_schema_search_path(schema)
         .to_owned();
-    let db = &Database::connect(connect_options)
+
+    let db = make_connection(connect_options)
         .await
         .expect("Fail to acquire database connection");
 
-    run_migrate(migrator, db, cli.command, cli.verbose)
+    run_migrate(migrator, &db, cli.command, cli.verbose)
         .await
         .unwrap_or_else(handle_error);
 }