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