Use ActiveEnum field as composite primary key (#1414)

* Use ActiveEnum field as part of composite primary key

* Try unique constraint failed

src/entity/active_enum.rs

@@ -1,4 +1,4 @@
-use crate::{ColumnDef, DbErr, Iterable, QueryResult, TryGetError, TryGetable};
+use crate::{ColumnDef, DbErr, Iterable, QueryResult, TryFromU64, TryGetError, TryGetable};
 use sea_query::{DynIden, Expr, Nullable, SimpleExpr, Value, ValueType};
 
 /// A Rust representation of enum defined in database.
@@ -160,6 +160,17 @@ where
     }
 }
 
+impl<T> TryFromU64 for T
+where
+    T: ActiveEnum,
+{
+    fn try_from_u64(_: u64) -> Result<Self, DbErr> {
+        Err(DbErr::ConvertFromU64(
+            "Fail to construct ActiveEnum from a u64, if your primary key consist of a ActiveEnum field, its auto increment should be set to false."
+        ))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate as sea_orm;

tests/common/features/mod.rs

@@ -16,6 +16,7 @@ pub mod satellite;
 pub mod schema;
 pub mod sea_orm_active_enums;
 pub mod self_join;
+pub mod teas;
 pub mod transaction_log;
 pub mod uuid_fmt;
 
@@ -36,5 +37,6 @@ pub use satellite::Entity as Satellite;
 pub use schema::*;
 pub use sea_orm_active_enums::*;
 pub use self_join::Entity as SelfJoin;
+pub use teas::Entity as Teas;
 pub use transaction_log::Entity as TransactionLog;
 pub use uuid_fmt::Entity as UuidFmt;

tests/common/features/schema.rs

@@ -44,6 +44,7 @@ pub async fn create_tables(db: &DatabaseConnection) -> Result<(), DbErr> {
     create_pi_table(db).await?;
     create_uuid_fmt_table(db).await?;
     create_edit_log_table(db).await?;
+    create_teas_table(db).await?;
 
     if DbBackend::Postgres == db_backend {
         create_collection_table(db).await?;
@@ -498,3 +499,19 @@ pub async fn create_edit_log_table(db: &DbConn) -> Result<ExecResult, DbErr> {
 
     create_table(db, &stmt, EditLog).await
 }
+
+pub async fn create_teas_table(db: &DbConn) -> Result<ExecResult, DbErr> {
+    let create_table_stmt = sea_query::Table::create()
+        .table(teas::Entity.table_ref())
+        .col(
+            ColumnDef::new(teas::Column::Id)
+                .enumeration(TeaEnum, [TeaVariant::EverydayTea, TeaVariant::BreakfastTea])
+                .not_null()
+                .primary_key(),
+        )
+        .col(ColumnDef::new(teas::Column::Category).string_len(1))
+        .col(ColumnDef::new(teas::Column::Color).integer())
+        .to_owned();
+
+    create_table(db, &create_table_stmt, Teas).await
+}

tests/common/features/teas.rs

@@ -0,0 +1,16 @@
+use super::sea_orm_active_enums::*;
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+#[sea_orm(table_name = "teas")]
+pub struct Model {
+    #[sea_orm(primary_key, auto_increment = false)]
+    pub id: Tea,
+    pub category: Option<Category>,
+    pub color: Option<Color>,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}

tests/enum_primary_key_tests.rs

@@ -0,0 +1,126 @@
+pub mod common;
+
+pub use common::{features::*, setup::*, TestContext};
+use pretty_assertions::assert_eq;
+use sea_orm::{
+    entity::prelude::*,
+    entity::*,
+    sea_query::{BinOper, Expr},
+    ActiveEnum as ActiveEnumTrait, DatabaseConnection,
+};
+
+#[sea_orm_macros::test]
+#[cfg(any(
+    feature = "sqlx-mysql",
+    feature = "sqlx-sqlite",
+    feature = "sqlx-postgres"
+))]
+async fn main() -> Result<(), DbErr> {
+    let ctx = TestContext::new("enum_primary_key_tests").await;
+    create_tables(&ctx.db).await?;
+    insert_teas(&ctx.db).await?;
+    ctx.delete().await;
+
+    Ok(())
+}
+
+pub async fn insert_teas(db: &DatabaseConnection) -> Result<(), DbErr> {
+    use teas::*;
+
+    let model = Model {
+        id: Tea::EverydayTea,
+        category: None,
+        color: None,
+    };
+
+    assert_eq!(
+        model,
+        ActiveModel {
+            id: Set(Tea::EverydayTea),
+            category: Set(None),
+            color: Set(None),
+            ..Default::default()
+        }
+        .insert(db)
+        .await?
+    );
+    assert_eq!(model, Entity::find().one(db).await?.unwrap());
+    assert_eq!(
+        model,
+        Entity::find()
+            .filter(Column::Id.is_not_null())
+            .filter(Column::Category.is_null())
+            .filter(Column::Color.is_null())
+            .one(db)
+            .await?
+            .unwrap()
+    );
+
+    // UNIQUE constraint failed
+    assert!(ActiveModel {
+        id: Set(Tea::EverydayTea),
+        category: Set(Some(Category::Big)),
+        color: Set(Some(Color::Black)),
+        ..Default::default()
+    }
+    .insert(db)
+    .await
+    .is_err());
+
+    // UNIQUE constraint failed
+    assert!(Entity::insert(ActiveModel {
+        id: Set(Tea::EverydayTea),
+        category: Set(Some(Category::Big)),
+        color: Set(Some(Color::Black)),
+        ..Default::default()
+    })
+    .exec(db)
+    .await
+    .is_err());
+
+    let _ = ActiveModel {
+        category: Set(Some(Category::Big)),
+        color: Set(Some(Color::Black)),
+        ..model.into_active_model()
+    }
+    .save(db)
+    .await?;
+
+    let model = Entity::find().one(db).await?.unwrap();
+    assert_eq!(
+        model,
+        Model {
+            id: Tea::EverydayTea,
+            category: Some(Category::Big),
+            color: Some(Color::Black),
+        }
+    );
+    assert_eq!(
+        model,
+        Entity::find()
+            .filter(Column::Id.eq(Tea::EverydayTea))
+            .filter(Column::Category.eq(Category::Big))
+            .filter(Column::Color.eq(Color::Black))
+            .one(db)
+            .await?
+            .unwrap()
+    );
+    assert_eq!(
+        model,
+        Entity::find()
+            .filter(
+                Expr::col(Column::Id)
+                    .binary(BinOper::In, Expr::tuple([Tea::EverydayTea.as_enum()]))
+            )
+            .one(db)
+            .await?
+            .unwrap()
+    );
+
+    let res = model.delete(db).await?;
+
+    assert_eq!(res.rows_affected, 1);
+    assert_eq!(Entity::find().one(db).await?, None);
+
+    Ok(())
+}