#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

@@ -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,
+}

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<DateTimeCrate> 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;

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<TokenStream> {
+    pub fn get_column_rs_types(&self, date_time_crate: &DateTimeCrate) -> Vec<TokenStream> {
         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::<Vec<_>>();
@@ -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()
         );
     }
 

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<u8>".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<u8>",
+            "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()
             );
         }

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<OutputFile> {
+    pub fn write_entities(&self, context: &EntityWriterContext) -> Vec<OutputFile> {
         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<TokenStream> {
+    pub fn gen_expanded_code_blocks(
+        entity: &Entity,
+        with_serde: &WithSerde,
+        date_time_crate: &DateTimeCrate,
+    ) -> Vec<TokenStream> {
         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<TokenStream> {
+    pub fn gen_compact_code_blocks(
+        entity: &Entity,
+        with_serde: &WithSerde,
+        date_time_crate: &DateTimeCrate,
+    ) -> Vec<TokenStream> {
         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<String> = 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<dyn Fn(&Entity, &WithSerde) -> Vec<TokenStream>>,
+        generator: Box<dyn Fn(&Entity, &WithSerde, &DateTimeCrate) -> Vec<TokenStream>>,
     ) -> io::Result<()> {
         let mut reader = BufReader::new(entity_serde_variant.0.as_bytes());
         let mut lines: Vec<String> = 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);