From 12c22d54a4dc6060213a0b4708d8c6e0d30b2c2a Mon Sep 17 00:00:00 2001
From: candiduslynx <candiduslynx@gmail.com>
Date: Mon, 27 Nov 2023 20:55:57 +0200
Subject: [PATCH] fix(test): Register extensions in tests

---
 .../io/cloudquery/helper/ArrowHelper.java     | 23 +------------------
 .../io/cloudquery/helper/ArrowHelperTest.java |  9 ++++++++
 2 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/lib/src/main/java/io/cloudquery/helper/ArrowHelper.java b/lib/src/main/java/io/cloudquery/helper/ArrowHelper.java
index be21059..10661df 100644
--- a/lib/src/main/java/io/cloudquery/helper/ArrowHelper.java
+++ b/lib/src/main/java/io/cloudquery/helper/ArrowHelper.java
@@ -8,9 +8,7 @@
 import io.cloudquery.schema.Resource;
 import io.cloudquery.schema.Table;
 import io.cloudquery.schema.Table.TableBuilder;
-import io.cloudquery.types.JSONType;
 import io.cloudquery.types.JSONType.JSONVector;
-import io.cloudquery.types.UUIDType;
 import io.cloudquery.types.UUIDType.UUIDVector;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -23,7 +21,6 @@
 import org.apache.arrow.vector.ipc.ArrowReader;
 import org.apache.arrow.vector.ipc.ArrowStreamReader;
 import org.apache.arrow.vector.ipc.ArrowStreamWriter;
-import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -264,30 +261,12 @@ private static Column getColumn(Field field) {
     boolean isIncrementalKey =
         Objects.equals(field.getMetadata().get(CQ_EXTENSION_INCREMENTAL), "true");
 
-    ArrowType fieldType = field.getType();
-    String extensionName =
-        field.getMetadata().get(ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME);
-    String extensionMetadata =
-        field.getMetadata().get(ArrowType.ExtensionType.EXTENSION_METADATA_KEY_METADATA);
-
-    // We need to scan our extension types manually because of
-    // https://github.com/apache/arrow/issues/38891
-    if (JSONType.EXTENSION_NAME.equals(extensionName)
-        && JSONType.INSTANCE.serialize().equals(extensionMetadata)
-        && JSONType.INSTANCE.storageType().equals(fieldType)) {
-      fieldType = JSONType.INSTANCE;
-    } else if (UUIDType.EXTENSION_NAME.equals(extensionName)
-        && UUIDType.INSTANCE.serialize().equals(extensionMetadata)
-        && UUIDType.INSTANCE.storageType().equals(fieldType)) {
-      fieldType = UUIDType.INSTANCE;
-    }
-
     return Column.builder()
         .name(field.getName())
         .unique(isUnique)
         .primaryKey(isPrimaryKey)
         .incrementalKey(isIncrementalKey)
-        .type(fieldType)
+        .type(field.getType())
         .build();
   }
 
diff --git a/lib/src/test/java/io/cloudquery/helper/ArrowHelperTest.java b/lib/src/test/java/io/cloudquery/helper/ArrowHelperTest.java
index 5c587e4..9c0d87e 100644
--- a/lib/src/test/java/io/cloudquery/helper/ArrowHelperTest.java
+++ b/lib/src/test/java/io/cloudquery/helper/ArrowHelperTest.java
@@ -7,6 +7,7 @@
 import io.cloudquery.schema.Column;
 import io.cloudquery.schema.Resource;
 import io.cloudquery.schema.Table;
+import io.cloudquery.types.Extensions;
 import io.cloudquery.types.JSONType;
 import io.cloudquery.types.UUIDType;
 import java.io.IOException;
@@ -128,6 +129,8 @@ public class ArrowHelperTest {
 
   @Test
   public void testToArrowSchema() {
+    Extensions.registerExtensions();
+
     Schema arrowSchema = ArrowHelper.toArrowSchema(TEST_TABLE);
 
     for (Column col : TEST_TABLE.getColumns()) {
@@ -197,6 +200,8 @@ public void testToArrowSchema() {
 
   @Test
   public void testFromArrowSchema() {
+    Extensions.registerExtensions();
+
     List<Field> fields =
         List.of(
             Field.nullable("string_column1", ArrowType.Utf8.INSTANCE),
@@ -218,6 +223,8 @@ public void testFromArrowSchema() {
 
   @Test
   public void testRoundTripTableEncoding() throws IOException {
+    Extensions.registerExtensions();
+
     ByteString byteString = ArrowHelper.encode(TEST_TABLE);
     Table table = ArrowHelper.decode(byteString);
 
@@ -237,6 +244,8 @@ public void testRoundTripTableEncoding() throws IOException {
 
   @Test
   public void testRoundTripResourceEncoding() throws Exception {
+    Extensions.registerExtensions();
+
     Resource resource = Resource.builder().table(TEST_TABLE).build();
     resource.set("pk", "test_pk");
     resource.set("big_int", -1024L);