From a1bf662fc3e2726633b00ab179894ec3c13b3abc Mon Sep 17 00:00:00 2001 From: Nahua Date: Sun, 10 Jul 2022 09:22:47 +0200 Subject: [PATCH] #661: Add flag to `sea-orm-cli` to generate code for time crate (#724) * Implement entity writer context and date time crate enum * Use entity writer context to render correct rust type according to chrono/time enum * Add unit test to get col rs type with time crate * Add date-time-crate flag for cli * Edit usage of updated get_rs_type method * Use arg enum for date time crate * Move date time crate arg enum conversion into commands module * Directly use arg enum * Fixing import issues for arg enum w/ help from Billy --- sea-orm-cli/src/cli.rs | 17 ++- sea-orm-cli/src/commands.rs | 24 +++- sea-orm-codegen/src/entity/base_entity.rs | 30 +++-- sea-orm-codegen/src/entity/column.rs | 86 ++++++++++++-- sea-orm-codegen/src/entity/writer.rs | 134 ++++++++++++++++------ 5 files changed, 228 insertions(+), 63 deletions(-) diff --git a/sea-orm-cli/src/cli.rs b/sea-orm-cli/src/cli.rs index df86fe68..4a918a62 100644 --- a/sea-orm-cli/src/cli.rs +++ b/sea-orm-cli/src/cli.rs @@ -1,4 +1,4 @@ -use clap::{ArgGroup, Parser, Subcommand}; +use clap::{ArgEnum, ArgGroup, Parser, Subcommand}; #[derive(Parser, Debug)] #[clap(version)] @@ -166,5 +166,20 @@ pub enum GenerateSubcommands { help = "Automatically derive serde Serialize / Deserialize traits for the entity (none, serialize, deserialize, both)" )] with_serde: String, + + #[clap( + arg_enum, + value_parser, + long, + default_value = "chrono", + help = "The datetime crate to use for generating entities." + )] + date_time_crate: DateTimeCrate, }, } + +#[derive(ArgEnum, Copy, Clone, Debug, PartialEq)] +pub enum DateTimeCrate { + Chrono, + Time, +} diff --git a/sea-orm-cli/src/commands.rs b/sea-orm-cli/src/commands.rs index c55b66ef..9ed51339 100644 --- a/sea-orm-cli/src/commands.rs +++ b/sea-orm-cli/src/commands.rs @@ -1,11 +1,13 @@ use chrono::Local; use regex::Regex; -use sea_orm_codegen::{EntityTransformer, OutputFile, WithSerde}; +use sea_orm_codegen::{ + EntityTransformer, EntityWriterContext, OutputFile, WithSerde, DateTimeCrate as CodegenDateTimeCrate, +}; use std::{error::Error, fmt::Display, fs, io::Write, path::Path, process::Command, str::FromStr}; use tracing_subscriber::{prelude::*, EnvFilter}; use url::Url; -use crate::{GenerateSubcommands, MigrateSubcommands}; +use crate::{DateTimeCrate, GenerateSubcommands, MigrateSubcommands}; pub async fn run_generate_command( command: GenerateSubcommands, @@ -23,6 +25,7 @@ pub async fn run_generate_command( database_schema, database_url, with_serde, + date_time_crate, } => { if verbose { let _ = tracing_subscriber::fmt() @@ -173,8 +176,12 @@ pub async fn run_generate_command( _ => unimplemented!("{} is not supported", url.scheme()), }; - let output = EntityTransformer::transform(table_stmts)? - .generate(expanded_format, WithSerde::from_str(&with_serde).unwrap()); + let writer_context = EntityWriterContext::new( + expanded_format, + WithSerde::from_str(&with_serde).unwrap(), + date_time_crate.into(), + ); + let output = EntityTransformer::transform(table_stmts)?.generate(&writer_context); let dir = Path::new(&output_dir); fs::create_dir_all(dir)?; @@ -375,6 +382,15 @@ where ::std::process::exit(1); } +impl From for CodegenDateTimeCrate { + fn from(date_time_crate: DateTimeCrate) -> CodegenDateTimeCrate { + match date_time_crate { + DateTimeCrate::Chrono => CodegenDateTimeCrate::Chrono, + DateTimeCrate::Time => CodegenDateTimeCrate::Time, + } + } +} + #[cfg(test)] mod tests { use clap::StructOpt; diff --git a/sea-orm-codegen/src/entity/base_entity.rs b/sea-orm-codegen/src/entity/base_entity.rs index 5fdf34de..a745b1f7 100644 --- a/sea-orm-codegen/src/entity/base_entity.rs +++ b/sea-orm-codegen/src/entity/base_entity.rs @@ -1,4 +1,4 @@ -use crate::{Column, ConjunctRelation, PrimaryKey, Relation}; +use crate::{Column, ConjunctRelation, DateTimeCrate, PrimaryKey, Relation}; use heck::{CamelCase, SnakeCase}; use proc_macro2::{Ident, TokenStream}; use quote::format_ident; @@ -43,11 +43,11 @@ impl Entity { .collect() } - pub fn get_column_rs_types(&self) -> Vec { + pub fn get_column_rs_types(&self, date_time_crate: &DateTimeCrate) -> Vec { self.columns .clone() .into_iter() - .map(|col| col.get_rs_type()) + .map(|col| col.get_rs_type(date_time_crate)) .collect() } @@ -100,7 +100,7 @@ impl Entity { format_ident!("{}", auto_increment) } - pub fn get_primary_key_rs_type(&self) -> TokenStream { + pub fn get_primary_key_rs_type(&self, date_time_crate: &DateTimeCrate) -> TokenStream { let types = self .primary_keys .iter() @@ -109,7 +109,7 @@ impl Entity { .iter() .find(|col| col.name.eq(&primary_key.name)) .unwrap() - .get_rs_type() + .get_rs_type(date_time_crate) .to_string() }) .collect::>(); @@ -149,7 +149,7 @@ impl Entity { #[cfg(test)] mod tests { - use crate::{Column, Entity, PrimaryKey, Relation, RelationType}; + use crate::{Column, DateTimeCrate, Entity, PrimaryKey, Relation, RelationType}; use quote::format_ident; use sea_query::{ColumnType, ForeignKeyAction}; @@ -260,10 +260,16 @@ mod tests { fn test_get_column_rs_types() { let entity = setup(); - for (i, elem) in entity.get_column_rs_types().into_iter().enumerate() { + for (i, elem) in entity + .get_column_rs_types(&DateTimeCrate::Chrono) + .into_iter() + .enumerate() + { assert_eq!( elem.to_string(), - entity.columns[i].get_rs_type().to_string() + entity.columns[i] + .get_rs_type(&DateTimeCrate::Chrono) + .to_string() ); } } @@ -363,8 +369,12 @@ mod tests { let entity = setup(); assert_eq!( - entity.get_primary_key_rs_type().to_string(), - entity.columns[0].get_rs_type().to_string() + entity + .get_primary_key_rs_type(&DateTimeCrate::Chrono) + .to_string(), + entity.columns[0] + .get_rs_type(&DateTimeCrate::Chrono) + .to_string() ); } diff --git a/sea-orm-codegen/src/entity/column.rs b/sea-orm-codegen/src/entity/column.rs index 89d998c5..40969d62 100644 --- a/sea-orm-codegen/src/entity/column.rs +++ b/sea-orm-codegen/src/entity/column.rs @@ -1,4 +1,4 @@ -use crate::util::escape_rust_keyword; +use crate::{util::escape_rust_keyword, DateTimeCrate}; use heck::{CamelCase, SnakeCase}; use proc_macro2::{Ident, TokenStream}; use quote::{format_ident, quote}; @@ -27,7 +27,7 @@ impl Column { self.name.to_snake_case() == self.name } - pub fn get_rs_type(&self) -> TokenStream { + pub fn get_rs_type(&self, date_time_crate: &DateTimeCrate) -> TokenStream { #[allow(unreachable_patterns)] let ident: TokenStream = match &self.col_type { ColumnType::Char(_) @@ -45,11 +45,27 @@ impl Column { ColumnType::Float(_) => "f32".to_owned(), ColumnType::Double(_) => "f64".to_owned(), ColumnType::Json | ColumnType::JsonBinary => "Json".to_owned(), - ColumnType::Date => "Date".to_owned(), - ColumnType::Time(_) => "Time".to_owned(), - ColumnType::DateTime(_) => "DateTime".to_owned(), - ColumnType::Timestamp(_) => "DateTimeUtc".to_owned(), - ColumnType::TimestampWithTimeZone(_) => "DateTimeWithTimeZone".to_owned(), + ColumnType::Date => match date_time_crate { + DateTimeCrate::Chrono => "Date".to_owned(), + DateTimeCrate::Time => "TimeDate".to_owned(), + }, + ColumnType::Time(_) => match date_time_crate { + DateTimeCrate::Chrono => "Time".to_owned(), + DateTimeCrate::Time => "TimeTime".to_owned(), + }, + ColumnType::DateTime(_) => match date_time_crate { + DateTimeCrate::Chrono => "DateTime".to_owned(), + DateTimeCrate::Time => "TimeDateTime".to_owned(), + }, + ColumnType::Timestamp(_) => match date_time_crate { + DateTimeCrate::Chrono => "DateTimeUtc".to_owned(), + // ColumnType::Timpestamp(_) => time::PrimitiveDateTime: https://docs.rs/sqlx/0.3.5/sqlx/postgres/types/index.html#time + DateTimeCrate::Time => "TimeDateTime".to_owned(), + }, + ColumnType::TimestampWithTimeZone(_) => match date_time_crate { + DateTimeCrate::Chrono => "DateTimeWithTimeZone".to_owned(), + DateTimeCrate::Time => "TimeDateTimeWithTimeZone".to_owned(), + }, ColumnType::Decimal(_) | ColumnType::Money(_) => "Decimal".to_owned(), ColumnType::Uuid => "Uuid".to_owned(), ColumnType::Binary(_) => "Vec".to_owned(), @@ -147,7 +163,7 @@ impl Column { pub fn get_info(&self) -> String { let mut info = String::new(); - let type_info = self.get_rs_type().to_string().replace(' ', ""); + let type_info = self.get_rs_type(&DateTimeCrate::Chrono).to_string().replace(' ', ""); let col_info = self.col_info(); write!( &mut info, @@ -210,7 +226,7 @@ impl From<&ColumnDef> for Column { #[cfg(test)] mod tests { - use crate::Column; + use crate::{Column, DateTimeCrate}; use proc_macro2::TokenStream; use quote::quote; use sea_query::{Alias, BlobSize, ColumnDef, ColumnType, SeaRc}; @@ -312,8 +328,9 @@ mod tests { } #[test] - fn test_get_rs_type() { + fn test_get_rs_type_with_chrono() { let columns = setup(); + let chrono_crate = DateTimeCrate::Chrono; let rs_types = vec![ "String", "String", @@ -339,11 +356,56 @@ mod tests { let rs_type: TokenStream = rs_type.parse().unwrap(); col.not_null = true; - assert_eq!(col.get_rs_type().to_string(), quote!(#rs_type).to_string()); + assert_eq!( + col.get_rs_type(&chrono_crate).to_string(), + quote!(#rs_type).to_string() + ); col.not_null = false; assert_eq!( - col.get_rs_type().to_string(), + col.get_rs_type(&chrono_crate).to_string(), + quote!(Option<#rs_type>).to_string() + ); + } + } + + #[test] + fn test_get_rs_type_with_time() { + let columns = setup(); + let time_crate = DateTimeCrate::Time; + let rs_types = vec![ + "String", + "String", + "i8", + "u8", + "i16", + "u16", + "i32", + "u32", + "i64", + "u64", + "f32", + "f64", + "Vec", + "bool", + "TimeDate", + "TimeTime", + "TimeDateTime", + "TimeDateTime", + "TimeDateTimeWithTimeZone", + ]; + for (mut col, rs_type) in columns.into_iter().zip(rs_types) { + let rs_type: TokenStream = rs_type.parse().unwrap(); + + col.not_null = true; + assert_eq!( + col.get_rs_type(&time_crate).to_string(), + quote!(#rs_type).to_string() + ); + + col.not_null = false; + assert_eq!( + col.get_rs_type(&time_crate).to_string(), quote!(Option<#rs_type>).to_string() ); } diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 922f4bce..720e7791 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/sea-orm-codegen/src/entity/writer.rs @@ -29,6 +29,19 @@ pub enum WithSerde { Both, } +#[derive(Debug)] +pub enum DateTimeCrate { + Chrono, + Time, +} + +#[derive(Debug)] +pub struct EntityWriterContext { + pub(crate) expanded_format: bool, + pub(crate) with_serde: WithSerde, + pub(crate) date_time_crate: DateTimeCrate, +} + impl WithSerde { pub fn extra_derive(&self) -> TokenStream { let mut extra_derive = match self { @@ -79,19 +92,33 @@ impl FromStr for WithSerde { } } +impl EntityWriterContext { + pub fn new( + expanded_format: bool, + with_serde: WithSerde, + date_time_crate: DateTimeCrate, + ) -> Self { + Self { + expanded_format, + with_serde, + date_time_crate, + } + } +} + impl EntityWriter { - pub fn generate(self, expanded_format: bool, with_serde: WithSerde) -> WriterOutput { + pub fn generate(self, context: &EntityWriterContext) -> WriterOutput { let mut files = Vec::new(); - files.extend(self.write_entities(expanded_format, &with_serde)); + files.extend(self.write_entities(context)); files.push(self.write_mod()); files.push(self.write_prelude()); if !self.enums.is_empty() { - files.push(self.write_sea_orm_active_enums(&with_serde)); + files.push(self.write_sea_orm_active_enums(&context.with_serde)); } WriterOutput { files } } - pub fn write_entities(&self, expanded_format: bool, with_serde: &WithSerde) -> Vec { + pub fn write_entities(&self, context: &EntityWriterContext) -> Vec { self.entities .iter() .map(|entity| { @@ -109,10 +136,18 @@ impl EntityWriter { let mut lines = Vec::new(); Self::write_doc_comment(&mut lines); - let code_blocks = if expanded_format { - Self::gen_expanded_code_blocks(entity, with_serde) + let code_blocks = if context.expanded_format { + Self::gen_expanded_code_blocks( + entity, + &context.with_serde, + &context.date_time_crate, + ) } else { - Self::gen_compact_code_blocks(entity, with_serde) + Self::gen_compact_code_blocks( + entity, + &context.with_serde, + &context.date_time_crate, + ) }; Self::write(&mut lines, code_blocks); OutputFile { @@ -196,17 +231,21 @@ impl EntityWriter { lines.push("".to_owned()); } - pub fn gen_expanded_code_blocks(entity: &Entity, with_serde: &WithSerde) -> Vec { + pub fn gen_expanded_code_blocks( + entity: &Entity, + with_serde: &WithSerde, + date_time_crate: &DateTimeCrate, + ) -> Vec { let mut imports = Self::gen_import(with_serde); imports.extend(Self::gen_import_active_enum(entity)); let mut code_blocks = vec![ imports, Self::gen_entity_struct(), Self::gen_impl_entity_name(entity), - Self::gen_model_struct(entity, with_serde), + Self::gen_model_struct(entity, with_serde, date_time_crate), Self::gen_column_enum(entity), Self::gen_primary_key_enum(entity), - Self::gen_impl_primary_key(entity), + Self::gen_impl_primary_key(entity, date_time_crate), Self::gen_relation_enum(entity), Self::gen_impl_column_trait(entity), Self::gen_impl_relation_trait(entity), @@ -217,10 +256,17 @@ impl EntityWriter { code_blocks } - pub fn gen_compact_code_blocks(entity: &Entity, with_serde: &WithSerde) -> Vec { + pub fn gen_compact_code_blocks( + entity: &Entity, + with_serde: &WithSerde, + date_time_crate: &DateTimeCrate, + ) -> Vec { let mut imports = Self::gen_import(with_serde); imports.extend(Self::gen_import_active_enum(entity)); - let mut code_blocks = vec![imports, Self::gen_compact_model_struct(entity, with_serde)]; + let mut code_blocks = vec![ + imports, + Self::gen_compact_model_struct(entity, with_serde, date_time_crate), + ]; let relation_defs = if entity.get_relation_enum_name().is_empty() { vec![ Self::gen_relation_enum(entity), @@ -299,9 +345,13 @@ impl EntityWriter { }) } - pub fn gen_model_struct(entity: &Entity, with_serde: &WithSerde) -> TokenStream { + pub fn gen_model_struct( + entity: &Entity, + with_serde: &WithSerde, + date_time_crate: &DateTimeCrate, + ) -> TokenStream { let column_names_snake_case = entity.get_column_names_snake_case(); - let column_rs_types = entity.get_column_rs_types(); + let column_rs_types = entity.get_column_rs_types(date_time_crate); let extra_derive = with_serde.extra_derive(); @@ -344,9 +394,9 @@ impl EntityWriter { } } - pub fn gen_impl_primary_key(entity: &Entity) -> TokenStream { + pub fn gen_impl_primary_key(entity: &Entity, date_time_crate: &DateTimeCrate) -> TokenStream { let primary_key_auto_increment = entity.get_primary_key_auto_increment(); - let value_type = entity.get_primary_key_rs_type(); + let value_type = entity.get_primary_key_rs_type(date_time_crate); quote! { impl PrimaryKeyTrait for PrimaryKey { type ValueType = #value_type; @@ -479,10 +529,14 @@ impl EntityWriter { } } - pub fn gen_compact_model_struct(entity: &Entity, with_serde: &WithSerde) -> TokenStream { + pub fn gen_compact_model_struct( + entity: &Entity, + with_serde: &WithSerde, + date_time_crate: &DateTimeCrate, + ) -> TokenStream { let table_name = entity.table_name.as_str(); let column_names_snake_case = entity.get_column_names_snake_case(); - let column_rs_types = entity.get_column_rs_types(); + let column_rs_types = entity.get_column_rs_types(date_time_crate); let primary_keys: Vec = entity .primary_keys .iter() @@ -561,8 +615,8 @@ impl EntityWriter { #[cfg(test)] mod tests { use crate::{ - Column, ConjunctRelation, Entity, EntityWriter, PrimaryKey, Relation, RelationType, - WithSerde, + Column, ConjunctRelation, DateTimeCrate, Entity, EntityWriter, PrimaryKey, Relation, + RelationType, WithSerde, }; use pretty_assertions::assert_eq; use proc_macro2::TokenStream; @@ -958,13 +1012,17 @@ mod tests { } let content = lines.join(""); let expected: TokenStream = content.parse().unwrap(); - let generated = EntityWriter::gen_expanded_code_blocks(entity, &crate::WithSerde::None) - .into_iter() - .skip(1) - .fold(TokenStream::new(), |mut acc, tok| { - acc.extend(tok); - acc - }); + let generated = EntityWriter::gen_expanded_code_blocks( + entity, + &crate::WithSerde::None, + &crate::DateTimeCrate::Chrono, + ) + .into_iter() + .skip(1) + .fold(TokenStream::new(), |mut acc, tok| { + acc.extend(tok); + acc + }); assert_eq!(expected.to_string(), generated.to_string()); } @@ -998,13 +1056,17 @@ mod tests { } let content = lines.join(""); let expected: TokenStream = content.parse().unwrap(); - let generated = EntityWriter::gen_compact_code_blocks(entity, &crate::WithSerde::None) - .into_iter() - .skip(1) - .fold(TokenStream::new(), |mut acc, tok| { - acc.extend(tok); - acc - }); + let generated = EntityWriter::gen_compact_code_blocks( + entity, + &crate::WithSerde::None, + &crate::DateTimeCrate::Chrono, + ) + .into_iter() + .skip(1) + .fold(TokenStream::new(), |mut acc, tok| { + acc.extend(tok); + acc + }); assert_eq!(expected.to_string(), generated.to_string()); } @@ -1091,7 +1153,7 @@ mod tests { fn assert_serde_variant_results( cake_entity: &Entity, entity_serde_variant: &(String, WithSerde), - generator: Box Vec>, + generator: Box Vec>, ) -> io::Result<()> { let mut reader = BufReader::new(entity_serde_variant.0.as_bytes()); let mut lines: Vec = Vec::new(); @@ -1105,7 +1167,7 @@ mod tests { } let content = lines.join(""); let expected: TokenStream = content.parse().unwrap(); - let generated = generator(cake_entity, &entity_serde_variant.1) + let generated = generator(cake_entity, &entity_serde_variant.1, &DateTimeCrate::Chrono) .into_iter() .fold(TokenStream::new(), |mut acc, tok| { acc.extend(tok);