Cont. "Delete all PostgreSQL types when calling fresh" (#991)

* 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 <ccw.billy.123@gmail.com>

* Refactoring

Co-authored-by: Denis Gavrilyuk <karpa4o4@gmail.com>

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::{
+    self, extension::postgres::Type, Alias, Expr, ForeignKey, Iden, JoinType, Query,
+    SelectStatement, SimpleExpr, Table,
+};
 use sea_orm::{
     ActiveModelTrait, ActiveValue, ColumnTrait, Condition, ConnectionTrait, DbBackend, DbConn,
     DbErr, EntityTrait, QueryFilter, QueryOrder, Schema, Statement,
@@ -146,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")?;
@@ -196,6 +181,21 @@ pub trait MigratorTrait: Send {
             info!("Table '{}' has been dropped", table_name);
         }
 
+        // Drop all types
+        if db_backend == DbBackend::Postgres {
+            info!("Dropping all types");
+            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")?;
+                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");
@@ -324,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
+}

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,
+}

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;
 mod m20220923_000001_seed_cake_table;
 
 pub struct Migrator;
@@ -14,6 +15,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),
             Box::new(m20220923_000001_seed_cake_table::Migration),
         ]
     }