Add max_connections option to CLI

sea-orm-cli/src/bin/main.rs

@@ -1,4 +1,3 @@
-
 use dotenv::dotenv;
 use sea_orm_cli::*;
 

sea-orm-cli/src/bin/sea.rs

@@ -16,4 +16,4 @@ async fn main() {
         ("migrate", Some(matches)) => run_migrate_command(matches).unwrap_or_else(handle_error),
         _ => unreachable!("You should never see this message"),
     }
-}
+}

sea-orm-cli/src/cli.rs

@@ -70,6 +70,13 @@ pub fn build_cli() -> App<'static, 'static> {
                         .help("Automatically derive serde Serialize / Deserialize traits for the entity (none, serialize, deserialize, both)")
                         .takes_value(true)
                         .default_value("none")
+                )
+                .arg(
+                    Arg::with_name("MAX_CONNECTIONS")
+                        .long("max-connections")
+                        .help("The maximum amount of connections to use when connecting to the database.")
+                        .takes_value(true)
+                        .default_value("1")
                 ),
         )
         .setting(AppSettings::SubcommandRequiredElseHelp);
@@ -108,4 +115,3 @@ pub fn build_cli() -> App<'static, 'static> {
         )
         .setting(AppSettings::SubcommandRequiredElseHelp)
 }
-

sea-orm-cli/src/commands.rs

@@ -1,4 +1,3 @@
-
 use clap::ArgMatches;
 use sea_orm_codegen::{EntityTransformer, OutputFile, WithSerde};
 use std::{error::Error, fmt::Display, fs, io::Write, path::Path, process::Command, str::FromStr};
@@ -22,6 +21,12 @@ pub async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dy
                     .try_init();
             }
 
+            let max_connections = args
+                .value_of("MAX_CONNECTIONS")
+                .map(str::parse::<u32>)
+                .transpose()?
+                .unwrap();
+
             // The database should be a valid URL that can be parsed
             // protocol://username:password@host/database_name
             let url = Url::parse(
@@ -98,9 +103,9 @@ pub async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dy
             let table_stmts = match url.scheme() {
                 "mysql" => {
                     use sea_schema::mysql::discovery::SchemaDiscovery;
-                    use sqlx::MySqlPool;
+                    use sqlx::MySql;
 
-                    let connection = MySqlPool::connect(url.as_str()).await?;
+                    let connection = connect::<MySql>(max_connections, url.as_str()).await?;
                     let schema_discovery = SchemaDiscovery::new(connection, database_name);
                     let schema = schema_discovery.discover().await;
                     schema
@@ -113,9 +118,9 @@ pub async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dy
                 }
                 "sqlite" => {
                     use sea_schema::sqlite::SchemaDiscovery;
-                    use sqlx::SqlitePool;
+                    use sqlx::Sqlite;
 
-                    let connection = SqlitePool::connect(url.as_str()).await?;
+                    let connection = connect::<Sqlite>(max_connections, url.as_str()).await?;
                     let schema_discovery = SchemaDiscovery::new(connection);
                     let schema = schema_discovery.discover().await?;
                     schema
@@ -128,10 +133,10 @@ pub async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dy
                 }
                 "postgres" | "postgresql" => {
                     use sea_schema::postgres::discovery::SchemaDiscovery;
-                    use sqlx::PgPool;
+                    use sqlx::Postgres;
 
                     let schema = args.value_of("DATABASE_SCHEMA").unwrap_or("public");
-                    let connection = PgPool::connect(url.as_str()).await?;
+                    let connection = connect::<Postgres>(max_connections, url.as_str()).await?;
                     let schema_discovery = SchemaDiscovery::new(connection, schema);
                     let schema = schema_discovery.discover().await;
                     schema
@@ -171,6 +176,17 @@ pub async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dy
     Ok(())
 }
 
+async fn connect<DB>(max_connections: u32, url: &str) -> Result<sqlx::Pool<DB>, Box<dyn Error>>
+where
+    DB: sqlx::Database,
+{
+    sqlx::pool::PoolOptions::<DB>::new()
+        .max_connections(max_connections)
+        .connect(url)
+        .await
+        .map_err(Into::into)
+}
+
 pub fn run_migrate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dyn Error>> {
     let migrate_subcommand = matches.subcommand();
     // If it's `migrate init`
@@ -257,8 +273,8 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use clap::AppSettings;
     use crate::cli;
+    use clap::AppSettings;
 
     #[test]
     #[should_panic(

sea-orm-cli/src/lib.rs

@@ -2,4 +2,4 @@ pub mod cli;
 pub mod commands;
 
 pub use cli::*;
-pub use commands::*;
+pub use commands::*;