From 80c72004d15afa8c4e6920abab76e91d55f5161a Mon Sep 17 00:00:00 2001 From: Billy Chan Date: Thu, 21 Oct 2021 15:40:22 +0800 Subject: [PATCH] Try Postgres enum --- Cargo.toml | 2 +- src/entity/active_enum.rs | 2 +- src/entity/column.rs | 7 ++++++ src/query/insert.rs | 19 +++++++++++---- src/query/select.rs | 20 ++++++++++++--- src/query/update.rs | 12 +++++++-- tests/active_enum_tests.rs | 7 +++++- tests/common/features/schema.rs | 43 +++++++++++++++++++++++++++++++-- 8 files changed, 96 insertions(+), 16 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 71690374..42f97b62 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,7 +30,7 @@ futures-util = { version = "^0.3" } log = { version = "^0.4", optional = true } rust_decimal = { version = "^1", optional = true } sea-orm-macros = { version = "^0.3.0", path = "sea-orm-macros", optional = true } -sea-query = { version = "^0.18.0", features = ["thread-safe"] } +sea-query = { version = "^0.18.0", git = "https://github.com/SeaQL/sea-query.git", branch = "sea-orm/active-enum-1", features = ["thread-safe"] } sea-strum = { version = "^0.21", features = ["derive", "sea-orm"] } serde = { version = "^1.0", features = ["derive"] } serde_json = { version = "^1", optional = true } diff --git a/src/entity/active_enum.rs b/src/entity/active_enum.rs index edca0aac..b42b31f7 100644 --- a/src/entity/active_enum.rs +++ b/src/entity/active_enum.rs @@ -71,7 +71,7 @@ use sea_query::{Nullable, Value, ValueType}; /// use sea_orm::entity::prelude::*; /// /// // Define the `Category` active enum -/// #[derive(Debug, Clone, PartialEq, DeriveActiveEnum)] +/// #[derive(Debug, Clone, PartialEq, EnumIter, DeriveActiveEnum)] /// #[sea_orm(rs_type = "String", db_type = "String(Some(1))")] /// pub enum Category { /// #[sea_orm(string_value = "B")] diff --git a/src/entity/column.rs b/src/entity/column.rs index 384be24a..32fe030c 100644 --- a/src/entity/column.rs +++ b/src/entity/column.rs @@ -242,6 +242,13 @@ impl ColumnType { indexed: false, } } + + pub(crate) fn get_enum_name(&self) -> Option<&String> { + match self { + ColumnType::Enum(s) => Some(s), + _ => None, + } + } } impl ColumnDef { diff --git a/src/query/insert.rs b/src/query/insert.rs index 5e504a0c..f7f7c7d3 100644 --- a/src/query/insert.rs +++ b/src/query/insert.rs @@ -1,9 +1,9 @@ use crate::{ - ActiveModelTrait, EntityName, EntityTrait, IntoActiveModel, Iterable, PrimaryKeyTrait, - QueryTrait, + ActiveModelTrait, ColumnTrait, EntityName, EntityTrait, IntoActiveModel, Iterable, + PrimaryKeyTrait, QueryTrait, }; use core::marker::PhantomData; -use sea_query::{InsertStatement, ValueTuple}; +use sea_query::{Alias, Expr, Func, InsertStatement, ValueTuple}; #[derive(Debug)] pub struct Insert @@ -124,6 +124,9 @@ where for (idx, col) in ::Column::iter().enumerate() { let av = am.take(col); let av_has_val = av.is_set() || av.is_unchanged(); + let col_def = col.def(); + let enum_name = col_def.get_column_type().get_enum_name(); + if columns_empty { self.columns.push(av_has_val); } else if self.columns[idx] != av_has_val { @@ -131,11 +134,17 @@ where } if av_has_val { columns.push(col); - values.push(av.into_value().unwrap()); + let val = av.into_value().unwrap(); + let expr = if let Some(enum_name) = enum_name { + Func::cast_as(val, Alias::new(&enum_name)) + } else { + Expr::val(val).into() + }; + values.push(expr); } } self.query.columns(columns); - self.query.values_panic(values); + self.query.exprs_panic(values); self } diff --git a/src/query/select.rs b/src/query/select.rs index 1b0c93c3..5a2e0cdc 100644 --- a/src/query/select.rs +++ b/src/query/select.rs @@ -2,7 +2,7 @@ use crate::{ColumnTrait, EntityTrait, Iterable, QueryFilter, QueryOrder, QuerySe use core::fmt::Debug; use core::marker::PhantomData; pub use sea_query::JoinType; -use sea_query::{DynIden, IntoColumnRef, SeaRc, SelectStatement, SimpleExpr}; +use sea_query::{Alias, DynIden, Expr, Func, IntoColumnRef, SeaRc, SelectStatement, SimpleExpr}; #[derive(Clone, Debug)] pub struct Select @@ -109,13 +109,25 @@ where } fn prepare_select(mut self) -> Self { - self.query.columns(self.column_list()); + self.query.exprs(self.column_list()); self } - fn column_list(&self) -> Vec<(DynIden, E::Column)> { + fn column_list(&self) -> Vec { let table = SeaRc::new(E::default()) as DynIden; - E::Column::iter().map(|col| (table.clone(), col)).collect() + E::Column::iter() + .map(|col| { + let col_def = col.def(); + let enum_name = col_def.get_column_type().get_enum_name(); + let col_expr = Expr::tbl(table.clone(), col); + let col_expr = if let Some(_) = enum_name { + Func::cast_expr_as(col_expr, Alias::new("text")) + } else { + col_expr.into() + }; + col_expr + }) + .collect() } fn prepare_from(mut self) -> Self { diff --git a/src/query/update.rs b/src/query/update.rs index fec7757c..2f9bdb3b 100644 --- a/src/query/update.rs +++ b/src/query/update.rs @@ -3,7 +3,7 @@ use crate::{ QueryTrait, }; use core::marker::PhantomData; -use sea_query::{IntoIden, SimpleExpr, UpdateStatement}; +use sea_query::{Alias, Expr, Func, IntoIden, SimpleExpr, UpdateStatement}; #[derive(Clone, Debug)] pub struct Update; @@ -104,9 +104,17 @@ where if ::PrimaryKey::from_column(col).is_some() { continue; } + let col_def = col.def(); + let enum_name = col_def.get_column_type().get_enum_name(); let av = self.model.get(col); if av.is_set() { - self.query.value(col, av.unwrap()); + let val = av.into_value().unwrap(); + let expr = if let Some(enum_name) = enum_name { + Func::cast_as(val, Alias::new(&enum_name)) + } else { + Expr::val(val).into() + }; + self.query.value_expr(col, expr); } } self diff --git a/tests/active_enum_tests.rs b/tests/active_enum_tests.rs index b0c72a82..ca9bf7d9 100644 --- a/tests/active_enum_tests.rs +++ b/tests/active_enum_tests.rs @@ -40,7 +40,7 @@ pub async fn insert_active_enum(db: &DatabaseConnection) -> Result<(), DbErr> { } ); - ActiveModel { + let am = ActiveModel { category: Set(Some(Category::Big)), color: Set(Some(Color::Black)), tea: Set(Some(Tea::EverydayTea)), @@ -59,5 +59,10 @@ pub async fn insert_active_enum(db: &DatabaseConnection) -> Result<(), DbErr> { } ); + let res = am.delete(db).await?; + + assert_eq!(res.rows_affected, 1); + assert_eq!(Entity::find().one(db).await?, None); + Ok(()) } diff --git a/tests/common/features/schema.rs b/tests/common/features/schema.rs index 823ccdfd..ffb47f26 100644 --- a/tests/common/features/schema.rs +++ b/tests/common/features/schema.rs @@ -3,7 +3,7 @@ pub use super::super::bakery_chain::*; use super::*; use crate::common::setup::create_table; use sea_orm::{error::*, sea_query, DatabaseConnection, DbConn, ExecResult}; -use sea_query::{ColumnDef, ForeignKeyCreateStatement}; +use sea_query::{Alias, ColumnDef, ForeignKeyCreateStatement}; pub async fn create_tables(db: &DatabaseConnection) -> Result<(), DbErr> { create_log_table(db).await?; @@ -103,6 +103,8 @@ pub async fn create_self_join_table(db: &DbConn) -> Result { } pub async fn create_active_enum_table(db: &DbConn) -> Result { + let tea_enum = Alias::new("tea"); + let stmt = sea_query::Table::create() .table(active_enum::Entity) .col( @@ -114,8 +116,45 @@ pub async fn create_active_enum_table(db: &DbConn) -> Result ) .col(ColumnDef::new(active_enum::Column::Category).string_len(1)) .col(ColumnDef::new(active_enum::Column::Color).integer()) - // .col(ColumnDef::new(active_enum::Column::Tea).custom(Alias::new("tea"))) + .col(ColumnDef::new(active_enum::Column::Tea).custom(tea_enum.clone())) .to_owned(); + match db { + #[cfg(feature = "sqlx-postgres")] + DatabaseConnection::SqlxPostgresPoolConnection(_) => { + use sea_orm::{ConnectionTrait, Statement}; + use sea_query::{extension::postgres::Type, PostgresQueryBuilder}; + + let drop_type_stmt = Type::drop() + .name(tea_enum.clone()) + .cascade() + .if_exists() + .to_owned(); + let (sql, values) = drop_type_stmt.build(PostgresQueryBuilder); + let stmt = Statement::from_sql_and_values(db.get_database_backend(), &sql, values); + db.execute(stmt).await?; + + let create_type_stmt = Type::create() + .as_enum(tea_enum) + .values(vec![Alias::new("EverydayTea"), Alias::new("BreakfastTea")]) + .to_owned(); + + // FIXME: This is not working + { + let (sql, values) = create_type_stmt.build(PostgresQueryBuilder); + let _stmt = Statement::from_sql_and_values(db.get_database_backend(), &sql, values); + } + + // But this is working... + let stmt = Statement::from_string( + db.get_database_backend(), + create_type_stmt.to_string(PostgresQueryBuilder), + ); + + db.execute(stmt).await?; + } + _ => {} + } + create_table(db, &stmt, ActiveEnum).await }