From 5fa2c1c28ced52cfdd2574c30aae3b37b5ed3026 Mon Sep 17 00:00:00 2001
From: Anshul Sanghi <anshul@anshulsanghi.tech>
Date: Wed, 2 Aug 2023 00:12:16 +0530
Subject: [PATCH] Add Support For PostgreSQL Arrays In FromQueryResult
 Implementation Of JsonValue (#1598)

* Add Support For PostgreSQL Arrays In FromQueryResult Implementation Of JsonValue

* Add support for root arrays in JSON in SeaORM #1517

* Refactoring

* Only when `postgres-array` is enabled

* Add test cases

---------

Co-authored-by: Billy Chan <ccw.billy.123@gmail.com>
---
 src/query/json.rs         | 23 ++++++++++++++++--
 tests/collection_tests.rs | 49 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/src/query/json.rs b/src/query/json.rs
index 066f69c24..66b4dc978 100644
--- a/src/query/json.rs
+++ b/src/query/json.rs
@@ -77,6 +77,7 @@ impl FromQueryResult for JsonValue {
             crate::QueryResultRow::SqlxPostgres(row) => {
                 use serde_json::json;
                 use sqlx::{postgres::types::Oid, Column, Postgres, Row, Type};
+
                 for column in row.columns() {
                     let col = if !column.name().starts_with(pre) {
                         continue;
@@ -84,13 +85,25 @@ impl FromQueryResult for JsonValue {
                         column.name().replacen(pre, "", 1)
                     };
                     let col_type = column.type_info();
+
                     macro_rules! match_postgres_type {
                         ( $type: ty ) => {
-                            if <$type as Type<Postgres>>::type_info().eq(col_type) {
-                                try_get_type!($type, col)
+                            match col_type.kind() {
+                                #[cfg(feature = "postgres-array")]
+                                sqlx::postgres::PgTypeKind::Array(_) => {
+                                    if <Vec<$type> as Type<Postgres>>::type_info().eq(col_type) {
+                                        try_get_type!(Vec<$type>, col);
+                                    }
+                                }
+                                _ => {
+                                    if <$type as Type<Postgres>>::type_info().eq(col_type) {
+                                        try_get_type!($type, col);
+                                    }
+                                }
                             }
                         };
                     }
+
                     match_postgres_type!(bool);
                     match_postgres_type!(i8);
                     match_postgres_type!(i16);
@@ -126,9 +139,15 @@ impl FromQueryResult for JsonValue {
                     match_postgres_type!(rust_decimal::Decimal);
                     #[cfg(feature = "with-json")]
                     try_get_type!(serde_json::Value, col);
+                    #[cfg(all(feature = "with-json", feature = "postgres-array"))]
+                    try_get_type!(Vec<serde_json::Value>, col);
                     try_get_type!(String, col);
+                    #[cfg(feature = "postgres-array")]
+                    try_get_type!(Vec<String>, col);
                     #[cfg(feature = "with-uuid")]
                     try_get_type!(uuid::Uuid, col);
+                    #[cfg(all(feature = "with-uuid", feature = "postgres-array"))]
+                    try_get_type!(Vec<uuid::Uuid>, col);
                     try_get_type!(Vec<u8>, col);
                 }
                 Ok(JsonValue::Object(map))
diff --git a/tests/collection_tests.rs b/tests/collection_tests.rs
index 2447bee9f..8bd0eafb6 100644
--- a/tests/collection_tests.rs
+++ b/tests/collection_tests.rs
@@ -5,6 +5,7 @@ use pretty_assertions::assert_eq;
 use sea_orm::{
     entity::prelude::*, entity::*, DatabaseConnection, DerivePartialModel, FromQueryResult,
 };
+use serde_json::json;
 
 #[sea_orm_macros::test]
 #[cfg(all(feature = "sqlx-postgres", feature = "postgres-array"))]
@@ -114,6 +115,54 @@ pub async fn insert_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
         }
     );
 
+    assert_eq!(
+        Entity::find_by_id(1).into_json().one(db).await?,
+        Some(json!({
+            "id": 1,
+            "name": "Collection 1",
+            "integers": [1, 2, 3],
+            "integers_opt": [1, 2, 3],
+            "teas": ["BreakfastTea"],
+            "teas_opt": ["BreakfastTea"],
+            "colors": [0],
+            "colors_opt": [0],
+            "uuid": [uuid],
+            "uuid_hyphenated": [uuid.hyphenated()],
+        }))
+    );
+
+    assert_eq!(
+        Entity::find_by_id(2).into_json().one(db).await?,
+        Some(json!({
+            "id": 2,
+            "name": "Collection 2",
+            "integers": [10, 9],
+            "integers_opt": null,
+            "teas": ["BreakfastTea"],
+            "teas_opt": null,
+            "colors": [0],
+            "colors_opt": null,
+            "uuid": [uuid],
+            "uuid_hyphenated": [uuid.hyphenated()],
+        }))
+    );
+
+    assert_eq!(
+        Entity::find_by_id(3).into_json().one(db).await?,
+        Some(json!({
+            "id": 3,
+            "name": "Collection 3",
+            "integers": [],
+            "integers_opt": [],
+            "teas": [],
+            "teas_opt": [],
+            "colors": [],
+            "colors_opt": [],
+            "uuid": [uuid],
+            "uuid_hyphenated": [uuid.hyphenated()],
+        }))
+    );
+
     Ok(())
 }