codegen: Fix SQLite integer and primary key generation

sea-orm-cli/src/cli.rs

@@ -330,6 +330,14 @@ pub enum GenerateSubcommands {
             long_help = "Generate empty ActiveModelBehavior impls."
         )]
         impl_active_model_behavior: bool,
+
+        #[arg(
+            long,
+            default_value_t,
+            value_enum,
+            help = "The Rust type of SQLite integers."
+        )]
+        sqlite_int_rs_type: SqliteIntRsType,
     },
 }
 
@@ -340,6 +348,22 @@ pub enum DateTimeCrate {
     Time,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum, Default)]
+pub enum SqliteIntRsType {
+    #[default]
+    I64,
+    I32,
+}
+
+impl From<SqliteIntRsType> for sea_orm_codegen::SqliteIntRsType {
+    fn from(value: SqliteIntRsType) -> Self {
+        match value {
+            SqliteIntRsType::I64 => sea_orm_codegen::SqliteIntRsType::I64,
+            SqliteIntRsType::I32 => sea_orm_codegen::SqliteIntRsType::I32,
+        }
+    }
+}
+
 /// Use this to build a local, version-controlled `sea-orm-cli` in dependent projects
 /// (see [example use case](https://github.com/SeaQL/sea-orm/discussions/1889)).
 #[cfg(feature = "codegen")]

sea-orm-cli/src/commands/generate.rs

@@ -1,7 +1,7 @@
 use core::time;
 use sea_orm_codegen::{
-    DateTimeCrate as CodegenDateTimeCrate, EntityTransformer, EntityWriterContext, OutputFile,
-    WithPrelude, WithSerde,
+    DatabaseBackend, DateTimeCrate as CodegenDateTimeCrate, EntityTransformer, EntityWriterContext,
+    OutputFile, WithPrelude, WithSerde,
 };
 use std::{error::Error, fs, io::Write, path::Path, process::Command, str::FromStr};
 use tracing_subscriber::{prelude::*, EnvFilter};
