From a321f0a566eade8bbfe9a105db1f6f974b662de2 Mon Sep 17 00:00:00 2001 From: Billy Chan Date: Tue, 19 Mar 2024 16:57:46 +0800 Subject: [PATCH] Upstream Changes - 3 (#2155) * Upstream Changes - 3 * refactor * revert --- sea-orm-cli/Cargo.toml | 2 +- sea-orm-cli/src/cli.rs | 3 +- sea-orm-cli/src/commands/generate.rs | 16 ++++++--- sea-orm-migration/Cargo.toml | 2 +- sea-orm-migration/src/manager.rs | 53 ++++++++++++++++++++-------- sea-orm-migration/src/migrator.rs | 9 ++--- 6 files changed, 58 insertions(+), 27 deletions(-) diff --git a/sea-orm-cli/Cargo.toml b/sea-orm-cli/Cargo.toml index 5b821f342..0cbac34ba 100644 --- a/sea-orm-cli/Cargo.toml +++ b/sea-orm-cli/Cargo.toml @@ -38,7 +38,7 @@ clap = { version = "4.3", features = ["env", "derive"], optional = true } dotenvy = { version = "0.15", default-features = false, optional = true } async-std = { version = "1.9", default-features = false, features = ["attributes", "tokio1"], optional = true } sea-orm-codegen = { version = "=1.0.0-rc.2", path = "../sea-orm-codegen", default-features = false, optional = true } -sea-schema = { version = "0.15.0-rc.3" } +sea-schema = { version = "0.15.0-rc.3", git = "https://github.com/SeaQL/sea-schema", branch = "upstream-changes-2" } sqlx = { version = "0.7", default-features = false, features = ["mysql", "postgres"], optional = true } tracing-subscriber = { version = "0.3.17", default-features = false, features = ["env-filter", "fmt"] } tracing = { version = "0.1", default-features = false } diff --git a/sea-orm-cli/src/cli.rs b/sea-orm-cli/src/cli.rs index 3bc9ae9c6..4c771da00 100644 --- a/sea-orm-cli/src/cli.rs +++ b/sea-orm-cli/src/cli.rs @@ -214,12 +214,11 @@ pub enum GenerateSubcommands { short = 's', long, env = "DATABASE_SCHEMA", - default_value = "public", long_help = "Database schema\n \ - For MySQL, this argument is ignored.\n \ - For PostgreSQL, this argument is optional with default value 'public'." )] - database_schema: String, + database_schema: Option, #[arg(short = 'u', long, env = "DATABASE_URL", help = "Database URL")] database_url: String, diff --git a/sea-orm-cli/src/commands/generate.rs b/sea-orm-cli/src/commands/generate.rs index 613a541fb..37415bd96 100644 --- a/sea-orm-cli/src/commands/generate.rs +++ b/sea-orm-cli/src/commands/generate.rs @@ -113,7 +113,8 @@ pub async fn run_generate_command( use sqlx::MySql; println!("Connecting to MySQL ..."); - let connection = connect::(max_connections, url.as_str(), None).await?; + let connection = + sqlx_connect::(max_connections, url.as_str(), None).await?; println!("Discovering schema ..."); let schema_discovery = SchemaDiscovery::new(connection, database_name); let schema = schema_discovery.discover().await?; @@ -132,7 +133,8 @@ pub async fn run_generate_command( use sqlx::Sqlite; println!("Connecting to SQLite ..."); - let connection = connect::(max_connections, url.as_str(), None).await?; + let connection = + sqlx_connect::(max_connections, url.as_str(), None).await?; println!("Discovering schema ..."); let schema_discovery = SchemaDiscovery::new(connection); let schema = schema_discovery.discover().await?; @@ -151,9 +153,13 @@ pub async fn run_generate_command( use sqlx::Postgres; println!("Connecting to Postgres ..."); - let schema = &database_schema; + let schema = database_schema + .as_ref() + .map(|s| s.as_str()) + .unwrap_or("public"); let connection = - connect::(max_connections, url.as_str(), Some(schema)).await?; + sqlx_connect::(max_connections, url.as_str(), Some(schema)) + .await?; println!("Discovering schema ..."); let schema_discovery = SchemaDiscovery::new(connection, schema); let schema = schema_discovery.discover().await?; @@ -214,7 +220,7 @@ pub async fn run_generate_command( Ok(()) } -async fn connect( +async fn sqlx_connect( max_connections: u32, url: &str, schema: Option<&str>, diff --git a/sea-orm-migration/Cargo.toml b/sea-orm-migration/Cargo.toml index e1e7f8a20..f7a84efd2 100644 --- a/sea-orm-migration/Cargo.toml +++ b/sea-orm-migration/Cargo.toml @@ -25,7 +25,7 @@ clap = { version = "4.3", features = ["env", "derive"], optional = true } dotenvy = { version = "0.15", default-features = false, optional = true } sea-orm = { version = "1.0.0-rc.2", path = "../", default-features = false, features = ["macros"] } sea-orm-cli = { version = "1.0.0-rc.2", path = "../sea-orm-cli", default-features = false, optional = true } -sea-schema = { version = "0.15.0-rc.3" } +sea-schema = { version = "0.15.0-rc.3", git = "https://github.com/SeaQL/sea-schema", branch = "upstream-changes-2" } tracing = { version = "0.1", default-features = false, features = ["log"] } tracing-subscriber = { version = "0.3.17", default-features = false, features = ["env-filter", "fmt"] } futures = { version = "0.3", default-features = false, features = ["std"] } diff --git a/sea-orm-migration/src/manager.rs b/sea-orm-migration/src/manager.rs index 06bfe2c5d..9118ccdde 100644 --- a/sea-orm-migration/src/manager.rs +++ b/sea-orm-migration/src/manager.rs @@ -11,6 +11,7 @@ use sea_schema::{mysql::MySql, postgres::Postgres, probe::SchemaProbe, sqlite::S /// Helper struct for writing migration scripts in migration file pub struct SchemaManager<'c> { conn: SchemaManagerConnection<'c>, + schema: Option, } impl<'c> SchemaManager<'c> { @@ -20,6 +21,7 @@ impl<'c> SchemaManager<'c> { { Self { conn: conn.into_schema_manager_connection(), + schema: None, } } @@ -38,6 +40,18 @@ impl<'c> SchemaManager<'c> { pub fn get_connection(&self) -> &SchemaManagerConnection<'c> { &self.conn } + + pub fn set_schema(&mut self, schema: T) -> &mut Self + where + T: Into, + { + self.schema = Some(schema.into()); + self + } + + pub fn get_schema(&self) -> &Option { + &self.schema + } } /// Schema Creation @@ -100,20 +114,7 @@ impl<'c> SchemaManager<'c> { where T: AsRef, { - let stmt = match self.conn.get_database_backend() { - DbBackend::MySql => MySql.has_table(table), - DbBackend::Postgres => Postgres.has_table(table), - DbBackend::Sqlite => Sqlite.has_table(table), - }; - - let builder = self.conn.get_database_backend(); - let res = self - .conn - .query_one(builder.build(&stmt)) - .await? - .ok_or_else(|| DbErr::Custom("Failed to check table exists".to_owned()))?; - - res.try_get("", "has_table") + has_table(&self.conn, self.schema.as_deref(), table).await } pub async fn has_column(&self, table: T, column: C) -> Result @@ -158,3 +159,27 @@ impl<'c> SchemaManager<'c> { res.try_get("", "has_index") } } + +pub(crate) async fn has_table( + conn: &C, + _schema: Option<&str>, + table: T, +) -> Result +where + C: ConnectionTrait, + T: AsRef, +{ + let stmt = match conn.get_database_backend() { + DbBackend::MySql => MySql.has_table(table), + DbBackend::Postgres => Postgres.has_table(table), + DbBackend::Sqlite => Sqlite.has_table(table), + }; + + let builder = conn.get_database_backend(); + let res = conn + .query_one(builder.build(&stmt)) + .await? + .ok_or_else(|| DbErr::Custom("Failed to check table exists".to_owned()))?; + + res.try_get("", "has_table") +} diff --git a/sea-orm-migration/src/migrator.rs b/sea-orm-migration/src/migrator.rs index 587951ed1..d58802884 100644 --- a/sea-orm-migration/src/migrator.rs +++ b/sea-orm-migration/src/migrator.rs @@ -164,10 +164,11 @@ pub trait MigratorTrait: Send { C: ConnectionTrait, { let builder = db.get_database_backend(); + let table_name = Self::migration_table_name(); let schema = Schema::new(builder); let mut stmt = schema .create_table_from_entity(seaql_migrations::Entity) - .table_name(Self::migration_table_name()); + .table_name(table_name); stmt.if_not_exists(); db.execute(builder.build(&stmt)).await.map(|_| ()) } @@ -441,9 +442,9 @@ where C: ConnectionTrait, { match db.get_database_backend() { - DbBackend::MySql => MySql::query_tables(), - DbBackend::Postgres => Postgres::query_tables(), - DbBackend::Sqlite => Sqlite::query_tables(), + DbBackend::MySql => MySql.query_tables(), + DbBackend::Postgres => Postgres.query_tables(), + DbBackend::Sqlite => Sqlite.query_tables(), } }