diff --git a/sea-orm-migration/src/manager.rs b/sea-orm-migration/src/manager.rs index fec6a22d..727e8340 100644 --- a/sea-orm-migration/src/manager.rs +++ b/sea-orm-migration/src/manager.rs @@ -5,8 +5,7 @@ use sea_orm::sea_query::{ TableRenameStatement, TableTruncateStatement, }; use sea_orm::{Condition, ConnectionTrait, DbBackend, DbConn, DbErr, Statement, StatementBuilder}; - -use super::query_tables; +use sea_schema::{mysql::MySql, postgres::Postgres, probe::SchemaProbe, sqlite::Sqlite}; /// Helper struct for writing migration scripts in migration file pub struct SchemaManager<'c> { @@ -95,11 +94,11 @@ impl<'c> SchemaManager<'c> { where T: AsRef, { - let mut stmt = Query::select(); - let mut subquery = query_tables(self.conn); - subquery.cond_where(Expr::col(Alias::new("table_name")).eq(table.as_ref())); - stmt.expr_as(Expr::cust("COUNT(*)"), Alias::new("rows")) - .from_subquery(subquery, Alias::new("subquery")); + 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 @@ -107,9 +106,8 @@ impl<'c> SchemaManager<'c> { .query_one(builder.build(&stmt)) .await? .ok_or_else(|| DbErr::Custom("Fail to check table exists".to_owned()))?; - let rows: i64 = res.try_get("", "rows")?; - Ok(rows > 0) + Ok(res.try_get("", "has_table")?) } pub async fn has_column(&self, table: T, column: C) -> Result @@ -117,51 +115,19 @@ impl<'c> SchemaManager<'c> { T: AsRef, C: AsRef, { - let db_backend = self.conn.get_database_backend(); - let found = match db_backend { - DbBackend::MySql | DbBackend::Postgres => { - let schema_name = match db_backend { - DbBackend::MySql => "DATABASE()", - DbBackend::Postgres => "CURRENT_SCHEMA()", - DbBackend::Sqlite => unreachable!(), - }; - let mut stmt = Query::select(); - stmt.expr_as(Expr::cust("COUNT(*)"), Alias::new("rows")) - .from((Alias::new("information_schema"), Alias::new("columns"))) - .cond_where( - Condition::all() - .add( - Expr::expr(Expr::cust(schema_name)) - .equals(Alias::new("columns"), Alias::new("table_schema")), - ) - .add(Expr::col(Alias::new("table_name")).eq(table.as_ref())) - .add(Expr::col(Alias::new("column_name")).eq(column.as_ref())), - ); - - let res = self - .conn - .query_one(db_backend.build(&stmt)) - .await? - .ok_or_else(|| DbErr::Custom("Fail to check column exists".to_owned()))?; - let rows: i64 = res.try_get("", "rows")?; - rows > 0 - } - DbBackend::Sqlite => { - let stmt = Statement::from_string( - db_backend, - format!("PRAGMA table_info({})", table.as_ref()), - ); - let results = self.conn.query_all(stmt).await?; - let mut found = false; - for res in results { - let name: String = res.try_get("", "name")?; - if name.as_str() == column.as_ref() { - found = true; - } - } - found - } + let stmt = match self.conn.get_database_backend() { + DbBackend::MySql => MySql::has_column(table, column), + DbBackend::Postgres => Postgres::has_column(table, column), + DbBackend::Sqlite => Sqlite::has_column(table, column), }; - Ok(found) + + let builder = self.conn.get_database_backend(); + let res = self + .conn + .query_one(builder.build(&stmt)) + .await? + .ok_or_else(|| DbErr::Custom("Fail to check table exists".to_owned()))?; + + Ok(res.try_get("", "has_column")?) } } diff --git a/sea-orm-migration/tests/main.rs b/sea-orm-migration/tests/main.rs index 9b7c27c7..4131bae4 100644 --- a/sea-orm-migration/tests/main.rs +++ b/sea-orm-migration/tests/main.rs @@ -1,15 +1,27 @@ mod migrator; use migrator::Migrator; -use sea_orm::{Database, DbErr}; +use sea_orm::{ConnectionTrait, Database, DbBackend, DbErr, Statement}; use sea_orm_migration::prelude::*; #[async_std::test] async fn main() -> Result<(), DbErr> { let url = std::env::var("DATABASE_URL").expect("Environment variable 'DATABASE_URL' not set"); let db_name = "sea_orm_migration"; - let url = format!("{}/{}", url, db_name); - let db = &Database::connect(&url).await?; + let db = Database::connect(&url).await?; + let db = &match db.get_database_backend() { + DbBackend::MySql | DbBackend::Postgres => { + db.execute(Statement::from_string( + db.get_database_backend(), + format!("CREATE DATABASE IF NOT EXISTS `{}`;", db_name), + )) + .await?; + + let url = format!("{}/{}", url, db_name); + Database::connect(&url).await? + } + DbBackend::Sqlite => db, + }; let manager = SchemaManager::new(db); println!("\nMigrator::status"); @@ -59,6 +71,9 @@ async fn main() -> Result<(), DbErr> { assert!(manager.has_table("cake").await?); assert!(manager.has_table("fruit").await?); + assert!(manager.has_column("cake", "name").await?); + assert!(manager.has_column("fruit", "cake_id").await?); + println!("\nMigrator::down"); Migrator::down(db, None).await?;