@@ -38,6 +38,7 @@ pub async fn run_generate_command(
             enum_extra_attributes,
             seaography,
             impl_active_model_behavior,
+            sqlite_int_rs_type,
         } => {
             if verbose {
                 let _ = tracing_subscriber::fmt()
@@ -111,7 +112,7 @@ pub async fn run_generate_command(
                 Default::default()
             };
 
-            let (schema_name, table_stmts) = match url.scheme() {
+            let (schema_name, table_stmts, backend) = match url.scheme() {
                 "mysql" => {
                     use sea_schema::mysql::discovery::SchemaDiscovery;
                     use sqlx::MySql;
@@ -132,7 +133,7 @@ pub async fn run_generate_command(
                         .filter(|schema| filter_skip_tables(&schema.info.name))
                         .map(|schema| schema.write())
                         .collect();
-                    (None, table_stmts)
+                    (None, table_stmts, DatabaseBackend::MySql)
                 }
                 "sqlite" => {
                     use sea_schema::sqlite::discovery::SchemaDiscovery;
@@ -161,7 +162,7 @@ pub async fn run_generate_command(
                         .filter(|schema| filter_skip_tables(&schema.name))
                         .map(|schema| schema.write())
                         .collect();
-                    (None, table_stmts)
+                    (None, table_stmts, DatabaseBackend::Sqlite)
                 }
                 "postgres" | "postgresql" => {
                     use sea_schema::postgres::discovery::SchemaDiscovery;
@@ -187,7 +188,7 @@ pub async fn run_generate_command(
                         .filter(|schema| filter_skip_tables(&schema.info.name))
                         .map(|schema| schema.write())
                         .collect();
-                    (database_schema, table_stmts)
+                    (database_schema, table_stmts, DatabaseBackend::Postgres)
                 }
                 _ => unimplemented!("{} is not supported", url.scheme()),
             };
@@ -209,6 +210,8 @@ pub async fn run_generate_command(
                 enum_extra_attributes,
                 seaography,
                 impl_active_model_behavior,
+                backend,
+                sqlite_int_rs_type.into(),
             );
             let output = EntityTransformer::transform(table_stmts)?.generate(&writer_context);
 

sea-orm-codegen/src/entity/base_entity.rs

@@ -4,9 +4,9 @@ use quote::format_ident;
 use quote::quote;
 use sea_query::ColumnType;
 
-use crate::{
-    util::escape_rust_keyword, Column, ConjunctRelation, DateTimeCrate, PrimaryKey, Relation,
-};
+use crate::{util::escape_rust_keyword, Column, ConjunctRelation, PrimaryKey, Relation};
+
+use super::EntityWriterContext;
 
 #[derive(Clone, Debug)]
 pub struct Entity {
@@ -48,11 +48,11 @@ impl Entity {
             .collect()
     }
 
-    pub fn get_column_rs_types(&self, date_time_crate: &DateTimeCrate) -> Vec<TokenStream> {
+    pub fn get_column_rs_types(&self, context: &EntityWriterContext) -> Vec<TokenStream> {
         self.columns
             .clone()
             .into_iter()
-            .map(|col| col.get_rs_type(date_time_crate))
+            .map(|col| col.get_rs_type(context))
             .collect()
     }
 
@@ -183,7 +183,7 @@ impl Entity {
         format_ident!("{}", auto_increment)
     }
 
-    pub fn get_primary_key_rs_type(&self, date_time_crate: &DateTimeCrate) -> TokenStream {
+    pub fn get_primary_key_rs_type(&self, context: &EntityWriterContext) -> TokenStream {
         let types = self
             .primary_keys
             .iter()
@@ -192,7 +192,7 @@ impl Entity {
                     .iter()
                     .find(|col| col.name.eq(&primary_key.name))
                     .unwrap()
-                    .get_rs_type(date_time_crate)
+                    .get_rs_type(context)
                     .to_string()
             })
             .collect::<Vec<_>>();
@@ -270,7 +270,9 @@ mod tests {
     use quote::{format_ident, quote};
     use sea_query::{ColumnType, ForeignKeyAction, StringLen};
 
-    use crate::{Column, DateTimeCrate, Entity, PrimaryKey, Relation, RelationType};
+    use crate::{
+        Column, DatabaseBackend, Entity, EntityWriterContext, PrimaryKey, Relation, RelationType,
+    };
 
     fn setup() -> Entity {
         Entity {
@@ -282,13 +284,15 @@ mod tests {
                     auto_increment: false,
                     not_null: false,
                     unique: false,
+                    primary_key: false,
                 },
                 Column {
                     name: "name".to_owned(),
                     col_type: ColumnType::String(StringLen::None),
                     auto_increment: false,
                     not_null: false,
                     unique: false,
+                    primary_key: false,
                 },
             ],
             relations: vec![
@@ -381,16 +385,15 @@ mod tests {
     fn test_get_column_rs_types() {
         let entity = setup();
 
-        for (i, elem) in entity
-            .get_column_rs_types(&DateTimeCrate::Chrono)
-            .into_iter()
-            .enumerate()
-        {
+        let context = EntityWriterContext {
+            db_backend: DatabaseBackend::Postgres,
+            ..Default::default()
+        };
+
+        for (i, elem) in entity.get_column_rs_types(&context).into_iter().enumerate() {
             assert_eq!(
                 elem.to_string(),
-                entity.columns[i]
-                    .get_rs_type(&DateTimeCrate::Chrono)
-                    .to_string()
+                entity.columns[i].get_rs_type(&context).to_string()
             );
         }
     }
@@ -489,13 +492,33 @@ mod tests {
     fn test_get_primary_key_rs_type() {
         let entity = setup();
 
+        let context = EntityWriterContext::default();
+
         assert_eq!(
-            entity
-                .get_primary_key_rs_type(&DateTimeCrate::Chrono)
-                .to_string(),
-            entity.columns[0]
-                .get_rs_type(&DateTimeCrate::Chrono)
-                .to_string()
+            entity.get_primary_key_rs_type(&context).to_string(),
+            entity.columns[0].get_rs_type(&context).to_string()
+        );
+    }
+
+    #[test]
+    fn test_get_sqlite_int_rs_type() {
+        let context = EntityWriterContext {
+            db_backend: DatabaseBackend::Sqlite,
+            ..Default::default()
+        };
+
+        assert_eq!(
+            "i64",
+            Column {
+                name: "id".to_owned(),
+                col_type: ColumnType::Integer,
+                auto_increment: true,
+                not_null: true,
+                unique: true,
+                primary_key: false,
+            }
+            .get_rs_type(&context)
+            .to_string()
         );
     }