From 83e48597764a0efb477a61e3f9a751e20acb6ce8 Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Tue, 22 Jun 2021 22:40:43 +0800
Subject: [PATCH] codegen ColumnDef null() and unique()

---
 examples/codegen/src/out/fruit.rs    |  2 +-
 sea-orm-codegen/src/entity/column.rs | 24 +++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/examples/codegen/src/out/fruit.rs b/examples/codegen/src/out/fruit.rs
index 6e190749..17f653ed 100644
--- a/examples/codegen/src/out/fruit.rs
+++ b/examples/codegen/src/out/fruit.rs
@@ -47,7 +47,7 @@ impl ColumnTrait for Column {
         match self {
             Self::Id => ColumnType::Integer.def(),
             Self::Name => ColumnType::String(Some(255u32)).def(),
-            Self::CakeId => ColumnType::Integer.def(),
+            Self::CakeId => ColumnType::Integer.def().null(),
         }
     }
 }
diff --git a/sea-orm-codegen/src/entity/column.rs b/sea-orm-codegen/src/entity/column.rs
index b575057f..e8a0b543 100644
--- a/sea-orm-codegen/src/entity/column.rs
+++ b/sea-orm-codegen/src/entity/column.rs
@@ -9,6 +9,7 @@ pub struct Column {
     pub(crate) col_type: ColumnType,
     pub(crate) auto_increment: bool,
     pub(crate) not_null: bool,
+    pub(crate) unique: bool,
 }
 
 impl Column {
@@ -50,7 +51,7 @@ impl Column {
     }
 
     pub fn get_def(&self) -> TokenStream {
-        match &self.col_type {
+        let mut col_def = match &self.col_type {
             ColumnType::Char(s) => match s {
                 Some(s) => quote! { ColumnType::Char(Some(#s)).def() },
                 None => quote! { ColumnType::Char(None).def() },
@@ -86,7 +87,18 @@ impl Column {
                 let s = s.to_string();
                 quote! { ColumnType::Custom(#s.to_owned()).def() }
             }
+        };
+        if !self.not_null {
+            col_def.extend(quote! {
+                .null()
+            });
         }
+        if self.unique {
+            col_def.extend(quote! {
+                .unique()
+            });
+        }
+        col_def
     }
 }
 
@@ -115,11 +127,21 @@ impl From<&ColumnDef> for Column {
             })
             .collect();
         let not_null = !not_nulls.is_empty();
+        let uniques: Vec<bool> = col_def
+            .get_column_spec()
+            .iter()
+            .filter_map(|spec| match spec {
+                ColumnSpec::UniqueKey => Some(true),
+                _ => None,
+            })
+            .collect();
+        let unique = !uniques.is_empty();
         Self {
             name,
             col_type,
             auto_increment,
             not_null,
+            unique,
         }
     }
 }