Skip to main content

sea_orm_codegen/entity/writer/
compact.rs

1use super::*;
2
3impl EntityWriter {
4    #[allow(clippy::too_many_arguments)]
5    pub fn gen_compact_code_blocks(
6        entity: &Entity,
7        with_serde: &WithSerde,
8        column_option: &ColumnOption,
9        schema_name: &Option<String>,
10        serde_skip_deserializing_primary_key: bool,
11        serde_skip_hidden_column: bool,
12        model_extra_derives: &TokenStream,
13        model_extra_attributes: &TokenStream,
14        _column_extra_derives: &TokenStream,
15        seaography: bool,
16        impl_active_model_behavior: bool,
17    ) -> Vec<TokenStream> {
18        let mut imports = Self::gen_import(with_serde);
19        let active_enums = Self::gen_import_active_enum(entity);
20        imports.extend(active_enums.imports);
21        let mut code_blocks = vec![
22            imports,
23            Self::gen_compact_model_struct(
24                entity,
25                with_serde,
26                column_option,
27                schema_name,
28                serde_skip_deserializing_primary_key,
29                serde_skip_hidden_column,
30                model_extra_derives,
31                model_extra_attributes,
32                &active_enums.type_idents,
33            ),
34            Self::gen_compact_relation_enum(entity),
35        ];
36        code_blocks.extend(Self::gen_impl_related(entity));
37        code_blocks.extend(Self::gen_impl_conjunct_related(entity));
38        if impl_active_model_behavior {
39            code_blocks.extend([Self::impl_active_model_behavior()]);
40        }
41        if seaography {
42            code_blocks.extend([Self::gen_related_entity(entity)]);
43        }
44        code_blocks
45    }
46
47    #[allow(clippy::too_many_arguments)]
48    pub fn gen_compact_model_struct(
49        entity: &Entity,
50        with_serde: &WithSerde,
51        column_option: &ColumnOption,
52        schema_name: &Option<String>,
53        serde_skip_deserializing_primary_key: bool,
54        serde_skip_hidden_column: bool,
55        model_extra_derives: &TokenStream,
56        model_extra_attributes: &TokenStream,
57        active_enum_type_idents: &ActiveEnumTypeIdents,
58    ) -> TokenStream {
59        let table_name = entity.table_name.as_str();
60        let column_names_snake_case = entity.get_column_names_snake_case();
61        let column_rs_types = Self::get_column_rs_types_with_enum_idents(
62            entity,
63            column_option,
64            active_enum_type_idents,
65        );
66        let if_eq_needed = entity.get_eq_needed();
67        let primary_keys: Vec<String> = entity
68            .primary_keys
69            .iter()
70            .map(|pk| pk.name.clone())
71            .collect();
72        let attrs: Vec<TokenStream> = entity
73            .columns
74            .iter()
75            .map(|col| {
76                let mut attrs: Punctuated<_, Comma> = Punctuated::new();
77                let is_primary_key = primary_keys.contains(&col.name);
78                if !col.is_snake_case_name() {
79                    let column_name = &col.name;
80                    attrs.push(quote! { column_name = #column_name });
81                }
82                if is_primary_key {
83                    attrs.push(quote! { primary_key });
84                    if !col.auto_increment {
85                        attrs.push(quote! { auto_increment = false });
86                    }
87                }
88                if let Some(ts) = col.get_col_type_attrs() {
89                    attrs.extend([ts]);
90                    if !col.not_null {
91                        attrs.push(quote! { nullable });
92                    }
93                };
94                if col.unique {
95                    attrs.push(quote! { unique });
96                } else if let Some(unique_key) = &col.unique_key {
97                    attrs.push(quote! { unique_key = #unique_key });
98                }
99                let mut ts = quote! {};
100                if !attrs.is_empty() {
101                    for (i, attr) in attrs.into_iter().enumerate() {
102                        if i > 0 {
103                            ts = quote! { #ts, };
104                        }
105                        ts = quote! { #ts #attr };
106                    }
107                    ts = quote! { #[sea_orm(#ts)] };
108                }
109                let serde_attribute = col.get_serde_attribute(
110                    is_primary_key,
111                    serde_skip_deserializing_primary_key,
112                    serde_skip_hidden_column,
113                );
114                ts = quote! {
115                    #ts
116                    #serde_attribute
117                };
118                ts
119            })
120            .collect();
121        let schema_name = match Self::gen_schema_name(schema_name) {
122            Some(schema_name) => quote! {
123                schema_name = #schema_name,
124            },
125            None => quote! {},
126        };
127        let extra_derive = with_serde.extra_derive();
128
129        quote! {
130            #[derive(Clone, Debug, PartialEq #if_eq_needed, DeriveEntityModel #extra_derive #model_extra_derives)]
131            #[sea_orm(
132                #schema_name
133                table_name = #table_name
134            )]
135            #model_extra_attributes
136            pub struct Model {
137                #(
138                    #attrs
139                    pub #column_names_snake_case: #column_rs_types,
140                )*
141            }
142        }
143    }
144
145    pub fn gen_compact_relation_enum(entity: &Entity) -> TokenStream {
146        let attrs = entity.get_relation_attrs();
147        let relation_enum_name = entity.get_relation_enum_name();
148        quote! {
149            #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
150            pub enum Relation {
151                #(
152                    #attrs
153                    #relation_enum_name,
154                )*
155            }
156        }
157    }
158}