caisin/sea-orm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[CLI] Generate entity Postgres connection with schema search path (#1212)

Browse Source
This commit is contained in:
Billy Chan 2022-11-09 15:01:44 +08:00 committed by GitHub
parent b9a07ee538
commit d6ca491d80
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 25 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
sea-orm-cli/src/commands/generate.rs
View File

@ -115,7 +115,7 @@ pub async fn run_generate_command(
use sea_schema::mysql::discovery::SchemaDiscovery; use sea_schema::mysql::discovery::SchemaDiscovery;
use sqlx::MySql; use sqlx::MySql;
let connection = connect::<MySql>(max_connections, url.as_str()).await?; let connection = connect::<MySql>(max_connections, url.as_str(), None).await?;
let schema_discovery = SchemaDiscovery::new(connection, database_name); let schema_discovery = SchemaDiscovery::new(connection, database_name);
let schema = schema_discovery.discover().await; let schema = schema_discovery.discover().await;
let table_stmts = schema let table_stmts = schema
@ -132,7 +132,7 @@ pub async fn run_generate_command(
use sea_schema::sqlite::discovery::SchemaDiscovery; use sea_schema::sqlite::discovery::SchemaDiscovery;
use sqlx::Sqlite; use sqlx::Sqlite;
let connection = connect::<Sqlite>(max_connections, url.as_str()).await?; let connection = connect::<Sqlite>(max_connections, url.as_str(), None).await?;
let schema_discovery = SchemaDiscovery::new(connection); let schema_discovery = SchemaDiscovery::new(connection);
let schema = schema_discovery.discover().await?; let schema = schema_discovery.discover().await?;
let table_stmts = schema let table_stmts = schema
@ -150,7 +150,8 @@ pub async fn run_generate_command(
use sqlx::Postgres; use sqlx::Postgres;
let schema = &database_schema; let schema = &database_schema;
let connection = connect::<Postgres>(max_connections, url.as_str()).await?; let connection =
connect::<Postgres>(max_connections, url.as_str(), Some(schema)).await?;
let schema_discovery = SchemaDiscovery::new(connection, schema); let schema_discovery = SchemaDiscovery::new(connection, schema);
let schema = schema_discovery.discover().await; let schema = schema_discovery.discover().await;
let table_stmts = schema let table_stmts = schema
@ -198,15 +199,30 @@ pub async fn run_generate_command(
Ok(()) Ok(())
} }
async fn connect<DB>(max_connections: u32, url: &str) -> Result<sqlx::Pool<DB>, Box<dyn Error>> async fn connect<DB>(
max_connections: u32,
url: &str,
schema: Option<&str>,
) -> Result<sqlx::Pool<DB>, Box<dyn Error>>
where where
DB: sqlx::Database, DB: sqlx::Database,
for<'a> &'a mut <DB as sqlx::Database>::Connection: sqlx::Executor<'a>,
{ {
sqlx::pool::PoolOptions::<DB>::new() let mut pool_options = sqlx::pool::PoolOptions::<DB>::new().max_connections(max_connections);
.max_connections(max_connections) // Set search_path for Postgres, E.g. Some("public") by default
.connect(url) // MySQL & SQLite connection initialize with schema `None`
.await if let Some(schema) = schema {
.map_err(Into::into) let sql = format!("SET search_path = '{}'", schema);
pool_options = pool_options.after_connect(move |conn, _| {
let sql = sql.clone();
Box::pin(async move {
sqlx::Executor::execute(conn, sql.as_str())
.await
.map(|_| ())
})
});
}
pool_options.connect(url).await.map_err(Into::into)
} }
impl From<DateTimeCrate> for CodegenDateTimeCrate { impl From<DateTimeCrate> for CodegenDateTimeCrate {