From 892c0fe57dca0783be6050eee455b62b3b460bd2 Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Tue, 2 Nov 2021 14:56:29 +0800
Subject: [PATCH] Support `Vec<u8>` primary key

---
 src/executor/query.rs                     |  2 +
 tests/byte_primary_key_tests.rs           | 74 +++++++++++++++++++++++
 tests/common/features/byte_primary_key.rs | 14 +++++
 tests/common/features/mod.rs              |  2 +
 tests/common/features/schema.rs           | 20 ++++++
 5 files changed, 112 insertions(+)
 create mode 100644 tests/byte_primary_key_tests.rs
 create mode 100644 tests/common/features/byte_primary_key.rs

diff --git a/src/executor/query.rs b/src/executor/query.rs
index 96455945..3c1a9184 100644
--- a/src/executor/query.rs
+++ b/src/executor/query.rs
@@ -539,5 +539,7 @@ macro_rules! try_from_u64_string {
 
 try_from_u64_string!(String);
 
+try_from_u64_err!(Vec<u8>);
+
 #[cfg(feature = "with-uuid")]
 try_from_u64_err!(uuid::Uuid);
diff --git a/tests/byte_primary_key_tests.rs b/tests/byte_primary_key_tests.rs
new file mode 100644
index 00000000..cb3ab775
--- /dev/null
+++ b/tests/byte_primary_key_tests.rs
@@ -0,0 +1,74 @@
+pub mod common;
+
+pub use common::{features::*, setup::*, TestContext};
+use pretty_assertions::assert_eq;
+use sea_orm::{entity::prelude::*, entity::*, 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("byte_primary_key_tests").await;
+    create_tables(&ctx.db).await?;
+    create_and_update(&ctx.db).await?;
+    ctx.delete().await;
+
+    Ok(())
+}
+
+pub async fn create_and_update(db: &DatabaseConnection) -> Result<(), DbErr> {
+    use common::features::byte_primary_key::*;
+
+    let model = Model {
+        id: vec![1],
+        value: "First Row".to_owned(),
+    };
+
+    let res = Entity::insert(model.clone().into_active_model())
+        .exec(db)
+        .await?;
+
+    assert_eq!(Entity::find().one(db).await?, Some(model.clone()));
+
+    assert_eq!(res.last_insert_id, model.id);
+
+    let updated_active_model = ActiveModel {
+        value: Set("First Row (Updated)".to_owned()),
+        ..model.clone().into_active_model()
+    };
+
+    let update_res = Entity::update(updated_active_model.clone())
+        .filter(Column::Id.eq(vec![1, 4]))
+        .exec(db)
+        .await;
+
+    assert_eq!(
+        update_res,
+        Err(DbErr::RecordNotFound(
+            "None of the database rows are affected".to_owned()
+        ))
+    );
+
+    let update_res = Entity::update(updated_active_model.clone())
+        .filter(Column::Id.eq(vec![1]))
+        .exec(db)
+        .await?;
+
+    assert_eq!(update_res, updated_active_model);
+
+    assert_eq!(
+        Entity::find()
+            .filter(Column::Id.eq(vec![1]))
+            .one(db)
+            .await?,
+        Some(Model {
+            id: vec![1],
+            value: "First Row (Updated)".to_owned(),
+        })
+    );
+
+    Ok(())
+}
diff --git a/tests/common/features/byte_primary_key.rs b/tests/common/features/byte_primary_key.rs
new file mode 100644
index 00000000..7578e3c8
--- /dev/null
+++ b/tests/common/features/byte_primary_key.rs
@@ -0,0 +1,14 @@
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
+#[sea_orm(table_name = "byte_primary_key")]
+pub struct Model {
+    #[sea_orm(primary_key, auto_increment = false)]
+    pub id: Vec<u8>,
+    pub value: String,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}
diff --git a/tests/common/features/mod.rs b/tests/common/features/mod.rs
index 1e709edf..186256ae 100644
--- a/tests/common/features/mod.rs
+++ b/tests/common/features/mod.rs
@@ -1,10 +1,12 @@
 pub mod applog;
+pub mod byte_primary_key;
 pub mod metadata;
 pub mod repository;
 pub mod schema;
 pub mod self_join;
 
 pub use applog::Entity as Applog;
+pub use byte_primary_key::Entity as BytePrimaryKey;
 pub use metadata::Entity as Metadata;
 pub use repository::Entity as Repository;
 pub use schema::*;
diff --git a/tests/common/features/schema.rs b/tests/common/features/schema.rs
index 7e6f6424..b6cecab0 100644
--- a/tests/common/features/schema.rs
+++ b/tests/common/features/schema.rs
@@ -10,6 +10,7 @@ pub async fn create_tables(db: &DatabaseConnection) -> Result<(), DbErr> {
     create_metadata_table(db).await?;
     create_repository_table(db).await?;
     create_self_join_table(db).await?;
+    create_byte_primary_key_table(db).await?;
 
     Ok(())
 }
@@ -100,3 +101,22 @@ pub async fn create_self_join_table(db: &DbConn) -> Result<ExecResult, DbErr> {
 
     create_table(db, &stmt, SelfJoin).await
 }
+
+pub async fn create_byte_primary_key_table(db: &DbConn) -> Result<ExecResult, DbErr> {
+    let stmt = sea_query::Table::create()
+        .table(byte_primary_key::Entity)
+        .col(
+            ColumnDef::new(byte_primary_key::Column::Id)
+                .binary()
+                .not_null()
+                .primary_key(),
+        )
+        .col(
+            ColumnDef::new(byte_primary_key::Column::Value)
+                .string()
+                .not_null(),
+        )
+        .to_owned();
+
+    create_table(db, &stmt, BytePrimaryKey).await
+}