From 73a57836faefa1c4fdbbee80d77626bcc606332e Mon Sep 17 00:00:00 2001
From: Niklas Korz <niklas@niklaskorz.de>
Date: Fri, 21 Jul 2023 15:45:11 +0200
Subject: [PATCH] sea-orm-codegen: PostgreSQL enum arrays (#1678)

* sea-orm-codegen: PostgreSQL enum arrays

* Refactoring

---------

Co-authored-by: Billy Chan <ccw.billy.123@gmail.com>
---
 sea-orm-codegen/src/entity/column.rs      | 7 +++++++
 sea-orm-codegen/src/entity/transformer.rs | 3 ++-
 sea-orm-codegen/src/entity/writer.rs      | 2 +-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/sea-orm-codegen/src/entity/column.rs b/sea-orm-codegen/src/entity/column.rs
index 50582c27f..f5ba9284b 100644
--- a/sea-orm-codegen/src/entity/column.rs
+++ b/sea-orm-codegen/src/entity/column.rs
@@ -249,6 +249,13 @@ impl Column {
             quote! {}
         }
     }
+
+    pub fn get_inner_col_type(&self) -> &ColumnType {
+        match &self.col_type {
+            ColumnType::Array(inner_col_type) => inner_col_type.as_ref(),
+            _ => &self.col_type,
+        }
+    }
 }
 
 impl From<ColumnDef> for Column {
diff --git a/sea-orm-codegen/src/entity/transformer.rs b/sea-orm-codegen/src/entity/transformer.rs
index 421a83bf1..20f8047cb 100644
--- a/sea-orm-codegen/src/entity/transformer.rs
+++ b/sea-orm-codegen/src/entity/transformer.rs
@@ -58,7 +58,8 @@ impl EntityTransformer {
                     col
                 })
                 .map(|col| {
-                    if let sea_query::ColumnType::Enum { name, variants } = &col.col_type {
+                    if let sea_query::ColumnType::Enum { name, variants } = col.get_inner_col_type()
+                    {
                         enums.insert(
                             name.to_string(),
                             ActiveEnum {
diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
index 36da98910..637e6eb28 100644
--- a/sea-orm-codegen/src/entity/writer.rs
+++ b/sea-orm-codegen/src/entity/writer.rs
@@ -461,7 +461,7 @@ impl EntityWriter {
             .columns
             .iter()
             .fold(TokenStream::new(), |mut ts, col| {
-                if let sea_query::ColumnType::Enum { name, .. } = &col.col_type {
+                if let sea_query::ColumnType::Enum { name, .. } = col.get_inner_col_type() {
                     let enum_name = format_ident!("{}", name.to_string().to_upper_camel_case());
                     ts.extend([quote! {
                         use super::sea_orm_active_enums::#enum_name;