Skip to main content

sea_orm_codegen/entity/
writer.rs

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