fix: check for extension type before built-in arrow types (#109)

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.

Author: mnorbury

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()

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;
 }

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);
   }
 

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),