sea_orm_codegen/entity/
writer.rs

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