From 5e9a849a87f0d2bbecc9fbfd0be560421af4f7a3 Mon Sep 17 00:00:00 2001 From: Denis Gavrilyuk Date: Thu, 25 Aug 2022 08:03:18 +0300 Subject: [PATCH 1/2] Delete all PostgreSQL types when calling fresh (#765) (#864) * Delete all PostgreSQL types when calling fresh (#765) * Test create db enum migration Co-authored-by: Billy Chan --- sea-orm-migration/src/migrator.rs | 38 ++++++++++++- .../m20220118_000004_create_tea_enum.rs | 53 +++++++++++++++++++ sea-orm-migration/tests/migrator/mod.rs | 2 + 3 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 sea-orm-migration/tests/migrator/m20220118_000004_create_tea_enum.rs diff --git a/sea-orm-migration/src/migrator.rs b/sea-orm-migration/src/migrator.rs index 5239d7598..8b4dc105a 100644 --- a/sea-orm-migration/src/migrator.rs +++ b/sea-orm-migration/src/migrator.rs @@ -3,7 +3,10 @@ use std::fmt::Display; use std::time::SystemTime; use tracing::info; -use sea_orm::sea_query::{Alias, Expr, ForeignKey, Query, SelectStatement, SimpleExpr, Table}; +use sea_orm::sea_query::{ + extension::postgres::Type, Alias, Expr, ForeignKey, JoinType, Query, SelectStatement, + SimpleExpr, Table, +}; use sea_orm::{ ActiveModelTrait, ActiveValue, ColumnTrait, Condition, ConnectionTrait, DbBackend, DbConn, DbErr, EntityTrait, QueryFilter, QueryOrder, Schema, Statement, @@ -196,6 +199,39 @@ pub trait MigratorTrait: Send { info!("Table '{}' has been dropped", table_name); } + // Drop all types + if db_backend == DbBackend::Postgres { + info!("Dropping all types"); + + let mut stmt = Query::select(); + stmt.columns([Alias::new("typname")]) + .from(Alias::new("pg_type")) + .join( + JoinType::LeftJoin, + Alias::new("pg_namespace"), + Expr::tbl(Alias::new("pg_namespace"), Alias::new("oid")) + .equals(Alias::new("pg_type"), Alias::new("typnamespace")), + ) + .cond_where( + Condition::all() + .add( + Expr::expr(get_current_schema(db)) + .equals(Alias::new("pg_namespace"), Alias::new("nspname")), + ) + .add(Expr::tbl(Alias::new("pg_type"), Alias::new("typelem")).eq(0)), + ); + + let rows = db.query_all(db_backend.build(&stmt)).await?; + for row in rows { + let type_name: String = row.try_get("", "typname")?; + info!("Dropping type '{}'", type_name); + let mut stmt = Type::drop(); + stmt.name(Alias::new(&type_name as &str)); + db.execute(db_backend.build(&stmt)).await?; + info!("Type '{}' has been dropped", type_name); + } + } + // Restore the foreign key check if db_backend == DbBackend::Sqlite { info!("Restoring foreign key check"); diff --git a/sea-orm-migration/tests/migrator/m20220118_000004_create_tea_enum.rs b/sea-orm-migration/tests/migrator/m20220118_000004_create_tea_enum.rs new file mode 100644 index 000000000..ba9b68e6b --- /dev/null +++ b/sea-orm-migration/tests/migrator/m20220118_000004_create_tea_enum.rs @@ -0,0 +1,53 @@ +use sea_orm_migration::prelude::{sea_query::extension::postgres::Type, *}; +use sea_orm_migration::sea_orm::{ConnectionTrait, DbBackend}; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + let db = manager.get_connection(); + + match db.get_database_backend() { + DbBackend::MySql | DbBackend::Sqlite => {} + DbBackend::Postgres => { + manager + .create_type( + Type::create() + .as_enum(Tea::Table) + .values([Tea::EverydayTea, Tea::BreakfastTea]) + .to_owned(), + ) + .await?; + } + } + + Ok(()) + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + let db = manager.get_connection(); + + match db.get_database_backend() { + DbBackend::MySql | DbBackend::Sqlite => {} + DbBackend::Postgres => { + manager + .drop_type(Type::drop().name(Tea::Table).to_owned()) + .await?; + } + } + + Ok(()) + } +} + +/// Learn more at https://docs.rs/sea-query#iden +#[derive(Iden)] +pub enum Tea { + Table, + #[iden = "EverydayTea"] + EverydayTea, + #[iden = "BreakfastTea"] + BreakfastTea, +} diff --git a/sea-orm-migration/tests/migrator/mod.rs b/sea-orm-migration/tests/migrator/mod.rs index fdb92d9eb..135670807 100644 --- a/sea-orm-migration/tests/migrator/mod.rs +++ b/sea-orm-migration/tests/migrator/mod.rs @@ -3,6 +3,7 @@ use sea_orm_migration::prelude::*; mod m20220118_000001_create_cake_table; mod m20220118_000002_create_fruit_table; mod m20220118_000003_seed_cake_table; +mod m20220118_000004_create_tea_enum; pub struct Migrator; @@ -13,6 +14,7 @@ impl MigratorTrait for Migrator { Box::new(m20220118_000001_create_cake_table::Migration), Box::new(m20220118_000002_create_fruit_table::Migration), Box::new(m20220118_000003_seed_cake_table::Migration), + Box::new(m20220118_000004_create_tea_enum::Migration), ] } } From ed2f4b19c3c2499f997d3cc7b60139d003202ed7 Mon Sep 17 00:00:00 2001 From: Billy Chan Date: Thu, 25 Aug 2022 13:56:08 +0800 Subject: [PATCH 2/2] Refactoring --- sea-orm-migration/src/migrator.rs | 120 ++++++++++++++++++++---------- 1 file changed, 80 insertions(+), 40 deletions(-) diff --git a/sea-orm-migration/src/migrator.rs b/sea-orm-migration/src/migrator.rs index 8b4dc105a..89a224756 100644 --- a/sea-orm-migration/src/migrator.rs +++ b/sea-orm-migration/src/migrator.rs @@ -4,8 +4,8 @@ use std::time::SystemTime; use tracing::info; use sea_orm::sea_query::{ - extension::postgres::Type, Alias, Expr, ForeignKey, JoinType, Query, SelectStatement, - SimpleExpr, Table, + self, extension::postgres::Type, Alias, Expr, ForeignKey, Iden, JoinType, Query, + SelectStatement, SimpleExpr, Table, }; use sea_orm::{ ActiveModelTrait, ActiveValue, ColumnTrait, Condition, ConnectionTrait, DbBackend, DbConn, @@ -149,25 +149,7 @@ pub trait MigratorTrait: Send { // Drop all foreign keys if db_backend == DbBackend::MySql { info!("Dropping all foreign keys"); - let mut stmt = Query::select(); - stmt.columns([Alias::new("TABLE_NAME"), Alias::new("CONSTRAINT_NAME")]) - .from(( - Alias::new("information_schema"), - Alias::new("table_constraints"), - )) - .cond_where( - Condition::all() - .add( - Expr::expr(get_current_schema(db)).equals( - Alias::new("table_constraints"), - Alias::new("table_schema"), - ), - ) - .add(Expr::expr(Expr::value("FOREIGN KEY")).equals( - Alias::new("table_constraints"), - Alias::new("constraint_type"), - )), - ); + let stmt = query_mysql_foreign_keys(db); let rows = db.query_all(db_backend.build(&stmt)).await?; for row in rows.into_iter() { let constraint_name: String = row.try_get("", "CONSTRAINT_NAME")?; @@ -202,25 +184,7 @@ pub trait MigratorTrait: Send { // Drop all types if db_backend == DbBackend::Postgres { info!("Dropping all types"); - - let mut stmt = Query::select(); - stmt.columns([Alias::new("typname")]) - .from(Alias::new("pg_type")) - .join( - JoinType::LeftJoin, - Alias::new("pg_namespace"), - Expr::tbl(Alias::new("pg_namespace"), Alias::new("oid")) - .equals(Alias::new("pg_type"), Alias::new("typnamespace")), - ) - .cond_where( - Condition::all() - .add( - Expr::expr(get_current_schema(db)) - .equals(Alias::new("pg_namespace"), Alias::new("nspname")), - ) - .add(Expr::tbl(Alias::new("pg_type"), Alias::new("typelem")).eq(0)), - ); - + let stmt = query_pg_types(db); let rows = db.query_all(db_backend.build(&stmt)).await?; for row in rows { let type_name: String = row.try_get("", "typname")?; @@ -360,3 +324,79 @@ pub(crate) fn get_current_schema(db: &DbConn) -> SimpleExpr { DbBackend::Sqlite => unimplemented!(), } } + +#[derive(Iden)] +enum InformationSchema { + #[iden = "information_schema"] + Schema, + #[iden = "TABLE_NAME"] + TableName, + #[iden = "CONSTRAINT_NAME"] + ConstraintName, + TableConstraints, + TableSchema, + ConstraintType, +} + +fn query_mysql_foreign_keys(db: &DbConn) -> SelectStatement { + let mut stmt = Query::select(); + stmt.columns([ + InformationSchema::TableName, + InformationSchema::ConstraintName, + ]) + .from(( + InformationSchema::Schema, + InformationSchema::TableConstraints, + )) + .cond_where( + Condition::all() + .add(Expr::expr(get_current_schema(db)).equals( + InformationSchema::TableConstraints, + InformationSchema::TableSchema, + )) + .add( + Expr::tbl( + InformationSchema::TableConstraints, + InformationSchema::ConstraintType, + ) + .eq("FOREIGN KEY"), + ), + ); + stmt +} + +#[derive(Iden)] +enum PgType { + Table, + Typname, + Typnamespace, + Typelem, +} + +#[derive(Iden)] +enum PgNamespace { + Table, + Oid, + Nspname, +} + +fn query_pg_types(db: &DbConn) -> SelectStatement { + let mut stmt = Query::select(); + stmt.column(PgType::Typname) + .from(PgType::Table) + .join( + JoinType::LeftJoin, + PgNamespace::Table, + Expr::tbl(PgNamespace::Table, PgNamespace::Oid) + .equals(PgType::Table, PgType::Typnamespace), + ) + .cond_where( + Condition::all() + .add( + Expr::expr(get_current_schema(db)) + .equals(PgNamespace::Table, PgNamespace::Nspname), + ) + .add(Expr::tbl(PgType::Table, PgType::Typelem).eq(0)), + ); + stmt +}