Upstream Changes - 3 (#2155)

billy1624 · web-flow · commit a321f0a566ea · 2024-03-19T16:57:46.000+08:00
* Upstream Changes - 3

* refactor

* revert
diff --git 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
@@ -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<String>,
 
         #[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
@@ -113,7 +113,8 @@ pub async fn run_generate_command(
                     use sqlx::MySql;
 
                     println!("Connecting to MySQL ...");
-                    let connection = connect::<MySql>(max_connections, url.as_str(), None).await?;
+                    let connection =
+                        sqlx_connect::<MySql>(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::<Sqlite>(max_connections, url.as_str(), None).await?;
+                    let connection =
+                        sqlx_connect::<Sqlite>(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::<Postgres>(max_connections, url.as_str(), Some(schema)).await?;
+                        sqlx_connect::<Postgres>(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<DB>(
+async fn sqlx_connect<DB>(
     max_connections: u32,
     url: &str,
     schema: Option<&str>,
diff --git 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
@@ -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<String>,
 }
 
 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<T>(&mut self, schema: T) -> &mut Self
+    where
+        T: Into<String>,
+    {
+        self.schema = Some(schema.into());
+        self
+    }
+
+    pub fn get_schema(&self) -> &Option<String> {
+        &self.schema
+    }
 }
 
 /// Schema Creation
@@ -100,20 +114,7 @@ impl<'c> SchemaManager<'c> {
     where
         T: AsRef<str>,
     {
-        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<T, C>(&self, table: T, column: C) -> Result<bool, DbErr>
@@ -158,3 +159,27 @@ impl<'c> SchemaManager<'c> {
         res.try_get("", "has_index")
     }
 }
+
+pub(crate) async fn has_table<C, T>(
+    conn: &C,
+    _schema: Option<&str>,
+    table: T,
+) -> Result<bool, DbErr>
+where
+    C: ConnectionTrait,
+    T: AsRef<str>,
+{
+    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
@@ -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(),
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -164,10 +164,11 @@ pub trait MigratorTrait: Send {`
`164`	`164`	`C: ConnectionTrait,`
`165`	`165`	`{`
`166`	`166`	`let builder = db.get_database_backend();`
	`167`	`+ let table_name = Self::migration_table_name();`
`167`	`168`	`let schema = Schema::new(builder);`
`168`	`169`	`let mut stmt = schema`
`169`	`170`	`.create_table_from_entity(seaql_migrations::Entity)`
`170`		`- .table_name(Self::migration_table_name());`
	`171`	`+ .table_name(table_name);`
`171`	`172`	`stmt.if_not_exists();`
`172`	`173`	`db.execute(builder.build(&stmt)).await.map(\|_\| ())`
`173`	`174`	`}`
`@@ -441,9 +442,9 @@ where`
`441`	`442`	`C: ConnectionTrait,`
`442`	`443`	`{`
`443`	`444`	`match db.get_database_backend() {`
`444`		`- DbBackend::MySql => MySql::query_tables(),`
`445`		`- DbBackend::Postgres => Postgres::query_tables(),`
`446`		`- DbBackend::Sqlite => Sqlite::query_tables(),`
	`445`	`+ DbBackend::MySql => MySql.query_tables(),`
	`446`	`+ DbBackend::Postgres => Postgres.query_tables(),`
	`447`	`+ DbBackend::Sqlite => Sqlite.query_tables(),`
`447`	`448`	`}`
`448`	`449`	`}`
`449`	`450`