Custom type for table primary key

Cargo.toml

@@ -34,7 +34,7 @@ tracing = { version = "0.1", default-features = false, features = ["attributes",
 rust_decimal = { version = "1", default-features = false, optional = true }
 bigdecimal = { version = "0.3", default-features = false, optional = true }
 sea-orm-macros = { version = "0.12.0-rc.4", path = "sea-orm-macros", default-features = false, features = ["strum"] }
-sea-query = { version = "0.29.0", default-features = false, features = ["thread-safe", "hashable-value", "backend-mysql", "backend-postgres", "backend-sqlite"] }
+sea-query = { path = "../sea-query", default-features = false, features = ["thread-safe", "hashable-value", "backend-mysql", "backend-postgres", "backend-sqlite"] }
 sea-query-binder = { version = "0.4.0-rc.2", default-features = false, optional = true }
 strum = { version = "0.25", default-features = false }
 serde = { version = "1.0", default-features = false }

sea-orm-cli/src/cli.rs

@@ -237,6 +237,13 @@ pub enum GenerateSubcommands {
         )]
         serde_skip_hidden_column: bool,
 
+        #[arg(
+            long,
+            default_value = "false",
+            help = "Custom Rust type for pk"
+        )]
+        custom_rust_type_pk: bool,
+
         #[arg(
             long,
             default_value = "false",

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

@@ -1,6 +1,6 @@
 use sea_orm_codegen::{
     DateTimeCrate as CodegenDateTimeCrate, EntityTransformer, EntityWriterContext, OutputFile,
-    WithSerde,
+    WithSerde, WithTablePkType,
 };
 use std::{error::Error, fs, io::Write, path::Path, process::Command, str::FromStr};
 use tracing_subscriber::{prelude::*, EnvFilter};
