Bulk load CDK: Add test with every single airbyte type (#48370)

airbytehq · Nov 5, 2024 · 9a51bce · 9a51bce
1 parent ddfa537
commit 9a51bce
Show file tree

Hide file tree

Showing 8 changed files with 432 additions and 13 deletions.
diff --git a/...kotlin/io/airbyte/cdk/load/mock_integration_test/MockBasicFunctionalityIntegrationTest.kt b/...kotlin/io/airbyte/cdk/load/mock_integration_test/MockBasicFunctionalityIntegrationTest.kt
@@ -8,6 +8,7 @@ import io.airbyte.cdk.load.test.util.NoopDestinationCleaner
 import io.airbyte.cdk.load.test.util.NoopExpectedRecordMapper
 import io.airbyte.cdk.load.test.util.NoopNameMapper
 import io.airbyte.cdk.load.write.BasicFunctionalityIntegrationTest
+import io.airbyte.cdk.load.write.Untyped
 import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
 
@@ -25,6 +26,7 @@ class MockBasicFunctionalityIntegrationTest :
         promoteUnionToObject = false,
         preserveUndeclaredFields = true,
         commitDataIncrementally = false,
+        allTypesBehavior = Untyped,
     ) {
     @Test
     override fun testBasicWrite() {
@@ -86,4 +88,9 @@ class MockBasicFunctionalityIntegrationTest :
     override fun testUnions() {
         super.testUnions()
     }
+
+    @Test
+    override fun testAllTypes() {
+        super.testAllTypes()
+    }
 }
diff --git a/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/AirbyteValue.kt b/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/AirbyteValue.kt
@@ -16,6 +16,7 @@ sealed interface AirbyteValue {
     companion object {
         fun from(value: Any?): AirbyteValue =
             when (value) {
+                is AirbyteValue -> value
                 null -> NullValue
                 is String -> StringValue(value)
                 is Boolean -> BooleanValue(value)

diff --git a/...k/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/json/AirbyteTypeToJsonSchema.kt b/...k/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/json/AirbyteTypeToJsonSchema.kt
@@ -64,7 +64,8 @@ class AirbyteTypeToJsonSchema {
                 val timestampNode = ofType("string").put("format", "date-time")
                 timestampNode.put("airbyte_type", "timestamp_without_timezone")
             }
-            is UnknownType -> throw IllegalArgumentException("Unknown type: $airbyteType")
+            // In case of unknown type, just return {} (i.e. the accept-all JsonSchema)
+            is UnknownType -> JsonNodeFactory.instance.objectNode()
         }
     }
 

diff --git a/...te-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/json/JsonToAirbyteValue.kt b/...te-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/json/JsonToAirbyteValue.kt
@@ -6,6 +6,7 @@ package io.airbyte.cdk.load.data.json
 
 import com.fasterxml.jackson.databind.JsonNode
 import io.airbyte.cdk.load.data.*
+import io.airbyte.cdk.util.Jsons
 import java.math.BigDecimal
 
 /**
@@ -34,7 +35,7 @@ class JsonToAirbyteValue {
                 is ObjectType -> toObject(json, schema)
                 is ObjectTypeWithoutSchema,
                 is ObjectTypeWithEmptySchema -> toObjectWithoutSchema(json)
-                is StringType -> StringValue(json.asText())
+                is StringType -> toString(json)
                 is TimeTypeWithTimezone,
                 is TimeTypeWithoutTimezone -> TimeValue(json.asText())
                 is TimestampTypeWithTimezone,
@@ -67,6 +68,14 @@ class JsonToAirbyteValue {
         return ArrayValue(json.map { fromJson(it) })
     }
 
+    private fun toString(json: JsonNode): StringValue {
+        return if (json.isTextual) {
+            StringValue(json.asText())
+        } else {
+            StringValue(Jsons.writeValueAsString(json))
+        }
+    }
+
     private fun toBoolean(json: JsonNode): BooleanValue {
         val boolVal =
             when {

diff --git a/...-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/RecordDiffer.kt b/...-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/RecordDiffer.kt
@@ -271,17 +271,25 @@ class RecordDiffer(
                 // Handle temporal types specifically, because they require explicit parsing
                 return when (v1) {
                     is DateValue ->
-                        LocalDate.parse(v1.value)
-                            .compareTo(LocalDate.parse((v2 as DateValue).value))
+                        try {
+                            LocalDate.parse(v1.value)
+                                .compareTo(LocalDate.parse((v2 as DateValue).value))
+                        } catch (e: Exception) {
+                            v1.value.compareTo((v2 as DateValue).value)
+                        }
                     is TimeValue -> {
                         try {
                             val time1 = LocalTime.parse(v1.value)
                             val time2 = LocalTime.parse((v2 as TimeValue).value)
                             time1.compareTo(time2)
                         } catch (e: Exception) {
-                            val time1 = OffsetTime.parse(v1.value)
-                            val time2 = OffsetTime.parse((v2 as TimeValue).value)
-                            time1.compareTo(time2)
+                            try {
+                                val time1 = OffsetTime.parse(v1.value)
+                                val time2 = OffsetTime.parse((v2 as TimeValue).value)
+                                time1.compareTo(time2)
+                            } catch (e: Exception) {
+                                v1.value.compareTo((v2 as TimeValue).value)
+                            }
                         }
                     }
                     is TimestampValue -> {
@@ -290,9 +298,13 @@ class RecordDiffer(
                             val ts2 = LocalDateTime.parse((v2 as TimestampValue).value)
                             ts1.compareTo(ts2)
                         } catch (e: Exception) {
-                            val ts1 = OffsetDateTime.parse(v1.value)
-                            val ts2 = OffsetDateTime.parse((v2 as TimestampValue).value)
-                            ts1.compareTo(ts2)
+                            try {
+                                val ts1 = OffsetDateTime.parse(v1.value)
+                                val ts2 = OffsetDateTime.parse((v2 as TimestampValue).value)
+                                ts1.compareTo(ts2)
+                            } catch (e: Exception) {
+                                v1.value.compareTo((v2 as TimestampValue).value)
+                            }
                         }
                     }
                     // otherwise, just be a terrible person.