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(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(migrator: M, make_connection: F) +where + M: MigratorTrait, + F: FnOnce(ConnectOptions) -> Fut, + Fut: Future>, { 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); }