diff --git a/sea-orm-cli/src/cli.rs b/sea-orm-cli/src/cli.rs index 4c771da00..0ac08ca57 100644 --- a/sea-orm-cli/src/cli.rs +++ b/sea-orm-cli/src/cli.rs @@ -202,6 +202,13 @@ pub enum GenerateSubcommands { )] max_connections: u32, + #[arg( + long, + default_value = "30", + long_help = "Acquire timeout in seconds of the connection used for schema discovery" + )] + acquire_timeout: u64, + #[arg( short = 'o', long, diff --git a/sea-orm-cli/src/commands/generate.rs b/sea-orm-cli/src/commands/generate.rs index 31e32af1a..f8bb4f11b 100644 --- a/sea-orm-cli/src/commands/generate.rs +++ b/sea-orm-cli/src/commands/generate.rs @@ -1,3 +1,4 @@ +use core::time; use sea_orm_codegen::{ DateTimeCrate as CodegenDateTimeCrate, EntityTransformer, EntityWriterContext, OutputFile, WithSerde, @@ -20,6 +21,7 @@ pub async fn run_generate_command( tables, ignore_tables, max_connections, + acquire_timeout, output_dir, database_schema, database_url, @@ -114,7 +116,9 @@ pub async fn run_generate_command( println!("Connecting to MySQL ..."); let connection = - sqlx_connect::(max_connections, url.as_str(), None).await?; + sqlx_connect::(max_connections, acquire_timeout, url.as_str(), None) + .await?; + println!("Discovering schema ..."); let schema_discovery = SchemaDiscovery::new(connection, database_name); let schema = schema_discovery.discover().await?; @@ -133,8 +137,14 @@ pub async fn run_generate_command( use sqlx::Sqlite; println!("Connecting to SQLite ..."); - let connection = - sqlx_connect::(max_connections, url.as_str(), None).await?; + let connection = sqlx_connect::( + max_connections, + acquire_timeout, + url.as_str(), + None, + ) + .await?; + println!("Discovering schema ..."); let schema_discovery = SchemaDiscovery::new(connection); let schema = schema_discovery @@ -157,9 +167,13 @@ pub async fn run_generate_command( println!("Connecting to Postgres ..."); let schema = database_schema.as_deref().unwrap_or("public"); - let connection = - sqlx_connect::(max_connections, url.as_str(), Some(schema)) - .await?; + let connection = sqlx_connect::( + max_connections, + acquire_timeout, + url.as_str(), + Some(schema), + ) + .await?; println!("Discovering schema ..."); let schema_discovery = SchemaDiscovery::new(connection, schema); let schema = schema_discovery.discover().await?; @@ -222,6 +236,7 @@ pub async fn run_generate_command( async fn sqlx_connect( max_connections: u32, + acquire_timeout: u64, url: &str, schema: Option<&str>, ) -> Result, Box> @@ -229,7 +244,9 @@ where DB: sqlx::Database, for<'a> &'a mut ::Connection: sqlx::Executor<'a>, { - let mut pool_options = sqlx::pool::PoolOptions::::new().max_connections(max_connections); + let mut pool_options = sqlx::pool::PoolOptions::::new() + .max_connections(max_connections) + .acquire_timeout(time::Duration::from_secs(acquire_timeout)); // Set search_path for Postgres, E.g. Some("public") by default // MySQL & SQLite connection initialize with schema `None` if let Some(schema) = schema {