Handle conjunction table

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

@@ -1,4 +1,4 @@
-use crate::{Column, PrimaryKey, Relation};
+use crate::{Column, ConjunctRelation, PrimaryKey, Relation};
 use heck::{CamelCase, SnakeCase};
 use proc_macro2::{Ident, TokenStream};
 use quote::format_ident;
@@ -8,6 +8,7 @@ pub struct Entity {
     pub(crate) table_name: String,
     pub(crate) columns: Vec<Column>,
     pub(crate) relations: Vec<Relation>,
+    pub(crate) conjunct_relations: Vec<ConjunctRelation>,
     pub(crate) primary_keys: Vec<PrimaryKey>,
 }
 
@@ -115,6 +116,27 @@ impl Entity {
         let auto_increment = self.columns.iter().any(|col| col.auto_increment);
         format_ident!("{}", auto_increment)
     }
+
+    pub fn get_conjunct_relations_via_snake_case(&self) -> Vec<Ident> {
+        self.conjunct_relations
+            .iter()
+            .map(|con_rel| con_rel.get_via_snake_case())
+            .collect()
+    }
+
+    pub fn get_conjunct_relations_to_snake_case(&self) -> Vec<Ident> {
+        self.conjunct_relations
+            .iter()
+            .map(|con_rel| con_rel.get_to_snake_case())
+            .collect()
+    }
+
+    pub fn get_conjunct_relations_to_camel_case(&self) -> Vec<Ident> {
+        self.conjunct_relations
+            .iter()
+            .map(|con_rel| con_rel.get_to_camel_case())
+            .collect()
+    }
 }
 
 #[cfg(test)]
@@ -156,6 +178,7 @@ mod tests {
                     rel_type: RelationType::HasOne,
                 },
             ],
+            conjunct_relations: vec![],
             primary_keys: vec![PrimaryKey {
                 name: "id".to_owned(),
             }],
@@ -349,4 +372,43 @@ mod tests {
             format_ident!("{}", true)
         );
     }
+
+    #[test]
+    fn test_get_conjunct_relations_via_snake_case() {
+        let entity = setup();
+
+        for (i, elem) in entity
+            .get_conjunct_relations_via_snake_case()
+            .into_iter()
+            .enumerate()
+        {
+            assert_eq!(elem, entity.conjunct_relations[i].get_via_snake_case());
+        }
+    }
+
+    #[test]
+    fn test_get_conjunct_relations_to_snake_case() {
+        let entity = setup();
+
+        for (i, elem) in entity
+            .get_conjunct_relations_to_snake_case()
+            .into_iter()
+            .enumerate()
+        {
+            assert_eq!(elem, entity.conjunct_relations[i].get_to_snake_case());
+        }
+    }
+
+    #[test]
+    fn test_get_conjunct_relations_to_camel_case() {
+        let entity = setup();
+
+        for (i, elem) in entity
+            .get_conjunct_relations_to_camel_case()
+            .into_iter()
+            .enumerate()
+        {
+            assert_eq!(elem, entity.conjunct_relations[i].get_to_camel_case());
+        }
+    }
 }

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