@@ -26,6 +26,7 @@ pub async fn run_generate_command(
             with_serde,
             serde_skip_deserializing_primary_key,
             serde_skip_hidden_column,
+            custom_rust_type_pk: bool,
             with_copy_enums,
             date_time_crate,
             lib,
@@ -165,6 +166,10 @@ pub async fn run_generate_command(
             let writer_context = EntityWriterContext::new(
                 expanded_format,
                 WithSerde::from_str(&with_serde).expect("Invalid serde derive option"),
+                match custom_rust_type_pk {
+                    true => WithTablePkType::Custom,
+                    false => WithTablePkType::None,
+                },
                 with_copy_enums,
                 date_time_crate.into(),
                 schema_name,

sea-orm-codegen/Cargo.toml

@@ -17,7 +17,7 @@ name = "sea_orm_codegen"
 path = "src/lib.rs"
 
 [dependencies]
-sea-query = { version = "0.29.0-rc.2", default-features = false, features = ["thread-safe"] }
+sea-query = { path = "../../sea-query", default-features = false, features = ["thread-safe"] }
 syn = { version = "2", default-features = false, features = ["parsing", "proc-macro", "derive", "printing"] }
 quote = { version = "1", default-features = false }
 heck = { version = "0.4", default-features = false }

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

@@ -4,6 +4,7 @@ use quote::format_ident;
 use quote::quote;
 use sea_query::ColumnType;
 
+use crate::WithTablePkType;
 use crate::{
     util::escape_rust_keyword, Column, ConjunctRelation, DateTimeCrate, PrimaryKey, Relation,
 };
@@ -48,11 +49,32 @@ impl Entity {
             .collect()
     }
 
-    pub fn get_column_rs_types(&self, date_time_crate: &DateTimeCrate) -> Vec<TokenStream> {
+    pub fn get_column_rs_types(
+        &self,
+        date_time_crate: &DateTimeCrate,
+        with_table_pk_type: &WithTablePkType,
+    ) -> Vec<TokenStream> {
         self.columns
             .clone()
             .into_iter()
-            .map(|col| col.get_rs_type(date_time_crate))
+            .map(|col| {
+                // If this is the primary key column, we need to check if we need
+                // to use the custom type for the table's primary key
+                if let ColumnType::CustomRustType { rust_ty, db_ty } = col.col_type {
+                    match with_table_pk_type {
+                        WithTablePkType::None => Column {
+                            col_type: *db_ty,
+                            ..col
+                        }
+                        .get_rs_type(date_time_crate),
+                        WithTablePkType::Custom => {
+                            quote! { #rust_ty }
+                        }
+                    }
+                } else {
+                    col.get_rs_type(date_time_crate)
+                }
+            })
             .collect()
     }
 
@@ -381,7 +403,7 @@ mod tests {
         let entity = setup();
 
         for (i, elem) in entity
-            .get_column_rs_types(&DateTimeCrate::Chrono)
+            .get_column_rs_types(&DateTimeCrate::Chrono, &crate::WithTablePkType::None)
             .into_iter()
             .enumerate()
         {

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

@@ -75,6 +75,7 @@ impl Column {
                 ColumnType::Array(column_type) => {
                     format!("Vec<{}>", write_rs_type(column_type, date_time_crate))
                 }
+                ColumnType::CustomRustType { rust_ty, db_ty } => rust_ty.to_owned(),
                 _ => unimplemented!(),
             }
         }

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

@@ -2,18 +2,21 @@ use crate::{
     util::unpack_table_ref, ActiveEnum, Column, ConjunctRelation, Entity, EntityWriter, Error,
     PrimaryKey, Relation, RelationType,
 };
-use sea_query::{ColumnSpec, TableCreateStatement};
+use heck::ToUpperCamelCase;
+use sea_query::{Alias, ColumnDef, ColumnSpec, SeaRc, TableCreateStatement};
 use std::collections::{BTreeMap, HashMap};
 
 #[derive(Clone, Debug)]
 pub struct EntityTransformer;
 
 impl EntityTransformer {
-    pub fn transform(table_create_stmts: Vec<TableCreateStatement>) -> Result<EntityWriter, Error> {
+    pub fn transform(
+        mut table_create_stmts: Vec<TableCreateStatement>,
+    ) -> Result<EntityWriter, Error> {
         let mut enums: BTreeMap<String, ActiveEnum> = BTreeMap::new();
         let mut inverse_relations: BTreeMap<String, Vec<Relation>> = BTreeMap::new();
         let mut entities = BTreeMap::new();
-        for table_create in table_create_stmts.into_iter() {
+        for table_create in table_create_stmts.iter_mut() {
             let table_name = match table_create.get_table_name() {
                 Some(table_ref) => match table_ref {
                     sea_query::TableRef::Table(t)
@@ -43,10 +46,37 @@ impl EntityTransformer {
                         primary_keys.push(PrimaryKey {
                             name: col_def.get_column_name(),
                         });
+
+                        // // Change this to a custom type
+                        // let curr_type = col_def.get_column_type();
+
+                        // let pk_custom_type_name = format!("{}PrimaryKey", table_name);
+
+                        // let new_col_type = sea_query::ColumnType::CustomRustType {
+                        //     rust_ty: todo!(),
+                        //     db_ty: todo!(),
+                        // };
+                        // let col_def = ColumnDef::new_with_type(
+                        //     SeaRc::new(Alias::new(col_def.get_column_name())),
+                        //     new_col_type,
+                        // );
                     }
-                    col_def.into()
+                    let col_def: Column = col_def.into();
+                    (col_def, primary_key)
                 })
-                .map(|mut col: Column| {
+                .map(|(mut col, primary_key)| {
+                    if primary_key {
+                        // Change this to a custom type
+                        let curr_type = col.col_type;
+
+                        let pk_custom_type_name = format!("{}PrimaryKey", table_name);
+
+                        col.col_type = sea_query::ColumnType::CustomRustType {
+                            rust_ty: pk_custom_type_name.to_upper_camel_case(),
+                            db_ty: Box::new(curr_type),
+                        };
+                    }
+
                     col.unique = table_create
                         .get_indexes()
                         .iter()
@@ -379,6 +409,7 @@ mod tests {
                 EntityWriter::gen_compact_code_blocks(
                     entity,
                     &crate::WithSerde::None,
+                    &crate::WithTablePkType::None,
                     &crate::DateTimeCrate::Chrono,
                     &None,
                     false,