From d7fc102338d76cfd76d3dbb236ec3d5a3573be29 Mon Sep 17 00:00:00 2001
From: Martin Norbury <martin.norbury@gmail.com>
Date: Wed, 30 Aug 2023 11:47:43 +0100
Subject: [PATCH] fix: check for extension type before built-in arrow types

The JSON type was not getting selected as the underlying storage type (Binary) is
the underlying arrow type.

Checking for extension type first will allow us to prefer our defined Arrow extensions over
the built-in types.
---
 .../main/java/io/cloudquery/memdb/MemDB.java  |  4 +++
 .../java/io/cloudquery/memdb/Table1Data.java  |  4 +++
 .../java/io/cloudquery/scalar/Scalar.java     | 32 +++++++++----------
 .../java/io/cloudquery/scalar/ScalarTest.java |  2 ++
 4 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/lib/src/main/java/io/cloudquery/memdb/MemDB.java b/lib/src/main/java/io/cloudquery/memdb/MemDB.java
index 26a100e..07f50e4 100644
--- a/lib/src/main/java/io/cloudquery/memdb/MemDB.java
+++ b/lib/src/main/java/io/cloudquery/memdb/MemDB.java
@@ -15,7 +15,9 @@
 import io.cloudquery.transformers.Tables;
 import io.cloudquery.transformers.TransformWithClass;
 import io.grpc.stub.StreamObserver;
+import java.time.LocalDateTime;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
 public class MemDB extends Plugin {
@@ -32,6 +34,8 @@ public void resolve(
                         Table1Data.builder()
                             .id(UUID.fromString("46b2b6e6-8f3e-4340-a721-4aa0786b1cc0"))
                             .name("name1")
+                            .timestamp(LocalDateTime.now())
+                            .json(Map.of("key1", "value1", "key2", "value2"))
                             .build());
                     stream.write(
                         Table1Data.builder()
diff --git a/lib/src/main/java/io/cloudquery/memdb/Table1Data.java b/lib/src/main/java/io/cloudquery/memdb/Table1Data.java
index f4ba5d0..786aff1 100644
--- a/lib/src/main/java/io/cloudquery/memdb/Table1Data.java
+++ b/lib/src/main/java/io/cloudquery/memdb/Table1Data.java
@@ -1,5 +1,7 @@
 package io.cloudquery.memdb;
 
+import java.time.LocalDateTime;
+import java.util.Map;
 import java.util.UUID;
 import lombok.Builder;
 import lombok.Getter;
@@ -9,4 +11,6 @@
 public class Table1Data {
   private UUID id;
   private String name;
+  private LocalDateTime timestamp;
+  private Map<String, String> json;
 }
diff --git a/lib/src/main/java/io/cloudquery/scalar/Scalar.java b/lib/src/main/java/io/cloudquery/scalar/Scalar.java
index c2680fd..ae9a775 100644
--- a/lib/src/main/java/io/cloudquery/scalar/Scalar.java
+++ b/lib/src/main/java/io/cloudquery/scalar/Scalar.java
@@ -1,5 +1,6 @@
 package io.cloudquery.scalar;
 
+import io.cloudquery.types.JSONType;
 import io.cloudquery.types.UUIDType;
 import java.util.Objects;
 import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -77,6 +78,21 @@ public final int hashCode() {
   public static final java.lang.String NULL_VALUE_STRING = "(null)";
 
   public static Scalar<?> fromArrowType(ArrowType arrowType) {
+    if (arrowType instanceof ArrowType.ExtensionType extensionType) {
+      switch (extensionType.extensionName()) {
+        case UUIDType.EXTENSION_NAME -> {
+          return new UUID();
+        }
+        case JSONType.EXTENSION_NAME -> {
+          return new JSON();
+        }
+          // TODO: Add support for these types when scalar available
+          // case INETType.EXTENSION_NAME -> {
+          //     return new INET();
+          // }
+      }
+    }
+
     switch (arrowType.getTypeID()) {
       case Timestamp -> {
         return new Timestamp();
@@ -110,22 +126,6 @@ public static Scalar<?> fromArrowType(ArrowType arrowType) {
       }
     }
 
-    if (arrowType instanceof ArrowType.ExtensionType extensionType) {
-      //noinspection SwitchStatementWithTooFewBranches
-      switch (extensionType.extensionName()) {
-        case UUIDType.EXTENSION_NAME -> {
-          return new UUID();
-        }
-          // TODO: Add support for these types when scalar available
-          // case JSONType.EXTENSION_NAME -> {
-          //     return new JSON();
-          // }
-          // case INETType.EXTENSION_NAME -> {
-          //     return new INET();
-          // }
-      }
-    }
-
     throw new UnsupportedOperationException("Unsupported type: " + arrowType);
   }
 
diff --git a/lib/src/test/java/io/cloudquery/scalar/ScalarTest.java b/lib/src/test/java/io/cloudquery/scalar/ScalarTest.java
index 3b21bf9..db5c9d8 100644
--- a/lib/src/test/java/io/cloudquery/scalar/ScalarTest.java
+++ b/lib/src/test/java/io/cloudquery/scalar/ScalarTest.java
@@ -2,6 +2,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
+import io.cloudquery.types.JSONType;
 import io.cloudquery.types.UUIDType;
 import java.time.ZoneOffset;
 import java.util.stream.Stream;
@@ -64,6 +65,7 @@ public static Stream<Arguments> testDataSource() {
 
         // Extension
         Arguments.of(new UUIDType(), UUID.class),
+        Arguments.of(new JSONType(), JSON.class),
 
         // Date
         Arguments.of(new ArrowType.Date(DateUnit.DAY), DateDay.class),