use crate::{ util::unpack_table_ref, ActiveEnum, Column, ConjunctRelation, Entity, EntityWriter, Error, PrimaryKey, Relation, RelationType, }; use sea_query::{ColumnSpec, TableCreateStatement}; use std::collections::BTreeMap; #[derive(Clone, Debug)] pub struct EntityTransformer; impl EntityTransformer { pub fn transform(table_create_stmts: Vec) -> Result { let mut enums: BTreeMap = BTreeMap::new(); let mut inverse_relations: BTreeMap> = BTreeMap::new(); let mut conjunct_relations: BTreeMap> = BTreeMap::new(); let mut entities = BTreeMap::new(); for table_create in table_create_stmts.into_iter() { let table_name = match table_create.get_table_name() { Some(table_ref) => match table_ref { sea_query::TableRef::Table(t) | sea_query::TableRef::SchemaTable(_, t) | sea_query::TableRef::DatabaseSchemaTable(_, _, t) | sea_query::TableRef::TableAlias(t, _) | sea_query::TableRef::SchemaTableAlias(_, t, _) | sea_query::TableRef::DatabaseSchemaTableAlias(_, _, t, _) => t.to_string(), _ => unimplemented!(), }, None => { return Err(Error::TransformError( "Table name should not be empty".into(), )) } }; let mut primary_keys: Vec = Vec::new(); let columns: Vec = table_create .get_columns() .iter() .map(|col_def| { let primary_key = col_def .get_column_spec() .iter() .any(|spec| matches!(spec, ColumnSpec::PrimaryKey)); if primary_key { primary_keys.push(PrimaryKey { name: col_def.get_column_name(), }); } col_def.into() }) .map(|mut col: Column| { col.unique = table_create .get_indexes() .iter() .filter(|index| index.is_unique_key()) .map(|index| index.get_index_spec().get_column_names()) .filter(|col_names| col_names.len() == 1 && col_names[0] == col.name) .count() > 0; col }) .map(|col| { if let sea_query::ColumnType::Enum { name, variants } = &col.col_type { enums.insert( name.to_string(), ActiveEnum { enum_name: name.clone(), values: variants.clone(), }, ); } col }) .collect(); let mut ref_table_counts: BTreeMap = BTreeMap::new(); let relations: Vec = table_create .get_foreign_key_create_stmts() .iter() .map(|fk_create_stmt| fk_create_stmt.get_foreign_key()) .map(|tbl_fk| { let ref_tbl = unpack_table_ref(tbl_fk.get_ref_table().unwrap()); if let Some(count) = ref_table_counts.get_mut(&ref_tbl) { if *count == 0 { *count = 1; } *count += 1; } else { ref_table_counts.insert(ref_tbl, 0); }; tbl_fk.into() }) .collect::>() .into_iter() .rev() .map(|mut rel: Relation| { rel.self_referencing = rel.ref_table == table_name; if let Some(count) = ref_table_counts.get_mut(&rel.ref_table) { rel.num_suffix = *count; if *count > 0 { *count -= 1; } } rel }) .rev() .collect(); primary_keys.extend( table_create .get_indexes() .iter() .filter(|index| index.is_primary_key()) .flat_map(|index| { index .get_index_spec() .get_column_names() .into_iter() .map(|name| PrimaryKey { name }) .collect::>() }), ); let entity = Entity { table_name: table_name.clone(), columns, relations: relations.clone(), conjunct_relations: vec![], primary_keys, }; entities.insert(table_name.clone(), entity.clone()); for (i, mut rel) in relations.into_iter().enumerate() { // This will produce a duplicated relation if rel.self_referencing { continue; } // This will cause compile error on the many side, // got relation variant but without Related implemented if rel.num_suffix > 0 { continue; } let is_conjunct_relation = entity.relations.len() == 2 && 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 column in rel.columns.iter() { if !entity .columns .iter() .filter(|col| col.unique) .any(|col| col.name.as_str() == column) { 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]); } } } } } for (tbl_name, relations) in inverse_relations.into_iter() { if let Some(entity) = entities.get_mut(&tbl_name) { for relation in relations.into_iter() { let duplicate_relation = entity .relations .iter() .any(|rel| rel.ref_table == relation.ref_table); if !duplicate_relation { entity.relations.push(relation); } } } } for (tbl_name, mut conjunct_relations) in conjunct_relations.into_iter() { if let Some(entity) = entities.get_mut(&tbl_name) { for relation in entity.relations.iter_mut() { // Skip `impl Related ... { fn to() ... }` implementation block, // if the same related entity is being referenced by a conjunct relation if conjunct_relations .iter() .any(|conjunct_relation| conjunct_relation.to == relation.ref_table) { relation.impl_related = false; } } entity.conjunct_relations.append(&mut conjunct_relations); } } Ok(EntityWriter { entities: entities .into_values() .map(|mut v| { v.relations.sort_by(|a, b| a.ref_table.cmp(&b.ref_table)); v }) .collect(), enums, }) } }