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