From 39c5a4d134c58aba88f215c78143cdb7a740818b Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Mon, 13 Sep 2021 18:33:07 +0800
Subject: [PATCH] Derive attributes

---
 sea-orm-codegen/src/entity/column.rs | 15 ++++++++++++++
 sea-orm-codegen/src/entity/writer.rs | 29 +++++++++++++++++++++++-----
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/sea-orm-codegen/src/entity/column.rs b/sea-orm-codegen/src/entity/column.rs
index 0b9eb8de..fd8e1936 100644
--- a/sea-orm-codegen/src/entity/column.rs
+++ b/sea-orm-codegen/src/entity/column.rs
@@ -2,6 +2,7 @@ use heck::{CamelCase, SnakeCase};
 use proc_macro2::{Ident, TokenStream};
 use quote::{format_ident, quote};
 use sea_query::{ColumnDef, ColumnSpec, ColumnType};
+use syn::punctuated::Punctuated;
 
 #[derive(Clone, Debug)]
 pub struct Column {
@@ -52,6 +53,20 @@ impl Column {
         }
     }
 
+    pub fn get_col_type_attrs(&self) -> Option<TokenStream> {
+        let col_type = match &self.col_type {
+            ColumnType::Float(Some(l)) => Some(format!("Float(Some({}))", l)),
+            ColumnType::Double(Some(l)) => Some(format!("Double(Some({}))", l)),
+            ColumnType::Decimal(Some((p, s))) => Some(format!("Decimal(Some(({}, {})))", p, s)),
+            ColumnType::Money(Some((p, s))) => Some(format!("Money(Some({}, {}))", p, s)),
+            ColumnType::Custom(iden) => {
+                Some(format!("Custom(\"{}\".to_owned())", iden.to_string()))
+            }
+            _ => None,
+        };
+        col_type.map(|ty| quote! { column_type = #ty })
+    }
+
     pub fn get_def(&self) -> TokenStream {
         let mut col_def = match &self.col_type {
             ColumnType::Char(s) => match s {
diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
index 210d43a0..5a47ce4c 100644
--- a/sea-orm-codegen/src/entity/writer.rs
+++ b/sea-orm-codegen/src/entity/writer.rs
@@ -1,6 +1,8 @@
 use crate::Entity;
 use proc_macro2::TokenStream;
 use quote::quote;
+use std::iter::FromIterator;
+use syn::{punctuated::Punctuated, token::Comma};
 
 #[derive(Clone, Debug)]
 pub struct EntityWriter {
@@ -332,13 +334,30 @@ impl EntityWriter {
             .columns
             .iter()
             .map(|col| {
-                if !primary_keys.contains(&col.name) {
-                    TokenStream::new()
-                } else {
-                    quote! {
-                        #[sea_orm(primary_key)]
+                let mut attrs: Punctuated<_, Comma> = Punctuated::new();
+                if primary_keys.contains(&col.name) {
+                    attrs.push(quote! { primary_key });
+                    if !col.auto_increment {
+                        attrs.push(quote! { auto_increment = false });
                     }
                 }
+                if let Some(ts) = col.get_col_type_attrs() {
+                    attrs.extend(vec![ts]);
+                };
+                if !attrs.is_empty() {
+                    let mut ts = TokenStream::new();
+                    for (i, attr) in attrs.into_iter().enumerate() {
+                        if i > 0 {
+                            ts = quote! { #ts, };
+                        }
+                        ts = quote! { #ts #attr };
+                    }
+                    quote! {
+                        #[sea_orm(#ts)]
+                    }
+                } else {
+                    TokenStream::new()
+                }
             })
             .collect();
         quote! {