@@ -0,0 +1,68 @@
+use heck::{CamelCase, SnakeCase};
+use proc_macro2::Ident;
+use quote::format_ident;
+
+#[derive(Clone, Debug)]
+pub struct ConjunctRelation {
+    pub(crate) via: String,
+    pub(crate) to: String,
+}
+
+impl ConjunctRelation {
+    pub fn get_via_snake_case(&self) -> Ident {
+        format_ident!("{}", self.via.to_snake_case())
+    }
+
+    pub fn get_to_snake_case(&self) -> Ident {
+        format_ident!("{}", self.to.to_snake_case())
+    }
+
+    pub fn get_to_camel_case(&self) -> Ident {
+        format_ident!("{}", self.to.to_camel_case())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::ConjunctRelation;
+
+    fn setup() -> Vec<ConjunctRelation> {
+        vec![
+            ConjunctRelation {
+                via: "cake_filling".to_owned(),
+                to: "cake".to_owned(),
+            },
+            ConjunctRelation {
+                via: "cake_filling".to_owned(),
+                to: "filling".to_owned(),
+            },
+        ]
+    }
+
+    #[test]
+    fn test_get_via_snake_case() {
+        let conjunct_relations = setup();
+        let via_vec = vec!["cake_filling", "cake_filling"];
+        for (con_rel, via) in conjunct_relations.into_iter().zip(via_vec) {
+            assert_eq!(con_rel.get_via_snake_case(), via);
+        }
+    }
+
+    #[test]
+    fn test_get_to_snake_case() {
+        let conjunct_relations = setup();
+        let to_vec = vec!["cake", "filling"];
+        for (con_rel, to) in conjunct_relations.into_iter().zip(to_vec) {
+            assert_eq!(con_rel.get_to_snake_case(), to);
+        }
+    }
+
+    #[test]
+    fn test_get_to_camel_case() {
+        let conjunct_relations = setup();
+        let to_vec = vec!["Cake", "Filling"];
+        for (con_rel, to) in conjunct_relations.into_iter().zip(to_vec) {
+            assert_eq!(con_rel.get_to_camel_case(), to);
+        }
+    }
+}

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

@@ -1,5 +1,6 @@
 mod base_entity;
 mod column;
+mod conjunct_relation;
 mod generator;
 mod primary_key;
 mod relation;
@@ -8,6 +9,7 @@ mod writer;
 
 pub use base_entity::*;
 pub use column::*;
+pub use conjunct_relation::*;
 pub use generator::*;
 pub use primary_key::*;
 pub use relation::*;

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

@@ -1,4 +1,6 @@
-use crate::{Column, Entity, EntityWriter, Error, PrimaryKey, Relation, RelationType};
+use crate::{
+    Column, ConjunctRelation, Entity, EntityWriter, Error, PrimaryKey, Relation, RelationType,
+};
 use sea_query::TableStatement;
 use sea_schema::mysql::def::Schema;
 use std::collections::HashMap;
@@ -11,7 +13,8 @@ pub struct EntityTransformer {
 impl EntityTransformer {
     pub fn transform(self) -> Result<EntityWriter, Error> {
         let mut inverse_relations: HashMap<String, Vec<Relation>> = HashMap::new();
-        let mut entities = Vec::new();
+        let mut conjunct_relations: HashMap<String, Vec<ConjunctRelation>> = HashMap::new();
+        let mut entities = HashMap::new();
         for table_ref in self.schema.tables.iter() {
             let table_stmt = table_ref.write();
             let table_create = match table_stmt {
@@ -63,42 +66,65 @@ impl EntityTransformer {
                 table_name: table_name.clone(),
                 columns,
                 relations: relations.clone().collect(),
+                conjunct_relations: vec![],
                 primary_keys,
             };
-            entities.push(entity);
-            for mut rel in relations.into_iter() {
-                let ref_table = rel.ref_table;
-                let mut unique = true;
-                for col in rel.columns.iter() {
-                    if !unique_columns.contains(col) {
-                        unique = false;
-                        break;
+            entities.insert(table_name.clone(), entity.clone());
+            for (i, mut rel) in relations.into_iter().enumerate() {
+                let is_conjunct_relation = entity.primary_keys.len() == entity.columns.len()
+                    && entity.primary_keys.len() == 2;
+                match is_conjunct_relation {
+                    true => {
+                        let another_rel = entity.relations.get((i == 0) as usize).unwrap();
+                        let conjunct_relation = ConjunctRelation {
+                            via: table_name.clone(),
+                            to: another_rel.ref_table.clone(),
+                        };
+                        if let Some(vec) = conjunct_relations.get_mut(&rel.ref_table) {
+                            vec.push(conjunct_relation);
+                        } else {
+                            conjunct_relations.insert(rel.ref_table, vec![conjunct_relation]);
+                        }
+                    }
+                    false => {
+                        let ref_table = rel.ref_table;
+                        let mut unique = true;
+                        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 {
+                            inverse_relations.insert(ref_table, vec![rel]);
+                        }
                     }
                 }
-                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 {
-                    inverse_relations.insert(ref_table, vec![rel]);
-                }
             }
         }
-        for (tbl_name, relations) in inverse_relations.iter() {
-            for ent in entities.iter_mut() {
-                if ent.table_name.eq(tbl_name) {
-                    ent.relations.append(relations.clone().as_mut());
-                }
+        for (tbl_name, mut relations) in inverse_relations.into_iter() {
+            if let Some(entity) = entities.get_mut(&tbl_name) {
+                entity.relations.append(&mut relations);
             }
         }
-        println!("{:#?}", entities);
-        Ok(EntityWriter { entities })
+        for (tbl_name, mut conjunct_relations) in conjunct_relations.into_iter() {
+            if let Some(entity) = entities.get_mut(&tbl_name) {
+                entity.conjunct_relations.append(&mut conjunct_relations);
+            }
+        }
+        Ok(EntityWriter {
+            entities: entities.into_iter().map(|(_, v)| v).collect(),
+        })
     }
 }

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

@@ -116,6 +116,7 @@ impl EntityWriter {
             Self::gen_impl_relation_trait(entity),
         ];
         code_blocks.extend(Self::gen_impl_related(entity));
+        code_blocks.extend(Self::gen_impl_conjunct_related(entity));
         code_blocks.extend(vec![Self::gen_impl_active_model_behavior()]);
         code_blocks
     }
@@ -239,7 +240,7 @@ impl EntityWriter {
         let camel = entity.get_relation_ref_tables_camel_case();
         let snake = entity.get_relation_ref_tables_snake_case();
         camel
-            .iter()
+            .into_iter()
             .zip(snake)
             .map(|(c, s)| {
                 quote! {
@@ -253,6 +254,31 @@ impl EntityWriter {
             .collect()
     }
 
+    pub fn gen_impl_conjunct_related(entity: &Entity) -> Vec<TokenStream> {
+        let table_name_camel_case = entity.get_table_name_camel_case_ident();
+        let via_snake_case = entity.get_conjunct_relations_via_snake_case();
+        let to_snake_case = entity.get_conjunct_relations_to_snake_case();
+        let to_camel_case = entity.get_conjunct_relations_to_camel_case();
+        via_snake_case
+            .into_iter()
+            .zip(to_snake_case)
+            .zip(to_camel_case)
+            .map(|((via_snake_case, to_snake_case), to_camel_case)| {
+                quote! {
+                    impl Related<super::#to_snake_case::Entity> for Entity {
+                        fn to() -> RelationDef {
+                            super::#via_snake_case::Relation::#to_camel_case.def()
+                        }
+
+                        fn via() -> Option<RelationDef> {
+                            Some(super::#via_snake_case::Relation::#table_name_camel_case.def().rev())
+                        }
+                    }
+                }
+            })
+            .collect()
+    }
+
     pub fn gen_impl_active_model_behavior() -> TokenStream {
         quote! {
             impl ActiveModelBehavior for ActiveModel {}
@@ -277,7 +303,9 @@ impl EntityWriter {
 
 #[cfg(test)]
 mod tests {
-    use crate::{Column, Entity, EntityWriter, PrimaryKey, Relation, RelationType};
+    use crate::{
+        Column, ConjunctRelation, Entity, EntityWriter, PrimaryKey, Relation, RelationType,
+    };
     use proc_macro2::TokenStream;
     use sea_query::ColumnType;
     use std::io::{self, BufRead, BufReader};
@@ -310,20 +338,16 @@ mod tests {
                         unique: false,
                     },
                 ],
-                relations: vec![
-                    Relation {
-                        ref_table: "cake_filling".to_owned(),
-                        columns: vec![],
-                        ref_columns: vec![],
-                        rel_type: RelationType::HasMany,
-                    },
-                    Relation {
-                        ref_table: "fruit".to_owned(),
-                        columns: vec![],
-                        ref_columns: vec![],
-                        rel_type: RelationType::HasMany,
-                    },
-                ],
+                relations: vec![Relation {
+                    ref_table: "fruit".to_owned(),
+                    columns: vec![],
+                    ref_columns: vec![],
+                    rel_type: RelationType::HasMany,
+                }],
+                conjunct_relations: vec![ConjunctRelation {
+                    via: "cake_filling".to_owned(),
+                    to: "filling".to_owned(),
+                }],
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
@@ -360,6 +384,7 @@ mod tests {
                         rel_type: RelationType::BelongsTo,
                     },
                 ],
+                conjunct_relations: vec![],
                 primary_keys: vec![
                     PrimaryKey {
                         name: "cake_id".to_owned(),
@@ -387,11 +412,10 @@ mod tests {
                         unique: false,
                     },
                 ],
-                relations: vec![Relation {
-                    ref_table: "cake_filling".to_owned(),
-                    columns: vec![],
-                    ref_columns: vec![],
-                    rel_type: RelationType::HasMany,
+                relations: vec![],
+                conjunct_relations: vec![ConjunctRelation {
+                    via: "cake_filling".to_owned(),
+                    to: "cake".to_owned(),
                 }],
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
@@ -436,6 +460,7 @@ mod tests {
                         rel_type: RelationType::HasMany,
                     },
                 ],
+                conjunct_relations: vec![],
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
@@ -471,6 +496,7 @@ mod tests {
                     ref_columns: vec!["id".to_owned()],
                     rel_type: RelationType::BelongsTo,
                 }],
+                conjunct_relations: vec![],
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],