Merge branch 'codegen-amendment'

examples/cli/src/entity/fruit.rs

@@ -48,7 +48,7 @@ impl ColumnTrait for Column {
         match self {
             Self::Id => ColumnType::Integer.def(),
             Self::Name => ColumnType::String(Some(255u32)).def(),
-            Self::CakeId => ColumnType::Integer.def(),
+            Self::CakeId => ColumnType::Integer.def().null(),
         }
     }
 }

sea-orm-codegen/src/entity/column.rs

@@ -9,6 +9,7 @@ pub struct Column {
     pub(crate) col_type: ColumnType,
     pub(crate) auto_increment: bool,
     pub(crate) not_null: bool,
+    pub(crate) unique: bool,
 }
 
 impl Column {
@@ -50,7 +51,7 @@ impl Column {
     }
 
     pub fn get_def(&self) -> TokenStream {
-        match &self.col_type {
+        let mut col_def = match &self.col_type {
             ColumnType::Char(s) => match s {
                 Some(s) => quote! { ColumnType::Char(Some(#s)).def() },
                 None => quote! { ColumnType::Char(None).def() },
@@ -86,7 +87,18 @@ impl Column {
                 let s = s.to_string();
                 quote! { ColumnType::Custom(#s.to_owned()).def() }
             }
+        };
+        if !self.not_null {
+            col_def.extend(quote! {
+                .null()
+            });
         }
+        if self.unique {
+            col_def.extend(quote! {
+                .unique()
+            });
+        }
+        col_def
     }
 }
 
@@ -115,11 +127,21 @@ impl From<&ColumnDef> for Column {
             })
             .collect();
         let not_null = !not_nulls.is_empty();
+        let uniques: Vec<bool> = col_def
+            .get_column_spec()
+            .iter()
+            .filter_map(|spec| match spec {
+                ColumnSpec::UniqueKey => Some(true),
+                _ => None,
+            })
+            .collect();
+        let unique = !uniques.is_empty();
         Self {
             name,
             col_type,
             auto_increment,
             not_null,
+            unique,
         }
     }
 }

sea-orm-codegen/src/entity/entity.rs

@@ -86,6 +86,10 @@ impl Entity {
             .collect()
     }
 
+    pub fn get_relation_defs(&self) -> Vec<TokenStream> {
+        self.relations.iter().map(|rel| rel.get_def()).collect()
+    }
+
     pub fn get_relation_rel_types(&self) -> Vec<Ident> {
         self.relations
             .iter()

sea-orm-codegen/src/entity/relation.rs

@@ -1,9 +1,15 @@
 use heck::{CamelCase, SnakeCase};
-use proc_macro2::Ident;
+use proc_macro2::{Ident, TokenStream};
-use quote::format_ident;
+use quote::{format_ident, quote};
-use sea_orm::RelationType;
 use sea_query::TableForeignKey;
 
+#[derive(Clone, Debug)]
+pub enum RelationType {
+    HasOne,
+    HasMany,
+    BelongsTo,
+}
+
 #[derive(Clone, Debug)]
 pub struct Relation {
     pub(crate) ref_table: String,
@@ -21,10 +27,33 @@ impl Relation {
         format_ident!("{}", self.ref_table.to_camel_case())
     }
 
+    pub fn get_def(&self) -> TokenStream {
+        let rel_type = self.get_rel_type();
+        let ref_table_snake_case = self.get_ref_table_snake_case();
+        match self.rel_type {
+            RelationType::HasOne | RelationType::HasMany => {
+                quote! {
+                    Entity::#rel_type(super::#ref_table_snake_case::Entity).into()
+                }
+            }
+            RelationType::BelongsTo => {
+                let column_camel_case = self.get_column_camel_case();
+                let ref_column_camel_case = self.get_ref_column_camel_case();
+                quote! {
+                    Entity::#rel_type(super::#ref_table_snake_case::Entity)
+                        .from(Column::#column_camel_case)
+                        .to(super::#ref_table_snake_case::Column::#ref_column_camel_case)
+                        .into()
+                }
+            }
+        }
+    }
+
     pub fn get_rel_type(&self) -> Ident {
         match self.rel_type {
             RelationType::HasOne => format_ident!("has_one"),
             RelationType::HasMany => format_ident!("has_many"),
+            RelationType::BelongsTo => format_ident!("belongs_to"),
         }
     }
 
@@ -49,7 +78,7 @@ impl From<&TableForeignKey> for Relation {
         };
         let columns = tbl_fk.get_columns();
         let ref_columns = tbl_fk.get_ref_columns();
-        let rel_type = RelationType::HasOne;
+        let rel_type = RelationType::BelongsTo;
         Self {
             ref_table,
             columns,

sea-orm-codegen/src/entity/transformer.rs

@@ -1,5 +1,4 @@
-use crate::{Entity, EntityWriter, Error, PrimaryKey, Relation};
+use crate::{Column, Entity, EntityWriter, Error, PrimaryKey, Relation, RelationType};
-use sea_orm::RelationType;
 use sea_query::TableStatement;
 use sea_schema::mysql::def::Schema;
 use std::{collections::HashMap, mem::swap};
@@ -31,11 +30,16 @@ impl EntityTransformer {
                     ))
                 }
             };
-            let columns = table_create
+            let columns: Vec<Column> = table_create
                 .get_columns()
                 .iter()
                 .map(|col_def| col_def.into())
                 .collect();
+            let unique_columns: Vec<String> = columns
+                .iter()
+                .filter(|col| col.unique)
+                .map(|col| col.name.clone())
+                .collect();
             let relations = table_create
                 .get_foreign_key_create_stmts()
                 .iter()
@@ -64,9 +68,22 @@ impl EntityTransformer {
             entities.push(entity);
             for mut rel in relations.into_iter() {
                 let ref_table = rel.ref_table;
-                swap(&mut rel.columns, &mut rel.ref_columns);
+                let mut unique = true;
-                rel.rel_type = RelationType::HasMany;
+                for col in rel.columns.iter() {
+                    if !unique_columns.contains(col) {
+                        unique = false;
+                        break;
+                    }
+                }
+                let rel_type = if unique {
+                    RelationType::HasOne
+                } else {
+                    RelationType::HasMany
+                };
+                rel.rel_type = rel_type;
                 rel.ref_table = table_name.clone();
+                rel.columns = Vec::new();
+                rel.ref_columns = Vec::new();
                 if let Some(vec) = inverse_relations.get_mut(&ref_table) {
                     vec.push(rel);
                 } else {

sea-orm-codegen/src/entity/writer.rs

@@ -217,20 +217,14 @@ impl EntityWriter {
 
     pub fn gen_impl_relation_trait(entity: &Entity) -> TokenStream {
         let relation_ref_tables_camel_case = entity.get_relation_ref_tables_camel_case();
-        let relation_rel_types = entity.get_relation_rel_types();
+        let relation_defs = entity.get_relation_defs();
-        let relation_ref_tables_snake_case = entity.get_relation_ref_tables_snake_case();
-        let relation_columns_camel_case = entity.get_relation_columns_camel_case();
-        let relation_ref_columns_camel_case = entity.get_relation_ref_columns_camel_case();
         let quoted = if relation_ref_tables_camel_case.is_empty() {
             quote! {
                 _ => panic!("No RelationDef"),
             }
         } else {
             quote! {
-                #(Self::#relation_ref_tables_camel_case => Entity::#relation_rel_types(super::#relation_ref_tables_snake_case::Entity)
+                #(Self::#relation_ref_tables_camel_case => #relation_defs),*
-                    .from(Column::#relation_columns_camel_case)
-                    .to(super::#relation_ref_tables_snake_case::Column::#relation_ref_columns_camel_case)
-                    .into()),*
             }
         };
         quote! {