feat: allow support time and chrono

Author: LeadcodeDev

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "sqlx-gen"
-version = "0.4.8"
+version = "0.4.9"
 edition = "2021"
 description = "Generate Rust structs from database schema introspection"
 license = "MIT"

crates/sqlx_gen/Cargo.toml

@@ -87,6 +87,10 @@ pub struct EntitiesArgs {
     #[arg(short = 'v', long)]
     pub views: bool,
 
+    /// Time crate to use for date/time types: chrono (default) or time
+    #[arg(long, default_value = "chrono")]
+    pub time_crate: TimeCrate,
+
     /// Print to stdout without writing files
     #[arg(short = 'n', long)]
     pub dry_run: bool,
@@ -195,6 +199,37 @@ pub enum DatabaseKind {
     Sqlite,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum TimeCrate {
+    #[default]
+    Chrono,
+    Time,
+}
+
+impl std::str::FromStr for TimeCrate {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "chrono" => Ok(Self::Chrono),
+            "time" => Ok(Self::Time),
+            other => Err(format!(
+                "Unknown time crate '{}'. Expected: chrono, time",
+                other
+            )),
+        }
+    }
+}
+
+impl std::fmt::Display for TimeCrate {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Chrono => write!(f, "chrono"),
+            Self::Time => write!(f, "time"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
 pub enum PoolVisibility {
     #[default]
@@ -291,6 +326,7 @@ mod tests {
             tables: None,
             exclude_tables: None,
             views: false,
+            time_crate: TimeCrate::Chrono,
             dry_run: false,
         }
     }

crates/sqlx_gen/src/cli.rs

@@ -4,7 +4,7 @@ use heck::{ToSnakeCase, ToUpperCamelCase};
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 
-use crate::cli::DatabaseKind;
+use crate::cli::{DatabaseKind, TimeCrate};
 use crate::codegen::{imports_for_derives, is_rust_keyword};
 use crate::introspect::{CompositeTypeInfo, SchemaInfo};
 use crate::typemap;
@@ -15,6 +15,7 @@ pub fn generate_composite(
     schema_info: &SchemaInfo,
     extra_derives: &[String],
     type_overrides: &HashMap<String, String>,
+    time_crate: TimeCrate,
 ) -> (TokenStream, BTreeSet<String>) {
     let mut imports = BTreeSet::new();
     for imp in imports_for_derives(extra_derives) {
@@ -51,7 +52,7 @@ pub fn generate_composite(
         .fields
         .iter()
         .map(|col| {
-            let rust_type = typemap::map_column(col, db_kind, schema_info, type_overrides);
+            let rust_type = typemap::map_column(col, db_kind, schema_info, type_overrides, time_crate);
             if let Some(imp) = &rust_type.needs_import {
                 imports.insert(imp.clone());
             }
@@ -131,7 +132,7 @@ mod tests {
 
     fn gen(composite: &CompositeTypeInfo) -> String {
         let schema = SchemaInfo::default();
-        let (tokens, _) = generate_composite(composite, DatabaseKind::Postgres, &schema, &[], &HashMap::new());
+        let (tokens, _) = generate_composite(composite, DatabaseKind::Postgres, &schema, &[], &HashMap::new(), TimeCrate::Chrono);
         parse_and_format(&tokens)
     }
 
@@ -141,7 +142,7 @@ mod tests {
         overrides: &HashMap<String, String>,
     ) -> (String, BTreeSet<String>) {
         let schema = SchemaInfo::default();
-        let (tokens, imports) = generate_composite(composite, DatabaseKind::Postgres, &schema, derives, overrides);
+        let (tokens, imports) = generate_composite(composite, DatabaseKind::Postgres, &schema, derives, overrides, TimeCrate::Chrono);
         (parse_and_format(&tokens), imports)
     }
 
@@ -187,7 +188,7 @@ mod tests {
             fields: vec![make_field("x", "float8", false)],
         };
         let schema = SchemaInfo::default();
-        let (tokens, _) = generate_composite(&c, DatabaseKind::Postgres, &schema, &[], &HashMap::new());
+        let (tokens, _) = generate_composite(&c, DatabaseKind::Postgres, &schema, &[], &HashMap::new(), TimeCrate::Chrono);
         let code = parse_and_format(&tokens);
         assert!(code.contains("sqlx(type_name = \"point\")"));
     }

crates/sqlx_gen/src/codegen/composite_gen.rs

@@ -4,7 +4,7 @@ use heck::ToUpperCamelCase;
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 
-use crate::cli::DatabaseKind;
+use crate::cli::{DatabaseKind, TimeCrate};
 use crate::introspect::{DomainInfo, SchemaInfo};
 use crate::typemap;
 
@@ -13,6 +13,7 @@ pub fn generate_domain(
     db_kind: DatabaseKind,
     schema_info: &SchemaInfo,
     type_overrides: &HashMap<String, String>,
+    time_crate: TimeCrate,
 ) -> (TokenStream, BTreeSet<String>) {
     let mut imports = BTreeSet::new();
     let alias_name = format_ident!("{}", domain.name.to_upper_camel_case());
@@ -34,7 +35,7 @@ pub fn generate_domain(
         column_default: None,
     };
 
-    let rust_type = typemap::map_column(&fake_col, db_kind, schema_info, type_overrides);
+    let rust_type = typemap::map_column(&fake_col, db_kind, schema_info, type_overrides, time_crate);
     if let Some(imp) = &rust_type.needs_import {
         imports.insert(imp.clone());
     }
@@ -69,13 +70,13 @@ mod tests {
 
     fn gen(domain: &DomainInfo) -> (String, BTreeSet<String>) {
         let schema = SchemaInfo::default();
-        let (tokens, imports) = generate_domain(domain, DatabaseKind::Postgres, &schema, &HashMap::new());
+        let (tokens, imports) = generate_domain(domain, DatabaseKind::Postgres, &schema, &HashMap::new(), TimeCrate::Chrono);
         (parse_and_format(&tokens), imports)
     }
 
     fn gen_with_overrides(domain: &DomainInfo, overrides: &HashMap<String, String>) -> (String, BTreeSet<String>) {
         let schema = SchemaInfo::default();
-        let (tokens, imports) = generate_domain(domain, DatabaseKind::Postgres, &schema, overrides);
+        let (tokens, imports) = generate_domain(domain, DatabaseKind::Postgres, &schema, overrides, TimeCrate::Chrono);
         (parse_and_format(&tokens), imports)
     }