From 484da8f6b61abf124be8ace694be95e731664f01 Mon Sep 17 00:00:00 2001 From: Billy Chan Date: Sat, 28 Aug 2021 20:00:59 +0800 Subject: [PATCH] WIP --- examples/async-std/src/example_cake.rs | 2 ++ .../async-std/src/example_cake_filling.rs | 2 ++ examples/async-std/src/example_filling.rs | 2 ++ examples/async-std/src/example_fruit.rs | 2 ++ examples/tokio/src/cake.rs | 2 ++ sea-orm-codegen/src/entity/base_entity.rs | 24 ++++++++++++++++++- sea-orm-codegen/src/entity/writer.rs | 3 +++ sea-orm-codegen/tests/entity/cake.rs | 2 ++ sea-orm-codegen/tests/entity/cake_filling.rs | 2 ++ sea-orm-codegen/tests/entity/filling.rs | 2 ++ sea-orm-codegen/tests/entity/fruit.rs | 2 ++ sea-orm-codegen/tests/entity/vendor.rs | 2 ++ src/entity/column.rs | 2 +- src/entity/primary_key.rs | 2 ++ src/tests_cfg/cake.rs | 2 ++ src/tests_cfg/cake_filling.rs | 2 ++ src/tests_cfg/cake_filling_price.rs | 2 ++ src/tests_cfg/filling.rs | 2 ++ src/tests_cfg/fruit.rs | 2 ++ tests/common/bakery_chain/baker.rs | 2 ++ tests/common/bakery_chain/bakery.rs | 2 ++ tests/common/bakery_chain/cake.rs | 2 ++ tests/common/bakery_chain/cakes_bakers.rs | 2 ++ tests/common/bakery_chain/customer.rs | 2 ++ tests/common/bakery_chain/lineitem.rs | 2 ++ tests/common/bakery_chain/order.rs | 2 ++ 26 files changed, 73 insertions(+), 2 deletions(-) diff --git a/examples/async-std/src/example_cake.rs b/examples/async-std/src/example_cake.rs index 475315e8..114347ee 100644 --- a/examples/async-std/src/example_cake.rs +++ b/examples/async-std/src/example_cake.rs @@ -27,6 +27,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/examples/async-std/src/example_cake_filling.rs b/examples/async-std/src/example_cake_filling.rs index 19de83e4..c24e7296 100644 --- a/examples/async-std/src/example_cake_filling.rs +++ b/examples/async-std/src/example_cake_filling.rs @@ -28,6 +28,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { false } diff --git a/examples/async-std/src/example_filling.rs b/examples/async-std/src/example_filling.rs index 925b92fc..2a39a7de 100644 --- a/examples/async-std/src/example_filling.rs +++ b/examples/async-std/src/example_filling.rs @@ -27,6 +27,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/examples/async-std/src/example_fruit.rs b/examples/async-std/src/example_fruit.rs index b875da24..8802e707 100644 --- a/examples/async-std/src/example_fruit.rs +++ b/examples/async-std/src/example_fruit.rs @@ -29,6 +29,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/examples/tokio/src/cake.rs b/examples/tokio/src/cake.rs index 0b1a4439..21b83331 100644 --- a/examples/tokio/src/cake.rs +++ b/examples/tokio/src/cake.rs @@ -27,6 +27,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/sea-orm-codegen/src/entity/base_entity.rs b/sea-orm-codegen/src/entity/base_entity.rs index f1e03b21..2cc6dbb4 100644 --- a/sea-orm-codegen/src/entity/base_entity.rs +++ b/sea-orm-codegen/src/entity/base_entity.rs @@ -117,6 +117,18 @@ impl Entity { format_ident!("{}", auto_increment) } + pub fn get_primary_key_rs_type(&self) -> TokenStream { + if let Some(primary_key) = self.primary_keys.first() { + self.columns + .iter() + .find(|col| col.name.eq(&primary_key.name)) + .unwrap() + .get_rs_type() + } else { + TokenStream::new() + } + } + pub fn get_conjunct_relations_via_snake_case(&self) -> Vec { self.conjunct_relations .iter() @@ -151,7 +163,7 @@ mod tests { columns: vec![ Column { name: "id".to_owned(), - col_type: ColumnType::String(None), + col_type: ColumnType::Integer(None), auto_increment: false, not_null: false, unique: false, @@ -373,6 +385,16 @@ mod tests { ); } + #[test] + fn test_get_primary_key_rs_type() { + let entity = setup(); + + assert_eq!( + entity.get_primary_key_rs_type().to_string(), + entity.columns[0].get_rs_type().to_string() + ); + } + #[test] fn test_get_conjunct_relations_via_snake_case() { let entity = setup(); diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 207864ba..7accac2d 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/sea-orm-codegen/src/entity/writer.rs @@ -173,8 +173,11 @@ impl EntityWriter { pub fn gen_impl_primary_key(entity: &Entity) -> TokenStream { let primary_key_auto_increment = entity.get_primary_key_auto_increment(); + let value_type = entity.get_primary_key_rs_type(); quote! { impl PrimaryKeyTrait for PrimaryKey { + type ValueType = #value_type; + fn auto_increment() -> bool { #primary_key_auto_increment } diff --git a/sea-orm-codegen/tests/entity/cake.rs b/sea-orm-codegen/tests/entity/cake.rs index 29f55ac6..55fa279f 100644 --- a/sea-orm-codegen/tests/entity/cake.rs +++ b/sea-orm-codegen/tests/entity/cake.rs @@ -29,6 +29,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/sea-orm-codegen/tests/entity/cake_filling.rs b/sea-orm-codegen/tests/entity/cake_filling.rs index d0f00560..227740c1 100644 --- a/sea-orm-codegen/tests/entity/cake_filling.rs +++ b/sea-orm-codegen/tests/entity/cake_filling.rs @@ -30,6 +30,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { false } diff --git a/sea-orm-codegen/tests/entity/filling.rs b/sea-orm-codegen/tests/entity/filling.rs index bedb1ab4..73d58152 100644 --- a/sea-orm-codegen/tests/entity/filling.rs +++ b/sea-orm-codegen/tests/entity/filling.rs @@ -29,6 +29,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/sea-orm-codegen/tests/entity/fruit.rs b/sea-orm-codegen/tests/entity/fruit.rs index 72c37c1b..12919021 100644 --- a/sea-orm-codegen/tests/entity/fruit.rs +++ b/sea-orm-codegen/tests/entity/fruit.rs @@ -31,6 +31,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/sea-orm-codegen/tests/entity/vendor.rs b/sea-orm-codegen/tests/entity/vendor.rs index 320400f4..ffc211a8 100644 --- a/sea-orm-codegen/tests/entity/vendor.rs +++ b/sea-orm-codegen/tests/entity/vendor.rs @@ -31,6 +31,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/src/entity/column.rs b/src/entity/column.rs index 16546057..611950f5 100644 --- a/src/entity/column.rs +++ b/src/entity/column.rs @@ -1,6 +1,6 @@ -use std::str::FromStr; use crate::{EntityName, IdenStatic, Iterable}; use sea_query::{DynIden, Expr, SeaRc, SelectStatement, SimpleExpr, Value}; +use std::str::FromStr; #[derive(Debug, Clone)] pub struct ColumnDef { diff --git a/src/entity/primary_key.rs b/src/entity/primary_key.rs index 81a28915..a36b04e0 100644 --- a/src/entity/primary_key.rs +++ b/src/entity/primary_key.rs @@ -2,6 +2,8 @@ use super::{ColumnTrait, IdenStatic, Iterable}; //LINT: composite primary key cannot auto increment pub trait PrimaryKeyTrait: IdenStatic + Iterable { + type ValueType: Sized; + fn auto_increment() -> bool; } diff --git a/src/tests_cfg/cake.rs b/src/tests_cfg/cake.rs index f8a35d6c..b4306313 100644 --- a/src/tests_cfg/cake.rs +++ b/src/tests_cfg/cake.rs @@ -28,6 +28,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/src/tests_cfg/cake_filling.rs b/src/tests_cfg/cake_filling.rs index b1151ee4..4336d2fa 100644 --- a/src/tests_cfg/cake_filling.rs +++ b/src/tests_cfg/cake_filling.rs @@ -29,6 +29,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { false } diff --git a/src/tests_cfg/cake_filling_price.rs b/src/tests_cfg/cake_filling_price.rs index c0bcbea1..bd2bc8eb 100644 --- a/src/tests_cfg/cake_filling_price.rs +++ b/src/tests_cfg/cake_filling_price.rs @@ -35,6 +35,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { false } diff --git a/src/tests_cfg/filling.rs b/src/tests_cfg/filling.rs index b439af7b..5e691980 100644 --- a/src/tests_cfg/filling.rs +++ b/src/tests_cfg/filling.rs @@ -41,6 +41,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/src/tests_cfg/fruit.rs b/src/tests_cfg/fruit.rs index 0511ae58..b0b8bc79 100644 --- a/src/tests_cfg/fruit.rs +++ b/src/tests_cfg/fruit.rs @@ -30,6 +30,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/tests/common/bakery_chain/baker.rs b/tests/common/bakery_chain/baker.rs index 0c63e721..071a1fbc 100644 --- a/tests/common/bakery_chain/baker.rs +++ b/tests/common/bakery_chain/baker.rs @@ -31,6 +31,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/tests/common/bakery_chain/bakery.rs b/tests/common/bakery_chain/bakery.rs index 61803329..2baa7b6b 100644 --- a/tests/common/bakery_chain/bakery.rs +++ b/tests/common/bakery_chain/bakery.rs @@ -29,6 +29,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/tests/common/bakery_chain/cake.rs b/tests/common/bakery_chain/cake.rs index 72e649ce..f699773c 100644 --- a/tests/common/bakery_chain/cake.rs +++ b/tests/common/bakery_chain/cake.rs @@ -35,6 +35,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/tests/common/bakery_chain/cakes_bakers.rs b/tests/common/bakery_chain/cakes_bakers.rs index 8106bbdf..853cf6da 100644 --- a/tests/common/bakery_chain/cakes_bakers.rs +++ b/tests/common/bakery_chain/cakes_bakers.rs @@ -28,6 +28,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { false } diff --git a/tests/common/bakery_chain/customer.rs b/tests/common/bakery_chain/customer.rs index ce4319ff..51668131 100644 --- a/tests/common/bakery_chain/customer.rs +++ b/tests/common/bakery_chain/customer.rs @@ -29,6 +29,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/tests/common/bakery_chain/lineitem.rs b/tests/common/bakery_chain/lineitem.rs index 45a6037f..1d338ed1 100644 --- a/tests/common/bakery_chain/lineitem.rs +++ b/tests/common/bakery_chain/lineitem.rs @@ -33,6 +33,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true } diff --git a/tests/common/bakery_chain/order.rs b/tests/common/bakery_chain/order.rs index 82b02dee..f2d48971 100644 --- a/tests/common/bakery_chain/order.rs +++ b/tests/common/bakery_chain/order.rs @@ -33,6 +33,8 @@ pub enum PrimaryKey { } impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + fn auto_increment() -> bool { true }