sea_orm_codegen/entity/
writer.rs

1use crate::{util::escape_rust_keyword, ActiveEnum, Entity};
2use heck::ToUpperCamelCase;
3use proc_macro2::TokenStream;
4use quote::{format_ident, quote};
5use std::{collections::BTreeMap, str::FromStr};
6use syn::{punctuated::Punctuated, token::Comma};
7use tracing::info;
8
9#[derive(Clone, Debug)]
10pub struct EntityWriter {
11    pub(crate) entities: Vec<Entity>,
12    pub(crate) enums: BTreeMap<String, ActiveEnum>,
13}
14
15pub struct WriterOutput {
16    pub files: Vec<OutputFile>,
17}
18
19pub struct OutputFile {
20    pub name: String,
21    pub content: String,
22}
23
24#[derive(Default, PartialEq, Eq, Clone, Copy, Debug)]
25pub enum WithPrelude {
26    #[default]
27    All,
28    None,
29    AllAllowUnusedImports,
30}
31
32#[derive(PartialEq, Eq, Debug)]
33pub enum WithSerde {
34    None,
35    Serialize,
36    Deserialize,
37    Both,
38}
39
40#[derive(Debug)]
41pub enum DateTimeCrate {
42    Chrono,
43    Time,
44}
45
46#[derive(Debug)]
47pub struct EntityWriterContext {
48    pub(crate) expanded_format: bool,
49    pub(crate) with_prelude: WithPrelude,
50    pub(crate) with_serde: WithSerde,
51    pub(crate) with_copy_enums: bool,
52    pub(crate) date_time_crate: DateTimeCrate,
53    pub(crate) schema_name: Option<String>,
54    pub(crate) lib: bool,
55    pub(crate) serde_skip_hidden_column: bool,
56    pub(crate) serde_skip_deserializing_primary_key: bool,
57    pub(crate) model_extra_derives: TokenStream,
58    pub(crate) model_extra_attributes: TokenStream,
59    pub(crate) enum_extra_derives: TokenStream,
60    pub(crate) enum_extra_attributes: TokenStream,
61    pub(crate) seaography: bool,
62    pub(crate) impl_active_model_behavior: bool,
63}
64
65impl WithSerde {
66    pub fn extra_derive(&self) -> TokenStream {
67        let mut extra_derive = match self {
68            Self::None => {
69                quote! {}
70            }
71            Self::Serialize => {
72                quote! {
73                    Serialize
74                }
75            }
76            Self::Deserialize => {
77                quote! {
78                    Deserialize
79                }
80            }
81            Self::Both => {
82                quote! {
83                    Serialize, Deserialize
84                }
85            }
86        };
87        if !extra_derive.is_empty() {
88            extra_derive = quote! { , #extra_derive }
89        }
90        extra_derive
91    }
92}
93
94/// Converts *_extra_derives argument to token stream
95pub(crate) fn bonus_derive<T, I>(extra_derives: I) -> TokenStream
96where
97    T: Into<String>,
98    I: IntoIterator<Item = T>,
99{
100    extra_derives.into_iter().map(Into::<String>::into).fold(
101        TokenStream::default(),
102        |acc, derive| {
103            let tokens: TokenStream = derive.parse().unwrap();
104            quote! { #acc, #tokens }
105        },
106    )
107}
108
109/// convert *_extra_attributes argument to token stream
110pub(crate) fn bonus_attributes<T, I>(attributes: I) -> TokenStream
111where
112    T: Into<String>,
113    I: IntoIterator<Item = T>,
114{
115    attributes.into_iter().map(Into::<String>::into).fold(
116        TokenStream::default(),
117        |acc, attribute| {
118            let tokens: TokenStream = attribute.parse().unwrap();
119            quote! {
120                #acc
121                #[#tokens]
122            }
123        },
124    )
125}
126
127impl FromStr for WithPrelude {
128    type Err = crate::Error;
129
130    fn from_str(s: &str) -> Result<Self, Self::Err> {
131        Ok(match s {
132            "none" => Self::None,
133            "all-allow-unused-imports" => Self::AllAllowUnusedImports,
134            "all" => Self::All,
135            v => {
136                return Err(crate::Error::TransformError(format!(
137                    "Unsupported enum variant '{v}'"
138                )))
139            }
140        })
141    }
142}
143
144impl FromStr for WithSerde {
145    type Err = crate::Error;
146
147    fn from_str(s: &str) -> Result<Self, Self::Err> {
148        Ok(match s {
149            "none" => Self::None,
150            "serialize" => Self::Serialize,
151            "deserialize" => Self::Deserialize,
152            "both" => Self::Both,
153            v => {
154                return Err(crate::Error::TransformError(format!(
155                    "Unsupported enum variant '{v}'"
156                )))
157            }
158        })
159    }
160}
161
162impl EntityWriterContext {
163    #[allow(clippy::too_many_arguments)]
164    pub fn new(
165        expanded_format: bool,
166        with_prelude: WithPrelude,
167        with_serde: WithSerde,
168        with_copy_enums: bool,
169        date_time_crate: DateTimeCrate,
170        schema_name: Option<String>,
171        lib: bool,
172        serde_skip_deserializing_primary_key: bool,
173        serde_skip_hidden_column: bool,
174        model_extra_derives: Vec<String>,
175        model_extra_attributes: Vec<String>,
176        enum_extra_derives: Vec<String>,
177        enum_extra_attributes: Vec<String>,
178        seaography: bool,
179        impl_active_model_behavior: bool,
180    ) -> Self {
181        Self {
182            expanded_format,
183            with_prelude,
184            with_serde,
185            with_copy_enums,
186            date_time_crate,
187            schema_name,
188            lib,
189            serde_skip_deserializing_primary_key,
190            serde_skip_hidden_column,
191            model_extra_derives: bonus_derive(model_extra_derives),
192            model_extra_attributes: bonus_attributes(model_extra_attributes),
193            enum_extra_derives: bonus_derive(enum_extra_derives),
194            enum_extra_attributes: bonus_attributes(enum_extra_attributes),
195            seaography,
196            impl_active_model_behavior,
197        }
198    }
199}
200
201impl EntityWriter {
202    pub fn generate(self, context: &EntityWriterContext) -> WriterOutput {
203        let mut files = Vec::new();
204        files.extend(self.write_entities(context));
205        let with_prelude = context.with_prelude != WithPrelude::None;
206        files.push(self.write_index_file(context.lib, with_prelude, context.seaography));
207        if with_prelude {
208            files.push(self.write_prelude(context.with_prelude));
209        }
210        if !self.enums.is_empty() {
211            files.push(self.write_sea_orm_active_enums(
212                &context.with_serde,
213                context.with_copy_enums,
214                &context.enum_extra_derives,
215                &context.enum_extra_attributes,
216            ));
217        }
218        WriterOutput { files }
219    }
220
221    pub fn write_entities(&self, context: &EntityWriterContext) -> Vec<OutputFile> {
222        self.entities
223            .iter()
224            .map(|entity| {
225                let entity_file = format!("{}.rs", entity.get_table_name_snake_case());
226                let column_info = entity
227                    .columns
228                    .iter()
229                    .map(|column| column.get_info(&context.date_time_crate))
230                    .collect::<Vec<String>>();
231                // Serde must be enabled to use this
232                let serde_skip_deserializing_primary_key = context
233                    .serde_skip_deserializing_primary_key
234                    && matches!(context.with_serde, WithSerde::Both | WithSerde::Deserialize);
235                let serde_skip_hidden_column = context.serde_skip_hidden_column
236                    && matches!(
237                        context.with_serde,
238                        WithSerde::Both | WithSerde::Serialize | WithSerde::Deserialize
239                    );
240
241                info!("Generating {}", entity_file);
242                for info in column_info.iter() {
243                    info!("    > {}", info);
244                }
245
246                let mut lines = Vec::new();
247                Self::write_doc_comment(&mut lines);
248                let code_blocks = if context.expanded_format {
249                    Self::gen_expanded_code_blocks(
250                        entity,
251                        &context.with_serde,
252                        &context.date_time_crate,
253                        &context.schema_name,
254                        serde_skip_deserializing_primary_key,
255                        serde_skip_hidden_column,
256                        &context.model_extra_derives,
257                        &context.model_extra_attributes,
258                        context.seaography,
259                        context.impl_active_model_behavior,
260                    )
261                } else {
262                    Self::gen_compact_code_blocks(
263                        entity,
264                        &context.with_serde,
265                        &context.date_time_crate,
266                        &context.schema_name,
267                        serde_skip_deserializing_primary_key,
268                        serde_skip_hidden_column,
269                        &context.model_extra_derives,
270                        &context.model_extra_attributes,
271                        context.seaography,
272                        context.impl_active_model_behavior,
273                    )
274                };
275                Self::write(&mut lines, code_blocks);
276                OutputFile {
277                    name: entity_file,
278                    content: lines.join("\n\n"),
279                }
280            })
281            .collect()
282    }
283
284    pub fn write_index_file(&self, lib: bool, prelude: bool, seaography: bool) -> OutputFile {
285        let mut lines = Vec::new();
286        Self::write_doc_comment(&mut lines);
287        let code_blocks: Vec<TokenStream> = self.entities.iter().map(Self::gen_mod).collect();
288        if prelude {
289            Self::write(
290                &mut lines,
291                vec![quote! {
292                    pub mod prelude;
293                }],
294            );
295            lines.push("".to_owned());
296        }
297        Self::write(&mut lines, code_blocks);
298        if !self.enums.is_empty() {
299            Self::write(
300                &mut lines,
301                vec![quote! {
302                    pub mod sea_orm_active_enums;
303                }],
304            );
305        }
306
307        if seaography {
308            lines.push("".to_owned());
309            let ts = Self::gen_seaography_entity_mod(&self.entities, &self.enums);
310            Self::write(&mut lines, vec![ts]);
311        }
312
313        let file_name = match lib {
314            true => "lib.rs".to_owned(),
315            false => "mod.rs".to_owned(),
316        };
317
318        OutputFile {
319            name: file_name,
320            content: lines.join("\n"),
321        }
322    }
323
324    pub fn write_prelude(&self, with_prelude: WithPrelude) -> OutputFile {
325        let mut lines = Vec::new();
326        Self::write_doc_comment(&mut lines);
327        if with_prelude == WithPrelude::AllAllowUnusedImports {
328            Self::write_allow_unused_imports(&mut lines)
329        }
330        let code_blocks = self.entities.iter().map(Self::gen_prelude_use).collect();
331        Self::write(&mut lines, code_blocks);
332        OutputFile {
333            name: "prelude.rs".to_owned(),
334            content: lines.join("\n"),
335        }
336    }
337
338    pub fn write_sea_orm_active_enums(
339        &self,
340        with_serde: &WithSerde,
341        with_copy_enums: bool,
342        extra_derives: &TokenStream,
343        extra_attributes: &TokenStream,
344    ) -> OutputFile {
345        let mut lines = Vec::new();
346        Self::write_doc_comment(&mut lines);
347        Self::write(&mut lines, vec![Self::gen_import(with_serde)]);
348        lines.push("".to_owned());
349        let code_blocks = self
350            .enums
351            .values()
352            .map(|active_enum| {
353                active_enum.impl_active_enum(
354                    with_serde,
355                    with_copy_enums,
356                    extra_derives,
357                    extra_attributes,
358                )
359            })
360            .collect();
361        Self::write(&mut lines, code_blocks);
362        OutputFile {
363            name: "sea_orm_active_enums.rs".to_owned(),
364            content: lines.join("\n"),
365        }
366    }
367
368    pub fn write(lines: &mut Vec<String>, code_blocks: Vec<TokenStream>) {
369        lines.extend(
370            code_blocks
371                .into_iter()
372                .map(|code_block| code_block.to_string())
373                .collect::<Vec<_>>(),
374        );
375    }
376
377    pub fn write_doc_comment(lines: &mut Vec<String>) {
378        let ver = env!("CARGO_PKG_VERSION");
379        let comments = vec![format!(
380            "//! `SeaORM` Entity, @generated by sea-orm-codegen {ver}"
381        )];
382        lines.extend(comments);
383        lines.push("".to_owned());
384    }
385
386    pub fn write_allow_unused_imports(lines: &mut Vec<String>) {
387        lines.extend(vec!["#![allow(unused_imports)]".to_string()]);
388        lines.push("".to_owned());
389    }
390
391    #[allow(clippy::too_many_arguments)]
392    pub fn gen_expanded_code_blocks(
393        entity: &Entity,
394        with_serde: &WithSerde,
395        date_time_crate: &DateTimeCrate,
396        schema_name: &Option<String>,
397        serde_skip_deserializing_primary_key: bool,
398        serde_skip_hidden_column: bool,
399        model_extra_derives: &TokenStream,
400        model_extra_attributes: &TokenStream,
401        seaography: bool,
402        impl_active_model_behavior: bool,
403    ) -> Vec<TokenStream> {
404        let mut imports = Self::gen_import(with_serde);
405        imports.extend(Self::gen_import_active_enum(entity));
406        let mut code_blocks = vec![
407            imports,
408            Self::gen_entity_struct(),
409            Self::gen_impl_entity_name(entity, schema_name),
410            Self::gen_model_struct(
411                entity,
412                with_serde,
413                date_time_crate,
414                serde_skip_deserializing_primary_key,
415                serde_skip_hidden_column,
416                model_extra_derives,
417                model_extra_attributes,
418            ),
419            Self::gen_column_enum(entity),
420            Self::gen_primary_key_enum(entity),
421            Self::gen_impl_primary_key(entity, date_time_crate),
422            Self::gen_relation_enum(entity),
423            Self::gen_impl_column_trait(entity),
424            Self::gen_impl_relation_trait(entity),
425        ];
426        code_blocks.extend(Self::gen_impl_related(entity));
427        code_blocks.extend(Self::gen_impl_conjunct_related(entity));
428        if impl_active_model_behavior {
429            code_blocks.extend([Self::impl_active_model_behavior()]);
430        }
431        if seaography {
432            code_blocks.extend([Self::gen_related_entity(entity)]);
433        }
434        code_blocks
435    }
436
437    #[allow(clippy::too_many_arguments)]
438    pub fn gen_compact_code_blocks(
439        entity: &Entity,
440        with_serde: &WithSerde,
441        date_time_crate: &DateTimeCrate,
442        schema_name: &Option<String>,
443        serde_skip_deserializing_primary_key: bool,
444        serde_skip_hidden_column: bool,
445        model_extra_derives: &TokenStream,
446        model_extra_attributes: &TokenStream,
447        seaography: bool,
448        impl_active_model_behavior: bool,
449    ) -> Vec<TokenStream> {
450        let mut imports = Self::gen_import(with_serde);
451        imports.extend(Self::gen_import_active_enum(entity));
452        let mut code_blocks = vec![
453            imports,
454            Self::gen_compact_model_struct(
455                entity,
456                with_serde,
457                date_time_crate,
458                schema_name,
459                serde_skip_deserializing_primary_key,
460                serde_skip_hidden_column,
461                model_extra_derives,
462                model_extra_attributes,
463            ),
464            Self::gen_compact_relation_enum(entity),
465        ];
466        code_blocks.extend(Self::gen_impl_related(entity));
467        code_blocks.extend(Self::gen_impl_conjunct_related(entity));
468        if impl_active_model_behavior {
469            code_blocks.extend([Self::impl_active_model_behavior()]);
470        }
471        if seaography {
472            code_blocks.extend([Self::gen_related_entity(entity)]);
473        }
474        code_blocks
475    }
476
477    pub fn gen_import(with_serde: &WithSerde) -> TokenStream {
478        let prelude_import = quote!(
479            use sea_orm::entity::prelude::*;
480        );
481
482        match with_serde {
483            WithSerde::None => prelude_import,
484            WithSerde::Serialize => {
485                quote! {
486                    #prelude_import
487                    use serde::Serialize;
488                }
489            }
490            WithSerde::Deserialize => {
491                quote! {
492                    #prelude_import
493                    use serde::Deserialize;
494                }
495            }
496            WithSerde::Both => {
497                quote! {
498                    #prelude_import
499                    use serde::{Deserialize,Serialize};
500                }
501            }
502        }
503    }
504
505    pub fn gen_entity_struct() -> TokenStream {
506        quote! {
507            #[derive(Copy, Clone, Default, Debug, DeriveEntity)]
508            pub struct Entity;
509        }
510    }
511
512    pub fn gen_impl_entity_name(entity: &Entity, schema_name: &Option<String>) -> TokenStream {
513        let schema_name = match Self::gen_schema_name(schema_name) {
514            Some(schema_name) => quote! {
515                fn schema_name(&self) -> Option<&str> {
516                    Some(#schema_name)
517                }
518            },
519            None => quote! {},
520        };
521        let table_name = entity.table_name.as_str();
522        let table_name = quote! {
523            fn table_name(&self) -> &str {
524                #table_name
525            }
526        };
527        quote! {
528            impl EntityName for Entity {
529                #schema_name
530                #table_name
531            }
532        }
533    }
534
535    pub fn gen_import_active_enum(entity: &Entity) -> TokenStream {
536        entity
537            .columns
538            .iter()
539            .fold(
540                (TokenStream::new(), Vec::new()),
541                |(mut ts, mut enums), col| {
542                    if let sea_query::ColumnType::Enum { name, .. } = col.get_inner_col_type() {
543                        if !enums.contains(&name) {
544                            enums.push(name);
545                            let enum_name =
546                                format_ident!("{}", name.to_string().to_upper_camel_case());
547                            ts.extend([quote! {
548                                use super::sea_orm_active_enums::#enum_name;
549                            }]);
550                        }
551                    }
552                    (ts, enums)
553                },
554            )
555            .0
556    }
557
558    pub fn gen_model_struct(
559        entity: &Entity,
560        with_serde: &WithSerde,
561        date_time_crate: &DateTimeCrate,
562        serde_skip_deserializing_primary_key: bool,
563        serde_skip_hidden_column: bool,
564        model_extra_derives: &TokenStream,
565        model_extra_attributes: &TokenStream,
566    ) -> TokenStream {
567        let column_names_snake_case = entity.get_column_names_snake_case();
568        let column_rs_types = entity.get_column_rs_types(date_time_crate);
569        let if_eq_needed = entity.get_eq_needed();
570        let serde_attributes = entity.get_column_serde_attributes(
571            serde_skip_deserializing_primary_key,
572            serde_skip_hidden_column,
573        );
574        let extra_derive = with_serde.extra_derive();
575
576        quote! {
577            #[derive(Clone, Debug, PartialEq, DeriveModel, DeriveActiveModel #if_eq_needed #extra_derive #model_extra_derives)]
578            #model_extra_attributes
579            pub struct Model {
580                #(
581                    #serde_attributes
582                    pub #column_names_snake_case: #column_rs_types,
583                )*
584            }
585        }
586    }
587
588    pub fn gen_column_enum(entity: &Entity) -> TokenStream {
589        let column_variants = entity.columns.iter().map(|col| {
590            let variant = col.get_name_camel_case();
591            let mut variant = quote! { #variant };
592            if !col.is_snake_case_name() {
593                let column_name = &col.name;
594                variant = quote! {
595                    #[sea_orm(column_name = #column_name)]
596                    #variant
597                };
598            }
599            variant
600        });
601        quote! {
602            #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
603            pub enum Column {
604                #(#column_variants,)*
605            }
606        }
607    }
608
609    pub fn gen_primary_key_enum(entity: &Entity) -> TokenStream {
610        let primary_key_names_camel_case = entity.get_primary_key_names_camel_case();
611        quote! {
612            #[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)]
613            pub enum PrimaryKey {
614                #(#primary_key_names_camel_case,)*
615            }
616        }
617    }
618
619    pub fn gen_impl_primary_key(entity: &Entity, date_time_crate: &DateTimeCrate) -> TokenStream {
620        let primary_key_auto_increment = entity.get_primary_key_auto_increment();
621        let value_type = entity.get_primary_key_rs_type(date_time_crate);
622        quote! {
623            impl PrimaryKeyTrait for PrimaryKey {
624                type ValueType = #value_type;
625
626                fn auto_increment() -> bool {
627                    #primary_key_auto_increment
628                }
629            }
630        }
631    }
632
633    pub fn gen_relation_enum(entity: &Entity) -> TokenStream {
634        let relation_enum_name = entity.get_relation_enum_name();
635        quote! {
636            #[derive(Copy, Clone, Debug, EnumIter)]
637            pub enum Relation {
638                #(#relation_enum_name,)*
639            }
640        }
641    }
642
643    pub fn gen_impl_column_trait(entity: &Entity) -> TokenStream {
644        let column_names_camel_case = entity.get_column_names_camel_case();
645        let column_defs = entity.get_column_defs();
646        quote! {
647            impl ColumnTrait for Column {
648                type EntityName = Entity;
649
650                fn def(&self) -> ColumnDef {
651                    match self {
652                        #(Self::#column_names_camel_case => #column_defs,)*
653                    }
654                }
655            }
656        }
657    }
658
659    pub fn gen_impl_relation_trait(entity: &Entity) -> TokenStream {
660        let relation_enum_name = entity.get_relation_enum_name();
661        let relation_defs = entity.get_relation_defs();
662        let quoted = if relation_enum_name.is_empty() {
663            quote! {
664                panic!("No RelationDef")
665            }
666        } else {
667            quote! {
668                match self {
669                    #(Self::#relation_enum_name => #relation_defs,)*
670                }
671            }
672        };
673        quote! {
674            impl RelationTrait for Relation {
675                fn def(&self) -> RelationDef {
676                    #quoted
677                }
678            }
679        }
680    }
681
682    pub fn gen_impl_related(entity: &Entity) -> Vec<TokenStream> {
683        entity
684            .relations
685            .iter()
686            .filter(|rel| !rel.self_referencing && rel.num_suffix == 0 && rel.impl_related)
687            .map(|rel| {
688                let enum_name = rel.get_enum_name();
689                let module_name = rel.get_module_name();
690                let inner = quote! {
691                    fn to() -> RelationDef {
692                        Relation::#enum_name.def()
693                    }
694                };
695                if module_name.is_some() {
696                    quote! {
697                        impl Related<super::#module_name::Entity> for Entity { #inner }
698                    }
699                } else {
700                    quote! {
701                        impl Related<Entity> for Entity { #inner }
702                    }
703                }
704            })
705            .collect()
706    }
707
708    /// Used to generate `enum RelatedEntity` that is useful to the Seaography project
709    pub fn gen_related_entity(entity: &Entity) -> TokenStream {
710        let related_enum_name = entity.get_related_entity_enum_name();
711        let related_attrs = entity.get_related_entity_attrs();
712
713        quote! {
714            #[derive(Copy, Clone, Debug, EnumIter, DeriveRelatedEntity)]
715            pub enum RelatedEntity {
716                #(
717                    #related_attrs
718                    #related_enum_name
719                ),*
720            }
721        }
722    }
723
724    pub fn gen_impl_conjunct_related(entity: &Entity) -> Vec<TokenStream> {
725        let table_name_camel_case = entity.get_table_name_camel_case_ident();
726        let via_snake_case = entity.get_conjunct_relations_via_snake_case();
727        let to_snake_case = entity.get_conjunct_relations_to_snake_case();
728        let to_upper_camel_case = entity.get_conjunct_relations_to_upper_camel_case();
729        via_snake_case
730            .into_iter()
731            .zip(to_snake_case)
732            .zip(to_upper_camel_case)
733            .map(|((via_snake_case, to_snake_case), to_upper_camel_case)| {
734                quote! {
735                    impl Related<super::#to_snake_case::Entity> for Entity {
736                        fn to() -> RelationDef {
737                            super::#via_snake_case::Relation::#to_upper_camel_case.def()
738                        }
739
740                        fn via() -> Option<RelationDef> {
741                            Some(super::#via_snake_case::Relation::#table_name_camel_case.def().rev())
742                        }
743                    }
744                }
745            })
746            .collect()
747    }
748
749    pub fn impl_active_model_behavior() -> TokenStream {
750        quote! {
751            impl ActiveModelBehavior for ActiveModel {}
752        }
753    }
754
755    pub fn gen_mod(entity: &Entity) -> TokenStream {
756        let table_name_snake_case_ident = format_ident!(
757            "{}",
758            escape_rust_keyword(entity.get_table_name_snake_case_ident())
759        );
760        quote! {
761            pub mod #table_name_snake_case_ident;
762        }
763    }
764
765    pub fn gen_seaography_entity_mod(
766        entities: &[Entity],
767        enums: &BTreeMap<String, ActiveEnum>,
768    ) -> TokenStream {
769        let mut ts = TokenStream::new();
770        for entity in entities {
771            let table_name_snake_case_ident = format_ident!(
772                "{}",
773                escape_rust_keyword(entity.get_table_name_snake_case_ident())
774            );
775            ts = quote! {
776                #ts
777                #table_name_snake_case_ident,
778            }
779        }
780        ts = quote! {
781            seaography::register_entity_modules!([
782                #ts
783            ]);
784        };
785
786        let mut enum_ts = TokenStream::new();
787        for active_enum in enums.values() {
788            let enum_name = &active_enum.enum_name.to_string();
789            let enum_iden = format_ident!("{}", enum_name.to_upper_camel_case());
790            enum_ts = quote! {
791                #enum_ts
792                sea_orm_active_enums::#enum_iden,
793            }
794        }
795        if !enum_ts.is_empty() {
796            ts = quote! {
797                #ts
798
799                seaography::register_active_enums!([
800                    #enum_ts
801                ]);
802            };
803        }
804        ts
805    }
806
807    pub fn gen_prelude_use(entity: &Entity) -> TokenStream {
808        let table_name_snake_case_ident = entity.get_table_name_snake_case_ident();
809        let table_name_camel_case_ident = entity.get_table_name_camel_case_ident();
810        quote! {
811            pub use super::#table_name_snake_case_ident::Entity as #table_name_camel_case_ident;
812        }
813    }
814
815    #[allow(clippy::too_many_arguments)]
816    pub fn gen_compact_model_struct(
817        entity: &Entity,
818        with_serde: &WithSerde,
819        date_time_crate: &DateTimeCrate,
820        schema_name: &Option<String>,
821        serde_skip_deserializing_primary_key: bool,
822        serde_skip_hidden_column: bool,
823        model_extra_derives: &TokenStream,
824        model_extra_attributes: &TokenStream,
825    ) -> TokenStream {
826        let table_name = entity.table_name.as_str();
827        let column_names_snake_case = entity.get_column_names_snake_case();
828        let column_rs_types = entity.get_column_rs_types(date_time_crate);
829        let if_eq_needed = entity.get_eq_needed();
830        let primary_keys: Vec<String> = entity
831            .primary_keys
832            .iter()
833            .map(|pk| pk.name.clone())
834            .collect();
835        let attrs: Vec<TokenStream> = entity
836            .columns
837            .iter()
838            .map(|col| {
839                let mut attrs: Punctuated<_, Comma> = Punctuated::new();
840                let is_primary_key = primary_keys.contains(&col.name);
841                if !col.is_snake_case_name() {
842                    let column_name = &col.name;
843                    attrs.push(quote! { column_name = #column_name });
844                }
845                if is_primary_key {
846                    attrs.push(quote! { primary_key });
847                    if !col.auto_increment {
848                        attrs.push(quote! { auto_increment = false });
849                    }
850                }
851                if let Some(ts) = col.get_col_type_attrs() {
852                    attrs.extend([ts]);
853                    if !col.not_null {
854                        attrs.push(quote! { nullable });
855                    }
856                };
857                if col.unique {
858                    attrs.push(quote! { unique });
859                }
860                let mut ts = quote! {};
861                if !attrs.is_empty() {
862                    for (i, attr) in attrs.into_iter().enumerate() {
863                        if i > 0 {
864                            ts = quote! { #ts, };
865                        }
866                        ts = quote! { #ts #attr };
867                    }
868                    ts = quote! { #[sea_orm(#ts)] };
869                }
870                let serde_attribute = col.get_serde_attribute(
871                    is_primary_key,
872                    serde_skip_deserializing_primary_key,
873                    serde_skip_hidden_column,
874                );
875                ts = quote! {
876                    #ts
877                    #serde_attribute
878                };
879                ts
880            })
881            .collect();
882        let schema_name = match Self::gen_schema_name(schema_name) {
883            Some(schema_name) => quote! {
884                schema_name = #schema_name,
885            },
886            None => quote! {},
887        };
888        let extra_derive = with_serde.extra_derive();
889
890        quote! {
891            #[derive(Clone, Debug, PartialEq, DeriveEntityModel #if_eq_needed #extra_derive #model_extra_derives)]
892            #[sea_orm(
893                #schema_name
894                table_name = #table_name
895            )]
896            #model_extra_attributes
897            pub struct Model {
898                #(
899                    #attrs
900                    pub #column_names_snake_case: #column_rs_types,
901                )*
902            }
903        }
904    }
905
906    pub fn gen_compact_relation_enum(entity: &Entity) -> TokenStream {
907        let relation_enum_name = entity.get_relation_enum_name();
908        let attrs = entity.get_relation_attrs();
909        quote! {
910            #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
911            pub enum Relation {
912                #(
913                    #attrs
914                    #relation_enum_name,
915                )*
916            }
917        }
918    }
919
920    pub fn gen_schema_name(schema_name: &Option<String>) -> Option<TokenStream> {
921        schema_name
922            .as_ref()
923            .map(|schema_name| quote! { #schema_name })
924    }
925}
926
927#[cfg(test)]
928mod tests {
929    use crate::{
930        entity::writer::{bonus_attributes, bonus_derive},
931        Column, ConjunctRelation, DateTimeCrate, Entity, EntityWriter, PrimaryKey, Relation,
932        RelationType, WithSerde,
933    };
934    use pretty_assertions::assert_eq;
935    use proc_macro2::TokenStream;
936    use quote::quote;
937    use sea_query::{Alias, ColumnType, ForeignKeyAction, RcOrArc, SeaRc, StringLen};
938    use std::io::{self, BufRead, BufReader, Read};
939
940    fn setup() -> Vec<Entity> {
941        vec![
942            Entity {
943                table_name: "cake".to_owned(),
944                columns: vec![
945                    Column {
946                        name: "id".to_owned(),
947                        col_type: ColumnType::Integer,
948                        auto_increment: true,
949                        not_null: true,
950                        unique: false,
951                    },
952                    Column {
953                        name: "name".to_owned(),
954                        col_type: ColumnType::Text,
955                        auto_increment: false,
956                        not_null: false,
957                        unique: false,
958                    },
959                ],
960                relations: vec![Relation {
961                    ref_table: "fruit".to_owned(),
962                    columns: vec![],
963                    ref_columns: vec![],
964                    rel_type: RelationType::HasMany,
965                    on_delete: None,
966                    on_update: None,
967                    self_referencing: false,
968                    num_suffix: 0,
969                    impl_related: true,
970                }],
971                conjunct_relations: vec![ConjunctRelation {
972                    via: "cake_filling".to_owned(),
973                    to: "filling".to_owned(),
974                }],
975                primary_keys: vec![PrimaryKey {
976                    name: "id".to_owned(),
977                }],
978            },
979            Entity {
980                table_name: "_cake_filling_".to_owned(),
981                columns: vec![
982                    Column {
983                        name: "cake_id".to_owned(),
984                        col_type: ColumnType::Integer,
985                        auto_increment: false,
986                        not_null: true,
987                        unique: false,
988                    },
989                    Column {
990                        name: "filling_id".to_owned(),
991                        col_type: ColumnType::Integer,
992                        auto_increment: false,
993                        not_null: true,
994                        unique: false,
995                    },
996                ],
997                relations: vec![
998                    Relation {
999                        ref_table: "cake".to_owned(),
1000                        columns: vec!["cake_id".to_owned()],
1001                        ref_columns: vec!["id".to_owned()],
1002                        rel_type: RelationType::BelongsTo,
1003                        on_delete: Some(ForeignKeyAction::Cascade),
1004                        on_update: Some(ForeignKeyAction::Cascade),
1005                        self_referencing: false,
1006                        num_suffix: 0,
1007                        impl_related: true,
1008                    },
1009                    Relation {
1010                        ref_table: "filling".to_owned(),
1011                        columns: vec!["filling_id".to_owned()],
1012                        ref_columns: vec!["id".to_owned()],
1013                        rel_type: RelationType::BelongsTo,
1014                        on_delete: Some(ForeignKeyAction::Cascade),
1015                        on_update: Some(ForeignKeyAction::Cascade),
1016                        self_referencing: false,
1017                        num_suffix: 0,
1018                        impl_related: true,
1019                    },
1020                ],
1021                conjunct_relations: vec![],
1022                primary_keys: vec![
1023                    PrimaryKey {
1024                        name: "cake_id".to_owned(),
1025                    },
1026                    PrimaryKey {
1027                        name: "filling_id".to_owned(),
1028                    },
1029                ],
1030            },
1031            Entity {
1032                table_name: "cake_filling_price".to_owned(),
1033                columns: vec![
1034                    Column {
1035                        name: "cake_id".to_owned(),
1036                        col_type: ColumnType::Integer,
1037                        auto_increment: false,
1038                        not_null: true,
1039                        unique: false,
1040                    },
1041                    Column {
1042                        name: "filling_id".to_owned(),
1043                        col_type: ColumnType::Integer,
1044                        auto_increment: false,
1045                        not_null: true,
1046                        unique: false,
1047                    },
1048                    Column {
1049                        name: "price".to_owned(),
1050                        col_type: ColumnType::Decimal(None),
1051                        auto_increment: false,
1052                        not_null: true,
1053                        unique: false,
1054                    },
1055                ],
1056                relations: vec![Relation {
1057                    ref_table: "cake_filling".to_owned(),
1058                    columns: vec!["cake_id".to_owned(), "filling_id".to_owned()],
1059                    ref_columns: vec!["cake_id".to_owned(), "filling_id".to_owned()],
1060                    rel_type: RelationType::BelongsTo,
1061                    on_delete: None,
1062                    on_update: None,
1063                    self_referencing: false,
1064                    num_suffix: 0,
1065                    impl_related: true,
1066                }],
1067                conjunct_relations: vec![],
1068                primary_keys: vec![
1069                    PrimaryKey {
1070                        name: "cake_id".to_owned(),
1071                    },
1072                    PrimaryKey {
1073                        name: "filling_id".to_owned(),
1074                    },
1075                ],
1076            },
1077            Entity {
1078                table_name: "filling".to_owned(),
1079                columns: vec![
1080                    Column {
1081                        name: "id".to_owned(),
1082                        col_type: ColumnType::Integer,
1083                        auto_increment: true,
1084                        not_null: true,
1085                        unique: false,
1086                    },
1087                    Column {
1088                        name: "name".to_owned(),
1089                        col_type: ColumnType::String(StringLen::N(255)),
1090                        auto_increment: false,
1091                        not_null: true,
1092                        unique: false,
1093                    },
1094                ],
1095                relations: vec![],
1096                conjunct_relations: vec![ConjunctRelation {
1097                    via: "cake_filling".to_owned(),
1098                    to: "cake".to_owned(),
1099                }],
1100                primary_keys: vec![PrimaryKey {
1101                    name: "id".to_owned(),
1102                }],
1103            },
1104            Entity {
1105                table_name: "fruit".to_owned(),
1106                columns: vec![
1107                    Column {
1108                        name: "id".to_owned(),
1109                        col_type: ColumnType::Integer,
1110                        auto_increment: true,
1111                        not_null: true,
1112                        unique: false,
1113                    },
1114                    Column {
1115                        name: "name".to_owned(),
1116                        col_type: ColumnType::String(StringLen::N(255)),
1117                        auto_increment: false,
1118                        not_null: true,
1119                        unique: false,
1120                    },
1121                    Column {
1122                        name: "cake_id".to_owned(),
1123                        col_type: ColumnType::Integer,
1124                        auto_increment: false,
1125                        not_null: false,
1126                        unique: false,
1127                    },
1128                ],
1129                relations: vec![
1130                    Relation {
1131                        ref_table: "cake".to_owned(),
1132                        columns: vec!["cake_id".to_owned()],
1133                        ref_columns: vec!["id".to_owned()],
1134                        rel_type: RelationType::BelongsTo,
1135                        on_delete: None,
1136                        on_update: None,
1137                        self_referencing: false,
1138                        num_suffix: 0,
1139                        impl_related: true,
1140                    },
1141                    Relation {
1142                        ref_table: "vendor".to_owned(),
1143                        columns: vec![],
1144                        ref_columns: vec![],
1145                        rel_type: RelationType::HasMany,
1146                        on_delete: None,
1147                        on_update: None,
1148                        self_referencing: false,
1149                        num_suffix: 0,
1150                        impl_related: true,
1151                    },
1152                ],
1153                conjunct_relations: vec![],
1154                primary_keys: vec![PrimaryKey {
1155                    name: "id".to_owned(),
1156                }],
1157            },
1158            Entity {
1159                table_name: "vendor".to_owned(),
1160                columns: vec![
1161                    Column {
1162                        name: "id".to_owned(),
1163                        col_type: ColumnType::Integer,
1164                        auto_increment: true,
1165                        not_null: true,
1166                        unique: false,
1167                    },
1168                    Column {
1169                        name: "_name_".to_owned(),
1170                        col_type: ColumnType::String(StringLen::N(255)),
1171                        auto_increment: false,
1172                        not_null: true,
1173                        unique: false,
1174                    },
1175                    Column {
1176                        name: "fruitId".to_owned(),
1177                        col_type: ColumnType::Integer,
1178                        auto_increment: false,
1179                        not_null: false,
1180                        unique: false,
1181                    },
1182                ],
1183                relations: vec![Relation {
1184                    ref_table: "fruit".to_owned(),
1185                    columns: vec!["fruitId".to_owned()],
1186                    ref_columns: vec!["id".to_owned()],
1187                    rel_type: RelationType::BelongsTo,
1188                    on_delete: None,
1189                    on_update: None,
1190                    self_referencing: false,
1191                    num_suffix: 0,
1192                    impl_related: true,
1193                }],
1194                conjunct_relations: vec![],
1195                primary_keys: vec![PrimaryKey {
1196                    name: "id".to_owned(),
1197                }],
1198            },
1199            Entity {
1200                table_name: "rust_keyword".to_owned(),
1201                columns: vec![
1202                    Column {
1203                        name: "id".to_owned(),
1204                        col_type: ColumnType::Integer,
1205                        auto_increment: true,
1206                        not_null: true,
1207                        unique: false,
1208                    },
1209                    Column {
1210                        name: "testing".to_owned(),
1211                        col_type: ColumnType::TinyInteger,
1212                        auto_increment: false,
1213                        not_null: true,
1214                        unique: false,
1215                    },
1216                    Column {
1217                        name: "rust".to_owned(),
1218                        col_type: ColumnType::TinyUnsigned,
1219                        auto_increment: false,
1220                        not_null: true,
1221                        unique: false,
1222                    },
1223                    Column {
1224                        name: "keywords".to_owned(),
1225                        col_type: ColumnType::SmallInteger,
1226                        auto_increment: false,
1227                        not_null: true,
1228                        unique: false,
1229                    },
1230                    Column {
1231                        name: "type".to_owned(),
1232                        col_type: ColumnType::SmallUnsigned,
1233                        auto_increment: false,
1234                        not_null: true,
1235                        unique: false,
1236                    },
1237                    Column {
1238                        name: "typeof".to_owned(),
1239                        col_type: ColumnType::Integer,
1240                        auto_increment: false,
1241                        not_null: true,
1242                        unique: false,
1243                    },
1244                    Column {
1245                        name: "crate".to_owned(),
1246                        col_type: ColumnType::Unsigned,
1247                        auto_increment: false,
1248                        not_null: true,
1249                        unique: false,
1250                    },
1251                    Column {
1252                        name: "self".to_owned(),
1253                        col_type: ColumnType::BigInteger,
1254                        auto_increment: false,
1255                        not_null: true,
1256                        unique: false,
1257                    },
1258                    Column {
1259                        name: "self_id1".to_owned(),
1260                        col_type: ColumnType::BigUnsigned,
1261                        auto_increment: false,
1262                        not_null: true,
1263                        unique: false,
1264                    },
1265                    Column {
1266                        name: "self_id2".to_owned(),
1267                        col_type: ColumnType::Integer,
1268                        auto_increment: false,
1269                        not_null: true,
1270                        unique: false,
1271                    },
1272                    Column {
1273                        name: "fruit_id1".to_owned(),
1274                        col_type: ColumnType::Integer,
1275                        auto_increment: false,
1276                        not_null: true,
1277                        unique: false,
1278                    },
1279                    Column {
1280                        name: "fruit_id2".to_owned(),
1281                        col_type: ColumnType::Integer,
1282                        auto_increment: false,
1283                        not_null: true,
1284                        unique: false,
1285                    },
1286                    Column {
1287                        name: "cake_id".to_owned(),
1288                        col_type: ColumnType::Integer,
1289                        auto_increment: false,
1290                        not_null: true,
1291                        unique: false,
1292                    },
1293                ],
1294                relations: vec![
1295                    Relation {
1296                        ref_table: "rust_keyword".to_owned(),
1297                        columns: vec!["self_id1".to_owned()],
1298                        ref_columns: vec!["id".to_owned()],
1299                        rel_type: RelationType::BelongsTo,
1300                        on_delete: None,
1301                        on_update: None,
1302                        self_referencing: true,
1303                        num_suffix: 1,
1304                        impl_related: true,
1305                    },
1306                    Relation {
1307                        ref_table: "rust_keyword".to_owned(),
1308                        columns: vec!["self_id2".to_owned()],
1309                        ref_columns: vec!["id".to_owned()],
1310                        rel_type: RelationType::BelongsTo,
1311                        on_delete: None,
1312                        on_update: None,
1313                        self_referencing: true,
1314                        num_suffix: 2,
1315                        impl_related: true,
1316                    },
1317                    Relation {
1318                        ref_table: "fruit".to_owned(),
1319                        columns: vec!["fruit_id1".to_owned()],
1320                        ref_columns: vec!["id".to_owned()],
1321                        rel_type: RelationType::BelongsTo,
1322                        on_delete: None,
1323                        on_update: None,
1324                        self_referencing: false,
1325                        num_suffix: 1,
1326                        impl_related: true,
1327                    },
1328                    Relation {
1329                        ref_table: "fruit".to_owned(),
1330                        columns: vec!["fruit_id2".to_owned()],
1331                        ref_columns: vec!["id".to_owned()],
1332                        rel_type: RelationType::BelongsTo,
1333                        on_delete: None,
1334                        on_update: None,
1335                        self_referencing: false,
1336                        num_suffix: 2,
1337                        impl_related: true,
1338                    },
1339                    Relation {
1340                        ref_table: "cake".to_owned(),
1341                        columns: vec!["cake_id".to_owned()],
1342                        ref_columns: vec!["id".to_owned()],
1343                        rel_type: RelationType::BelongsTo,
1344                        on_delete: None,
1345                        on_update: None,
1346                        self_referencing: false,
1347                        num_suffix: 0,
1348                        impl_related: true,
1349                    },
1350                ],
1351                conjunct_relations: vec![],
1352                primary_keys: vec![PrimaryKey {
1353                    name: "id".to_owned(),
1354                }],
1355            },
1356            Entity {
1357                table_name: "cake_with_float".to_owned(),
1358                columns: vec![
1359                    Column {
1360                        name: "id".to_owned(),
1361                        col_type: ColumnType::Integer,
1362                        auto_increment: true,
1363                        not_null: true,
1364                        unique: false,
1365                    },
1366                    Column {
1367                        name: "name".to_owned(),
1368                        col_type: ColumnType::Text,
1369                        auto_increment: false,
1370                        not_null: false,
1371                        unique: false,
1372                    },
1373                    Column {
1374                        name: "price".to_owned(),
1375                        col_type: ColumnType::Float,
1376                        auto_increment: false,
1377                        not_null: false,
1378                        unique: false,
1379                    },
1380                ],
1381                relations: vec![Relation {
1382                    ref_table: "fruit".to_owned(),
1383                    columns: vec![],
1384                    ref_columns: vec![],
1385                    rel_type: RelationType::HasMany,
1386                    on_delete: None,
1387                    on_update: None,
1388                    self_referencing: false,
1389                    num_suffix: 0,
1390                    impl_related: true,
1391                }],
1392                conjunct_relations: vec![ConjunctRelation {
1393                    via: "cake_filling".to_owned(),
1394                    to: "filling".to_owned(),
1395                }],
1396                primary_keys: vec![PrimaryKey {
1397                    name: "id".to_owned(),
1398                }],
1399            },
1400            Entity {
1401                table_name: "cake_with_double".to_owned(),
1402                columns: vec![
1403                    Column {
1404                        name: "id".to_owned(),
1405                        col_type: ColumnType::Integer,
1406                        auto_increment: true,
1407                        not_null: true,
1408                        unique: false,
1409                    },
1410                    Column {
1411                        name: "name".to_owned(),
1412                        col_type: ColumnType::Text,
1413                        auto_increment: false,
1414                        not_null: false,
1415                        unique: false,
1416                    },
1417                    Column {
1418                        name: "price".to_owned(),
1419                        col_type: ColumnType::Double,
1420                        auto_increment: false,
1421                        not_null: false,
1422                        unique: false,
1423                    },
1424                ],
1425                relations: vec![Relation {
1426                    ref_table: "fruit".to_owned(),
1427                    columns: vec![],
1428                    ref_columns: vec![],
1429                    rel_type: RelationType::HasMany,
1430                    on_delete: None,
1431                    on_update: None,
1432                    self_referencing: false,
1433                    num_suffix: 0,
1434                    impl_related: true,
1435                }],
1436                conjunct_relations: vec![ConjunctRelation {
1437                    via: "cake_filling".to_owned(),
1438                    to: "filling".to_owned(),
1439                }],
1440                primary_keys: vec![PrimaryKey {
1441                    name: "id".to_owned(),
1442                }],
1443            },
1444            Entity {
1445                table_name: "collection".to_owned(),
1446                columns: vec![
1447                    Column {
1448                        name: "id".to_owned(),
1449                        col_type: ColumnType::Integer,
1450                        auto_increment: true,
1451                        not_null: true,
1452                        unique: false,
1453                    },
1454                    Column {
1455                        name: "integers".to_owned(),
1456                        col_type: ColumnType::Array(RcOrArc::new(ColumnType::Integer)),
1457                        auto_increment: false,
1458                        not_null: true,
1459                        unique: false,
1460                    },
1461                    Column {
1462                        name: "integers_opt".to_owned(),
1463                        col_type: ColumnType::Array(RcOrArc::new(ColumnType::Integer)),
1464                        auto_increment: false,
1465                        not_null: false,
1466                        unique: false,
1467                    },
1468                ],
1469                relations: vec![],
1470                conjunct_relations: vec![],
1471                primary_keys: vec![PrimaryKey {
1472                    name: "id".to_owned(),
1473                }],
1474            },
1475            Entity {
1476                table_name: "collection_float".to_owned(),
1477                columns: vec![
1478                    Column {
1479                        name: "id".to_owned(),
1480                        col_type: ColumnType::Integer,
1481                        auto_increment: true,
1482                        not_null: true,
1483                        unique: false,
1484                    },
1485                    Column {
1486                        name: "floats".to_owned(),
1487                        col_type: ColumnType::Array(RcOrArc::new(ColumnType::Float)),
1488                        auto_increment: false,
1489                        not_null: true,
1490                        unique: false,
1491                    },
1492                    Column {
1493                        name: "doubles".to_owned(),
1494                        col_type: ColumnType::Array(RcOrArc::new(ColumnType::Double)),
1495                        auto_increment: false,
1496                        not_null: true,
1497                        unique: false,
1498                    },
1499                ],
1500                relations: vec![],
1501                conjunct_relations: vec![],
1502                primary_keys: vec![PrimaryKey {
1503                    name: "id".to_owned(),
1504                }],
1505            },
1506            Entity {
1507                table_name: "parent".to_owned(),
1508                columns: vec![
1509                    Column {
1510                        name: "id1".to_owned(),
1511                        col_type: ColumnType::Integer,
1512                        auto_increment: false,
1513                        not_null: true,
1514                        unique: false,
1515                    },
1516                    Column {
1517                        name: "id2".to_owned(),
1518                        col_type: ColumnType::Integer,
1519                        auto_increment: false,
1520                        not_null: true,
1521                        unique: false,
1522                    },
1523                ],
1524                relations: vec![Relation {
1525                    ref_table: "child".to_owned(),
1526                    columns: vec![],
1527                    ref_columns: vec![],
1528                    rel_type: RelationType::HasMany,
1529                    on_delete: None,
1530                    on_update: None,
1531                    self_referencing: false,
1532                    num_suffix: 0,
1533                    impl_related: true,
1534                }],
1535                conjunct_relations: vec![],
1536                primary_keys: vec![
1537                    PrimaryKey {
1538                        name: "id1".to_owned(),
1539                    },
1540                    PrimaryKey {
1541                        name: "id2".to_owned(),
1542                    },
1543                ],
1544            },
1545            Entity {
1546                table_name: "child".to_owned(),
1547                columns: vec![
1548                    Column {
1549                        name: "id".to_owned(),
1550                        col_type: ColumnType::Integer,
1551                        auto_increment: true,
1552                        not_null: true,
1553                        unique: false,
1554                    },
1555                    Column {
1556                        name: "parent_id1".to_owned(),
1557                        col_type: ColumnType::Integer,
1558                        auto_increment: false,
1559                        not_null: true,
1560                        unique: false,
1561                    },
1562                    Column {
1563                        name: "parent_id2".to_owned(),
1564                        col_type: ColumnType::Integer,
1565                        auto_increment: false,
1566                        not_null: true,
1567                        unique: false,
1568                    },
1569                ],
1570                relations: vec![Relation {
1571                    ref_table: "parent".to_owned(),
1572                    columns: vec!["parent_id1".to_owned(), "parent_id2".to_owned()],
1573                    ref_columns: vec!["id1".to_owned(), "id2".to_owned()],
1574                    rel_type: RelationType::BelongsTo,
1575                    on_delete: None,
1576                    on_update: None,
1577                    self_referencing: false,
1578                    num_suffix: 0,
1579                    impl_related: true,
1580                }],
1581                conjunct_relations: vec![],
1582                primary_keys: vec![PrimaryKey {
1583                    name: "id".to_owned(),
1584                }],
1585            },
1586        ]
1587    }
1588
1589    fn parse_from_file<R>(inner: R) -> io::Result<TokenStream>
1590    where
1591        R: Read,
1592    {
1593        let mut reader = BufReader::new(inner);
1594        let mut lines: Vec<String> = Vec::new();
1595
1596        reader.read_until(b';', &mut Vec::new())?;
1597
1598        let mut line = String::new();
1599        while reader.read_line(&mut line)? > 0 {
1600            lines.push(line.to_owned());
1601            line.clear();
1602        }
1603        let content = lines.join("");
1604        Ok(content.parse().unwrap())
1605    }
1606
1607    #[test]
1608    fn test_gen_expanded_code_blocks() -> io::Result<()> {
1609        let entities = setup();
1610        const ENTITY_FILES: [&str; 13] = [
1611            include_str!("../../tests/expanded/cake.rs"),
1612            include_str!("../../tests/expanded/cake_filling.rs"),
1613            include_str!("../../tests/expanded/cake_filling_price.rs"),
1614            include_str!("../../tests/expanded/filling.rs"),
1615            include_str!("../../tests/expanded/fruit.rs"),
1616            include_str!("../../tests/expanded/vendor.rs"),
1617            include_str!("../../tests/expanded/rust_keyword.rs"),
1618            include_str!("../../tests/expanded/cake_with_float.rs"),
1619            include_str!("../../tests/expanded/cake_with_double.rs"),
1620            include_str!("../../tests/expanded/collection.rs"),
1621            include_str!("../../tests/expanded/collection_float.rs"),
1622            include_str!("../../tests/expanded/parent.rs"),
1623            include_str!("../../tests/expanded/child.rs"),
1624        ];
1625        const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 13] = [
1626            include_str!("../../tests/expanded_with_schema_name/cake.rs"),
1627            include_str!("../../tests/expanded_with_schema_name/cake_filling.rs"),
1628            include_str!("../../tests/expanded_with_schema_name/cake_filling_price.rs"),
1629            include_str!("../../tests/expanded_with_schema_name/filling.rs"),
1630            include_str!("../../tests/expanded_with_schema_name/fruit.rs"),
1631            include_str!("../../tests/expanded_with_schema_name/vendor.rs"),
1632            include_str!("../../tests/expanded_with_schema_name/rust_keyword.rs"),
1633            include_str!("../../tests/expanded_with_schema_name/cake_with_float.rs"),
1634            include_str!("../../tests/expanded_with_schema_name/cake_with_double.rs"),
1635            include_str!("../../tests/expanded_with_schema_name/collection.rs"),
1636            include_str!("../../tests/expanded_with_schema_name/collection_float.rs"),
1637            include_str!("../../tests/expanded_with_schema_name/parent.rs"),
1638            include_str!("../../tests/expanded_with_schema_name/child.rs"),
1639        ];
1640
1641        assert_eq!(entities.len(), ENTITY_FILES.len());
1642
1643        for (i, entity) in entities.iter().enumerate() {
1644            assert_eq!(
1645                parse_from_file(ENTITY_FILES[i].as_bytes())?.to_string(),
1646                EntityWriter::gen_expanded_code_blocks(
1647                    entity,
1648                    &crate::WithSerde::None,
1649                    &crate::DateTimeCrate::Chrono,
1650                    &None,
1651                    false,
1652                    false,
1653                    &TokenStream::new(),
1654                    &TokenStream::new(),
1655                    false,
1656                    true,
1657                )
1658                .into_iter()
1659                .skip(1)
1660                .fold(TokenStream::new(), |mut acc, tok| {
1661                    acc.extend(tok);
1662                    acc
1663                })
1664                .to_string()
1665            );
1666            assert_eq!(
1667                parse_from_file(ENTITY_FILES_WITH_SCHEMA_NAME[i].as_bytes())?.to_string(),
1668                EntityWriter::gen_expanded_code_blocks(
1669                    entity,
1670                    &crate::WithSerde::None,
1671                    &crate::DateTimeCrate::Chrono,
1672                    &Some("schema_name".to_owned()),
1673                    false,
1674                    false,
1675                    &TokenStream::new(),
1676                    &TokenStream::new(),
1677                    false,
1678                    true,
1679                )
1680                .into_iter()
1681                .skip(1)
1682                .fold(TokenStream::new(), |mut acc, tok| {
1683                    acc.extend(tok);
1684                    acc
1685                })
1686                .to_string()
1687            );
1688        }
1689
1690        Ok(())
1691    }
1692
1693    #[test]
1694    fn test_gen_compact_code_blocks() -> io::Result<()> {
1695        let entities = setup();
1696        const ENTITY_FILES: [&str; 13] = [
1697            include_str!("../../tests/compact/cake.rs"),
1698            include_str!("../../tests/compact/cake_filling.rs"),
1699            include_str!("../../tests/compact/cake_filling_price.rs"),
1700            include_str!("../../tests/compact/filling.rs"),
1701            include_str!("../../tests/compact/fruit.rs"),
1702            include_str!("../../tests/compact/vendor.rs"),
1703            include_str!("../../tests/compact/rust_keyword.rs"),
1704            include_str!("../../tests/compact/cake_with_float.rs"),
1705            include_str!("../../tests/compact/cake_with_double.rs"),
1706            include_str!("../../tests/compact/collection.rs"),
1707            include_str!("../../tests/compact/collection_float.rs"),
1708            include_str!("../../tests/compact/parent.rs"),
1709            include_str!("../../tests/compact/child.rs"),
1710        ];
1711        const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 13] = [
1712            include_str!("../../tests/compact_with_schema_name/cake.rs"),
1713            include_str!("../../tests/compact_with_schema_name/cake_filling.rs"),
1714            include_str!("../../tests/compact_with_schema_name/cake_filling_price.rs"),
1715            include_str!("../../tests/compact_with_schema_name/filling.rs"),
1716            include_str!("../../tests/compact_with_schema_name/fruit.rs"),
1717            include_str!("../../tests/compact_with_schema_name/vendor.rs"),
1718            include_str!("../../tests/compact_with_schema_name/rust_keyword.rs"),
1719            include_str!("../../tests/compact_with_schema_name/cake_with_float.rs"),
1720            include_str!("../../tests/compact_with_schema_name/cake_with_double.rs"),
1721            include_str!("../../tests/compact_with_schema_name/collection.rs"),
1722            include_str!("../../tests/compact_with_schema_name/collection_float.rs"),
1723            include_str!("../../tests/compact_with_schema_name/parent.rs"),
1724            include_str!("../../tests/compact_with_schema_name/child.rs"),
1725        ];
1726
1727        assert_eq!(entities.len(), ENTITY_FILES.len());
1728
1729        for (i, entity) in entities.iter().enumerate() {
1730            assert_eq!(
1731                parse_from_file(ENTITY_FILES[i].as_bytes())?.to_string(),
1732                EntityWriter::gen_compact_code_blocks(
1733                    entity,
1734                    &crate::WithSerde::None,
1735                    &crate::DateTimeCrate::Chrono,
1736                    &None,
1737                    false,
1738                    false,
1739                    &TokenStream::new(),
1740                    &TokenStream::new(),
1741                    false,
1742                    true,
1743                )
1744                .into_iter()
1745                .skip(1)
1746                .fold(TokenStream::new(), |mut acc, tok| {
1747                    acc.extend(tok);
1748                    acc
1749                })
1750                .to_string()
1751            );
1752            assert_eq!(
1753                parse_from_file(ENTITY_FILES_WITH_SCHEMA_NAME[i].as_bytes())?.to_string(),
1754                EntityWriter::gen_compact_code_blocks(
1755                    entity,
1756                    &crate::WithSerde::None,
1757                    &crate::DateTimeCrate::Chrono,
1758                    &Some("schema_name".to_owned()),
1759                    false,
1760                    false,
1761                    &TokenStream::new(),
1762                    &TokenStream::new(),
1763                    false,
1764                    true,
1765                )
1766                .into_iter()
1767                .skip(1)
1768                .fold(TokenStream::new(), |mut acc, tok| {
1769                    acc.extend(tok);
1770                    acc
1771                })
1772                .to_string()
1773            );
1774        }
1775
1776        Ok(())
1777    }
1778
1779    #[test]
1780    fn test_gen_with_serde() -> io::Result<()> {
1781        let cake_entity = setup().get(0).unwrap().clone();
1782
1783        assert_eq!(cake_entity.get_table_name_snake_case(), "cake");
1784
1785        // Compact code blocks
1786        assert_eq!(
1787            comparable_file_string(include_str!("../../tests/compact_with_serde/cake_none.rs"))?,
1788            generated_to_string(EntityWriter::gen_compact_code_blocks(
1789                &cake_entity,
1790                &WithSerde::None,
1791                &DateTimeCrate::Chrono,
1792                &None,
1793                false,
1794                false,
1795                &TokenStream::new(),
1796                &TokenStream::new(),
1797                false,
1798                true,
1799            ))
1800        );
1801        assert_eq!(
1802            comparable_file_string(include_str!(
1803                "../../tests/compact_with_serde/cake_serialize.rs"
1804            ))?,
1805            generated_to_string(EntityWriter::gen_compact_code_blocks(
1806                &cake_entity,
1807                &WithSerde::Serialize,
1808                &DateTimeCrate::Chrono,
1809                &None,
1810                false,
1811                false,
1812                &TokenStream::new(),
1813                &TokenStream::new(),
1814                false,
1815                true,
1816            ))
1817        );
1818        assert_eq!(
1819            comparable_file_string(include_str!(
1820                "../../tests/compact_with_serde/cake_deserialize.rs"
1821            ))?,
1822            generated_to_string(EntityWriter::gen_compact_code_blocks(
1823                &cake_entity,
1824                &WithSerde::Deserialize,
1825                &DateTimeCrate::Chrono,
1826                &None,
1827                true,
1828                false,
1829                &TokenStream::new(),
1830                &TokenStream::new(),
1831                false,
1832                true,
1833            ))
1834        );
1835        assert_eq!(
1836            comparable_file_string(include_str!("../../tests/compact_with_serde/cake_both.rs"))?,
1837            generated_to_string(EntityWriter::gen_compact_code_blocks(
1838                &cake_entity,
1839                &WithSerde::Both,
1840                &DateTimeCrate::Chrono,
1841                &None,
1842                true,
1843                false,
1844                &TokenStream::new(),
1845                &TokenStream::new(),
1846                false,
1847                true,
1848            ))
1849        );
1850
1851        // Expanded code blocks
1852        assert_eq!(
1853            comparable_file_string(include_str!("../../tests/expanded_with_serde/cake_none.rs"))?,
1854            generated_to_string(EntityWriter::gen_expanded_code_blocks(
1855                &cake_entity,
1856                &WithSerde::None,
1857                &DateTimeCrate::Chrono,
1858                &None,
1859                false,
1860                false,
1861                &TokenStream::new(),
1862                &TokenStream::new(),
1863                false,
1864                true,
1865            ))
1866        );
1867        assert_eq!(
1868            comparable_file_string(include_str!(
1869                "../../tests/expanded_with_serde/cake_serialize.rs"
1870            ))?,
1871            generated_to_string(EntityWriter::gen_expanded_code_blocks(
1872                &cake_entity,
1873                &WithSerde::Serialize,
1874                &DateTimeCrate::Chrono,
1875                &None,
1876                false,
1877                false,
1878                &TokenStream::new(),
1879                &TokenStream::new(),
1880                false,
1881                true,
1882            ))
1883        );
1884        assert_eq!(
1885            comparable_file_string(include_str!(
1886                "../../tests/expanded_with_serde/cake_deserialize.rs"
1887            ))?,
1888            generated_to_string(EntityWriter::gen_expanded_code_blocks(
1889                &cake_entity,
1890                &WithSerde::Deserialize,
1891                &DateTimeCrate::Chrono,
1892                &None,
1893                true,
1894                false,
1895                &TokenStream::new(),
1896                &TokenStream::new(),
1897                false,
1898                true,
1899            ))
1900        );
1901        assert_eq!(
1902            comparable_file_string(include_str!("../../tests/expanded_with_serde/cake_both.rs"))?,
1903            generated_to_string(EntityWriter::gen_expanded_code_blocks(
1904                &cake_entity,
1905                &WithSerde::Both,
1906                &DateTimeCrate::Chrono,
1907                &None,
1908                true,
1909                false,
1910                &TokenStream::new(),
1911                &TokenStream::new(),
1912                false,
1913                true,
1914            ))
1915        );
1916
1917        Ok(())
1918    }
1919
1920    #[test]
1921    fn test_gen_with_seaography() -> io::Result<()> {
1922        let cake_entity = Entity {
1923            table_name: "cake".to_owned(),
1924            columns: vec![
1925                Column {
1926                    name: "id".to_owned(),
1927                    col_type: ColumnType::Integer,
1928                    auto_increment: true,
1929                    not_null: true,
1930                    unique: false,
1931                },
1932                Column {
1933                    name: "name".to_owned(),
1934                    col_type: ColumnType::Text,
1935                    auto_increment: false,
1936                    not_null: false,
1937                    unique: false,
1938                },
1939                Column {
1940                    name: "base_id".to_owned(),
1941                    col_type: ColumnType::Integer,
1942                    auto_increment: false,
1943                    not_null: false,
1944                    unique: false,
1945                },
1946            ],
1947            relations: vec![
1948                Relation {
1949                    ref_table: "fruit".to_owned(),
1950                    columns: vec![],
1951                    ref_columns: vec![],
1952                    rel_type: RelationType::HasMany,
1953                    on_delete: None,
1954                    on_update: None,
1955                    self_referencing: false,
1956                    num_suffix: 0,
1957                    impl_related: true,
1958                },
1959                Relation {
1960                    ref_table: "cake".to_owned(),
1961                    columns: vec![],
1962                    ref_columns: vec![],
1963                    rel_type: RelationType::HasOne,
1964                    on_delete: None,
1965                    on_update: None,
1966                    self_referencing: true,
1967                    num_suffix: 0,
1968                    impl_related: true,
1969                },
1970            ],
1971            conjunct_relations: vec![ConjunctRelation {
1972                via: "cake_filling".to_owned(),
1973                to: "filling".to_owned(),
1974            }],
1975            primary_keys: vec![PrimaryKey {
1976                name: "id".to_owned(),
1977            }],
1978        };
1979
1980        assert_eq!(cake_entity.get_table_name_snake_case(), "cake");
1981
1982        // Compact code blocks
1983        assert_eq!(
1984            comparable_file_string(include_str!("../../tests/with_seaography/cake.rs"))?,
1985            generated_to_string(EntityWriter::gen_compact_code_blocks(
1986                &cake_entity,
1987                &WithSerde::None,
1988                &DateTimeCrate::Chrono,
1989                &None,
1990                false,
1991                false,
1992                &TokenStream::new(),
1993                &TokenStream::new(),
1994                true,
1995                true,
1996            ))
1997        );
1998
1999        // Expanded code blocks
2000        assert_eq!(
2001            comparable_file_string(include_str!("../../tests/with_seaography/cake_expanded.rs"))?,
2002            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2003                &cake_entity,
2004                &WithSerde::None,
2005                &DateTimeCrate::Chrono,
2006                &None,
2007                false,
2008                false,
2009                &TokenStream::new(),
2010                &TokenStream::new(),
2011                true,
2012                true,
2013            ))
2014        );
2015
2016        Ok(())
2017    }
2018
2019    #[test]
2020    fn test_gen_with_seaography_mod() -> io::Result<()> {
2021        use crate::ActiveEnum;
2022        use sea_query::IntoIden;
2023
2024        let entities = setup();
2025        let enums = vec![
2026            (
2027                "coinflip_result_type",
2028                ActiveEnum {
2029                    enum_name: Alias::new("coinflip_result_type").into_iden(),
2030                    values: vec!["HEADS", "TAILS"]
2031                        .into_iter()
2032                        .map(|variant| Alias::new(variant).into_iden())
2033                        .collect(),
2034                },
2035            ),
2036            (
2037                "media_type",
2038                ActiveEnum {
2039                    enum_name: Alias::new("media_type").into_iden(),
2040                    values: vec![
2041                        "UNKNOWN",
2042                        "BITMAP",
2043                        "DRAWING",
2044                        "AUDIO",
2045                        "VIDEO",
2046                        "MULTIMEDIA",
2047                        "OFFICE",
2048                        "TEXT",
2049                        "EXECUTABLE",
2050                        "ARCHIVE",
2051                        "3D",
2052                    ]
2053                    .into_iter()
2054                    .map(|variant| Alias::new(variant).into_iden())
2055                    .collect(),
2056                },
2057            ),
2058        ]
2059        .into_iter()
2060        .map(|(k, v)| (k.to_string(), v))
2061        .collect();
2062
2063        assert_eq!(
2064            comparable_file_string(include_str!("../../tests/with_seaography/mod.rs"))?,
2065            generated_to_string(vec![EntityWriter::gen_seaography_entity_mod(
2066                &entities, &enums,
2067            )])
2068        );
2069
2070        Ok(())
2071    }
2072
2073    #[test]
2074    fn test_gen_with_derives() -> io::Result<()> {
2075        let mut cake_entity = setup().get_mut(0).unwrap().clone();
2076
2077        assert_eq!(cake_entity.get_table_name_snake_case(), "cake");
2078
2079        // Compact code blocks
2080        assert_eq!(
2081            comparable_file_string(include_str!(
2082                "../../tests/compact_with_derives/cake_none.rs"
2083            ))?,
2084            generated_to_string(EntityWriter::gen_compact_code_blocks(
2085                &cake_entity,
2086                &WithSerde::None,
2087                &DateTimeCrate::Chrono,
2088                &None,
2089                false,
2090                false,
2091                &TokenStream::new(),
2092                &TokenStream::new(),
2093                false,
2094                true,
2095            ))
2096        );
2097        assert_eq!(
2098            comparable_file_string(include_str!("../../tests/compact_with_derives/cake_one.rs"))?,
2099            generated_to_string(EntityWriter::gen_compact_code_blocks(
2100                &cake_entity,
2101                &WithSerde::None,
2102                &DateTimeCrate::Chrono,
2103                &None,
2104                false,
2105                false,
2106                &bonus_derive(["ts_rs::TS"]),
2107                &TokenStream::new(),
2108                false,
2109                true,
2110            ))
2111        );
2112        assert_eq!(
2113            comparable_file_string(include_str!(
2114                "../../tests/compact_with_derives/cake_multiple.rs"
2115            ))?,
2116            generated_to_string(EntityWriter::gen_compact_code_blocks(
2117                &cake_entity,
2118                &WithSerde::None,
2119                &DateTimeCrate::Chrono,
2120                &None,
2121                false,
2122                false,
2123                &bonus_derive(["ts_rs::TS", "utoipa::ToSchema"]),
2124                &TokenStream::new(),
2125                false,
2126                true,
2127            ))
2128        );
2129
2130        // Expanded code blocks
2131        assert_eq!(
2132            comparable_file_string(include_str!(
2133                "../../tests/expanded_with_derives/cake_none.rs"
2134            ))?,
2135            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2136                &cake_entity,
2137                &WithSerde::None,
2138                &DateTimeCrate::Chrono,
2139                &None,
2140                false,
2141                false,
2142                &TokenStream::new(),
2143                &TokenStream::new(),
2144                false,
2145                true,
2146            ))
2147        );
2148        assert_eq!(
2149            comparable_file_string(include_str!(
2150                "../../tests/expanded_with_derives/cake_one.rs"
2151            ))?,
2152            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2153                &cake_entity,
2154                &WithSerde::None,
2155                &DateTimeCrate::Chrono,
2156                &None,
2157                false,
2158                false,
2159                &bonus_derive(["ts_rs::TS"]),
2160                &TokenStream::new(),
2161                false,
2162                true,
2163            ))
2164        );
2165        assert_eq!(
2166            comparable_file_string(include_str!(
2167                "../../tests/expanded_with_derives/cake_multiple.rs"
2168            ))?,
2169            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2170                &cake_entity,
2171                &WithSerde::None,
2172                &DateTimeCrate::Chrono,
2173                &None,
2174                false,
2175                false,
2176                &bonus_derive(["ts_rs::TS", "utoipa::ToSchema"]),
2177                &TokenStream::new(),
2178                false,
2179                true,
2180            ))
2181        );
2182
2183        // Make the `name` column of `cake` entity as hidden column
2184        cake_entity.columns[1].name = "_name".into();
2185
2186        assert_serde_variant_results(
2187            &cake_entity,
2188            &(
2189                include_str!("../../tests/compact_with_serde/cake_serialize_with_hidden_column.rs"),
2190                WithSerde::Serialize,
2191                None,
2192            ),
2193            Box::new(EntityWriter::gen_compact_code_blocks),
2194        )?;
2195        assert_serde_variant_results(
2196            &cake_entity,
2197            &(
2198                include_str!(
2199                    "../../tests/expanded_with_serde/cake_serialize_with_hidden_column.rs"
2200                ),
2201                WithSerde::Serialize,
2202                None,
2203            ),
2204            Box::new(EntityWriter::gen_expanded_code_blocks),
2205        )?;
2206
2207        Ok(())
2208    }
2209
2210    #[allow(clippy::type_complexity)]
2211    fn assert_serde_variant_results(
2212        cake_entity: &Entity,
2213        entity_serde_variant: &(&str, WithSerde, Option<String>),
2214        generator: Box<
2215            dyn Fn(
2216                &Entity,
2217                &WithSerde,
2218                &DateTimeCrate,
2219                &Option<String>,
2220                bool,
2221                bool,
2222                &TokenStream,
2223                &TokenStream,
2224                bool,
2225                bool,
2226            ) -> Vec<TokenStream>,
2227        >,
2228    ) -> io::Result<()> {
2229        let mut reader = BufReader::new(entity_serde_variant.0.as_bytes());
2230        let mut lines: Vec<String> = Vec::new();
2231        let serde_skip_deserializing_primary_key = matches!(
2232            entity_serde_variant.1,
2233            WithSerde::Both | WithSerde::Deserialize
2234        );
2235        let serde_skip_hidden_column = matches!(entity_serde_variant.1, WithSerde::Serialize);
2236
2237        reader.read_until(b'\n', &mut Vec::new())?;
2238
2239        let mut line = String::new();
2240        while reader.read_line(&mut line)? > 0 {
2241            lines.push(line.to_owned());
2242            line.clear();
2243        }
2244        let content = lines.join("");
2245        let expected: TokenStream = content.parse().unwrap();
2246        println!("{:?}", entity_serde_variant.1);
2247        let generated = generator(
2248            cake_entity,
2249            &entity_serde_variant.1,
2250            &DateTimeCrate::Chrono,
2251            &entity_serde_variant.2,
2252            serde_skip_deserializing_primary_key,
2253            serde_skip_hidden_column,
2254            &TokenStream::new(),
2255            &TokenStream::new(),
2256            false,
2257            true,
2258        )
2259        .into_iter()
2260        .fold(TokenStream::new(), |mut acc, tok| {
2261            acc.extend(tok);
2262            acc
2263        });
2264
2265        assert_eq!(expected.to_string(), generated.to_string());
2266        Ok(())
2267    }
2268
2269    #[test]
2270    fn test_gen_with_attributes() -> io::Result<()> {
2271        let cake_entity = setup().get(0).unwrap().clone();
2272
2273        assert_eq!(cake_entity.get_table_name_snake_case(), "cake");
2274
2275        // Compact code blocks
2276        assert_eq!(
2277            comparable_file_string(include_str!(
2278                "../../tests/compact_with_attributes/cake_none.rs"
2279            ))?,
2280            generated_to_string(EntityWriter::gen_compact_code_blocks(
2281                &cake_entity,
2282                &WithSerde::None,
2283                &DateTimeCrate::Chrono,
2284                &None,
2285                false,
2286                false,
2287                &TokenStream::new(),
2288                &TokenStream::new(),
2289                false,
2290                true,
2291            ))
2292        );
2293        assert_eq!(
2294            comparable_file_string(include_str!(
2295                "../../tests/compact_with_attributes/cake_one.rs"
2296            ))?,
2297            generated_to_string(EntityWriter::gen_compact_code_blocks(
2298                &cake_entity,
2299                &WithSerde::None,
2300                &DateTimeCrate::Chrono,
2301                &None,
2302                false,
2303                false,
2304                &TokenStream::new(),
2305                &bonus_attributes([r#"serde(rename_all = "camelCase")"#]),
2306                false,
2307                true,
2308            ))
2309        );
2310        assert_eq!(
2311            comparable_file_string(include_str!(
2312                "../../tests/compact_with_attributes/cake_multiple.rs"
2313            ))?,
2314            generated_to_string(EntityWriter::gen_compact_code_blocks(
2315                &cake_entity,
2316                &WithSerde::None,
2317                &DateTimeCrate::Chrono,
2318                &None,
2319                false,
2320                false,
2321                &TokenStream::new(),
2322                &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]),
2323                false,
2324                true,
2325            ))
2326        );
2327
2328        // Expanded code blocks
2329        assert_eq!(
2330            comparable_file_string(include_str!(
2331                "../../tests/expanded_with_attributes/cake_none.rs"
2332            ))?,
2333            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2334                &cake_entity,
2335                &WithSerde::None,
2336                &DateTimeCrate::Chrono,
2337                &None,
2338                false,
2339                false,
2340                &TokenStream::new(),
2341                &TokenStream::new(),
2342                false,
2343                true,
2344            ))
2345        );
2346        assert_eq!(
2347            comparable_file_string(include_str!(
2348                "../../tests/expanded_with_attributes/cake_one.rs"
2349            ))?,
2350            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2351                &cake_entity,
2352                &WithSerde::None,
2353                &DateTimeCrate::Chrono,
2354                &None,
2355                false,
2356                false,
2357                &TokenStream::new(),
2358                &bonus_attributes([r#"serde(rename_all = "camelCase")"#]),
2359                false,
2360                true,
2361            ))
2362        );
2363        assert_eq!(
2364            comparable_file_string(include_str!(
2365                "../../tests/expanded_with_attributes/cake_multiple.rs"
2366            ))?,
2367            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2368                &cake_entity,
2369                &WithSerde::None,
2370                &DateTimeCrate::Chrono,
2371                &None,
2372                false,
2373                false,
2374                &TokenStream::new(),
2375                &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]),
2376                false,
2377                true,
2378            ))
2379        );
2380
2381        Ok(())
2382    }
2383
2384    fn generated_to_string(generated: Vec<TokenStream>) -> String {
2385        generated
2386            .into_iter()
2387            .fold(TokenStream::new(), |mut acc, tok| {
2388                acc.extend(tok);
2389                acc
2390            })
2391            .to_string()
2392    }
2393
2394    fn comparable_file_string(file: &str) -> io::Result<String> {
2395        let mut reader = BufReader::new(file.as_bytes());
2396        let mut lines: Vec<String> = Vec::new();
2397
2398        reader.read_until(b'\n', &mut Vec::new())?;
2399
2400        let mut line = String::new();
2401        while reader.read_line(&mut line)? > 0 {
2402            lines.push(line.to_owned());
2403            line.clear();
2404        }
2405        let content = lines.join("");
2406        let expected: TokenStream = content.parse().unwrap();
2407
2408        Ok(expected.to_string())
2409    }
2410
2411    #[test]
2412    fn test_gen_postgres() -> io::Result<()> {
2413        let entities = vec![
2414            // This tests that the JsonBinary column type is annotated
2415            // correctly in compact entity form. More information can be found
2416            // in this issue:
2417            //
2418            // https://github.com/SeaQL/sea-orm/issues/1344
2419            Entity {
2420                table_name: "task".to_owned(),
2421                columns: vec![
2422                    Column {
2423                        name: "id".to_owned(),
2424                        col_type: ColumnType::Integer,
2425                        auto_increment: true,
2426                        not_null: true,
2427                        unique: false,
2428                    },
2429                    Column {
2430                        name: "payload".to_owned(),
2431                        col_type: ColumnType::Json,
2432                        auto_increment: false,
2433                        not_null: true,
2434                        unique: false,
2435                    },
2436                    Column {
2437                        name: "payload_binary".to_owned(),
2438                        col_type: ColumnType::JsonBinary,
2439                        auto_increment: false,
2440                        not_null: true,
2441                        unique: false,
2442                    },
2443                ],
2444                relations: vec![],
2445                conjunct_relations: vec![],
2446                primary_keys: vec![PrimaryKey {
2447                    name: "id".to_owned(),
2448                }],
2449            },
2450        ];
2451        const ENTITY_FILES: [&str; 1] = [include_str!("../../tests/postgres/binary_json.rs")];
2452
2453        const ENTITY_FILES_EXPANDED: [&str; 1] =
2454            [include_str!("../../tests/postgres/binary_json_expanded.rs")];
2455
2456        assert_eq!(entities.len(), ENTITY_FILES.len());
2457
2458        for (i, entity) in entities.iter().enumerate() {
2459            assert_eq!(
2460                parse_from_file(ENTITY_FILES[i].as_bytes())?.to_string(),
2461                EntityWriter::gen_compact_code_blocks(
2462                    entity,
2463                    &crate::WithSerde::None,
2464                    &crate::DateTimeCrate::Chrono,
2465                    &None,
2466                    false,
2467                    false,
2468                    &TokenStream::new(),
2469                    &TokenStream::new(),
2470                    false,
2471                    true,
2472                )
2473                .into_iter()
2474                .skip(1)
2475                .fold(TokenStream::new(), |mut acc, tok| {
2476                    acc.extend(tok);
2477                    acc
2478                })
2479                .to_string()
2480            );
2481            assert_eq!(
2482                parse_from_file(ENTITY_FILES_EXPANDED[i].as_bytes())?.to_string(),
2483                EntityWriter::gen_expanded_code_blocks(
2484                    entity,
2485                    &crate::WithSerde::None,
2486                    &crate::DateTimeCrate::Chrono,
2487                    &Some("schema_name".to_owned()),
2488                    false,
2489                    false,
2490                    &TokenStream::new(),
2491                    &TokenStream::new(),
2492                    false,
2493                    true,
2494                )
2495                .into_iter()
2496                .skip(1)
2497                .fold(TokenStream::new(), |mut acc, tok| {
2498                    acc.extend(tok);
2499                    acc
2500                })
2501                .to_string()
2502            );
2503        }
2504
2505        Ok(())
2506    }
2507
2508    #[test]
2509    fn test_gen_import_active_enum() -> io::Result<()> {
2510        let entities = vec![
2511            Entity {
2512                table_name: "tea_pairing".to_owned(),
2513                columns: vec![
2514                    Column {
2515                        name: "id".to_owned(),
2516                        col_type: ColumnType::Integer,
2517                        auto_increment: true,
2518                        not_null: true,
2519                        unique: false,
2520                    },
2521                    Column {
2522                        name: "first_tea".to_owned(),
2523                        col_type: ColumnType::Enum {
2524                            name: SeaRc::new(Alias::new("tea_enum")),
2525                            variants: vec![
2526                                SeaRc::new(Alias::new("everyday_tea")),
2527                                SeaRc::new(Alias::new("breakfast_tea")),
2528                            ],
2529                        },
2530                        auto_increment: false,
2531                        not_null: true,
2532                        unique: false,
2533                    },
2534                    Column {
2535                        name: "second_tea".to_owned(),
2536                        col_type: ColumnType::Enum {
2537                            name: SeaRc::new(Alias::new("tea_enum")),
2538                            variants: vec![
2539                                SeaRc::new(Alias::new("everyday_tea")),
2540                                SeaRc::new(Alias::new("breakfast_tea")),
2541                            ],
2542                        },
2543                        auto_increment: false,
2544                        not_null: true,
2545                        unique: false,
2546                    },
2547                ],
2548                relations: vec![],
2549                conjunct_relations: vec![],
2550                primary_keys: vec![PrimaryKey {
2551                    name: "id".to_owned(),
2552                }],
2553            },
2554            Entity {
2555                table_name: "tea_pairing_with_size".to_owned(),
2556                columns: vec![
2557                    Column {
2558                        name: "id".to_owned(),
2559                        col_type: ColumnType::Integer,
2560                        auto_increment: true,
2561                        not_null: true,
2562                        unique: false,
2563                    },
2564                    Column {
2565                        name: "first_tea".to_owned(),
2566                        col_type: ColumnType::Enum {
2567                            name: SeaRc::new(Alias::new("tea_enum")),
2568                            variants: vec![
2569                                SeaRc::new(Alias::new("everyday_tea")),
2570                                SeaRc::new(Alias::new("breakfast_tea")),
2571                            ],
2572                        },
2573                        auto_increment: false,
2574                        not_null: true,
2575                        unique: false,
2576                    },
2577                    Column {
2578                        name: "second_tea".to_owned(),
2579                        col_type: ColumnType::Enum {
2580                            name: SeaRc::new(Alias::new("tea_enum")),
2581                            variants: vec![
2582                                SeaRc::new(Alias::new("everyday_tea")),
2583                                SeaRc::new(Alias::new("breakfast_tea")),
2584                            ],
2585                        },
2586                        auto_increment: false,
2587                        not_null: true,
2588                        unique: false,
2589                    },
2590                    Column {
2591                        name: "size".to_owned(),
2592                        col_type: ColumnType::Enum {
2593                            name: SeaRc::new(Alias::new("tea_size")),
2594                            variants: vec![
2595                                SeaRc::new(Alias::new("small")),
2596                                SeaRc::new(Alias::new("medium")),
2597                                SeaRc::new(Alias::new("huge")),
2598                            ],
2599                        },
2600                        auto_increment: false,
2601                        not_null: true,
2602                        unique: false,
2603                    },
2604                ],
2605                relations: vec![],
2606                conjunct_relations: vec![],
2607                primary_keys: vec![PrimaryKey {
2608                    name: "id".to_owned(),
2609                }],
2610            },
2611        ];
2612
2613        assert_eq!(
2614            quote!(
2615                use super::sea_orm_active_enums::TeaEnum;
2616            )
2617            .to_string(),
2618            EntityWriter::gen_import_active_enum(&entities[0]).to_string()
2619        );
2620
2621        assert_eq!(
2622            quote!(
2623                use super::sea_orm_active_enums::TeaEnum;
2624                use super::sea_orm_active_enums::TeaSize;
2625            )
2626            .to_string(),
2627            EntityWriter::gen_import_active_enum(&entities[1]).to_string()
2628        );
2629
2630        Ok(())
2631    }
2632}