Bulk Load CDK: fix airbytevalue serialize

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/AirbyteValue.kt

@@ -4,7 +4,13 @@
 
 package io.airbyte.cdk.load.data
 
+import com.fasterxml.jackson.annotation.JsonValue
+import com.fasterxml.jackson.core.JsonGenerator
 import com.fasterxml.jackson.databind.JsonNode
+import com.fasterxml.jackson.databind.JsonSerializer
+import com.fasterxml.jackson.databind.SerializerProvider
+import com.fasterxml.jackson.databind.annotation.JsonSerialize
+import com.fasterxml.jackson.databind.node.NullNode
 import java.math.BigDecimal
 import java.math.BigInteger
 import java.time.LocalDate
@@ -48,6 +54,11 @@ sealed interface AirbyteValue {
 // (mostly the date/timestamp/time types - everything else is fine)
 data object NullValue : AirbyteValue, Comparable<NullValue> {
     override fun compareTo(other: NullValue): Int = 0
+
+    // make sure that we serialize this as a NullNode, rather than an empty object.
+    // We can't just return `null`, because jackson treats that as an error
+    // and falls back to its normal serialization behavior.
+    @JsonValue fun toJson(): NullNode = NullNode.instance
 }
 
 @JvmInline
@@ -75,34 +86,39 @@ value class NumberValue(val value: BigDecimal) : AirbyteValue, Comparable<Number
 value class DateValue(val value: LocalDate) : AirbyteValue, Comparable<DateValue> {
     constructor(date: String) : this(LocalDate.parse(date))
     override fun compareTo(other: DateValue): Int = value.compareTo(other.value)
+    @JsonValue fun toJson() = value.toString()
 }
 
 @JvmInline
 value class TimestampWithTimezoneValue(val value: OffsetDateTime) :
     AirbyteValue, Comparable<TimestampWithTimezoneValue> {
     constructor(timestamp: String) : this(OffsetDateTime.parse(timestamp))
     override fun compareTo(other: TimestampWithTimezoneValue): Int = value.compareTo(other.value)
+    @JsonValue fun toJson() = value.toString()
 }
 
 @JvmInline
 value class TimestampWithoutTimezoneValue(val value: LocalDateTime) :
     AirbyteValue, Comparable<TimestampWithoutTimezoneValue> {
     constructor(timestamp: String) : this(LocalDateTime.parse(timestamp))
     override fun compareTo(other: TimestampWithoutTimezoneValue): Int = value.compareTo(other.value)
+    @JsonValue fun toJson() = value.toString()
 }
 
 @JvmInline
 value class TimeWithTimezoneValue(val value: OffsetTime) :
     AirbyteValue, Comparable<TimeWithTimezoneValue> {
     constructor(time: String) : this(OffsetTime.parse(time))
     override fun compareTo(other: TimeWithTimezoneValue): Int = value.compareTo(other.value)
+    @JsonValue fun toJson() = value.toString()
 }
 
 @JvmInline
 value class TimeWithoutTimezoneValue(val value: LocalTime) :
     AirbyteValue, Comparable<TimeWithoutTimezoneValue> {
     constructor(time: String) : this(LocalTime.parse(time))
     override fun compareTo(other: TimeWithoutTimezoneValue): Int = value.compareTo(other.value)
+    @JsonValue fun toJson() = value.toString()
 }
 
 @JvmInline
@@ -112,12 +128,28 @@ value class ArrayValue(val values: List<AirbyteValue>) : AirbyteValue {
     }
 }
 
+// jackson can't figure out how to serialize this class,
+// so write a custom serializer that just serializes the map directly.
+// For some reason, the more obvious `@JsonValue fun toJson() = values`
+// doesn't work either.
+@JsonSerialize(using = ObjectValueSerializer::class)
 @JvmInline
 value class ObjectValue(val values: LinkedHashMap<String, AirbyteValue>) : AirbyteValue {
+    @JsonValue fun toJson() = values
     companion object {
         fun from(map: Map<String, Any?>): ObjectValue =
             ObjectValue(map.mapValuesTo(linkedMapOf()) { (_, v) -> AirbyteValue.from(v) })
     }
 }
 
+private class ObjectValueSerializer : JsonSerializer<ObjectValue>() {
+    override fun serialize(
+        value: ObjectValue,
+        gen: JsonGenerator,
+        serializers: SerializerProvider,
+    ) {
+        gen.writePOJO(value.values)
+    }
+}
+
 @JvmInline value class UnknownValue(val value: JsonNode) : AirbyteValue

airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/data/json/AirbyteValueToJsonTest.kt

@@ -4,14 +4,33 @@
 
 package io.airbyte.cdk.load.data.json
 
+import io.airbyte.cdk.load.data.AirbyteValue
 import io.airbyte.cdk.load.data.ArrayValue
 import io.airbyte.cdk.load.data.BooleanValue
+import io.airbyte.cdk.load.data.DateValue
 import io.airbyte.cdk.load.data.IntegerValue
+import io.airbyte.cdk.load.data.NullValue
 import io.airbyte.cdk.load.data.NumberValue
 import io.airbyte.cdk.load.data.ObjectValue
 import io.airbyte.cdk.load.data.StringValue
+import io.airbyte.cdk.load.data.TimeWithTimezoneValue
+import io.airbyte.cdk.load.data.TimeWithoutTimezoneValue
+import io.airbyte.cdk.load.data.TimestampWithTimezoneValue
+import io.airbyte.cdk.load.data.TimestampWithoutTimezoneValue
+import io.airbyte.cdk.load.data.UnknownValue
+import io.airbyte.cdk.load.util.Jsons
+import io.airbyte.cdk.load.util.serializeToString
+import java.math.BigDecimal
+import java.math.BigInteger
+import java.time.LocalDate
+import java.time.LocalDateTime
+import java.time.LocalTime
+import java.time.OffsetDateTime
+import java.time.OffsetTime
+import kotlin.test.assertEquals
 import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertDoesNotThrow
 
 class AirbyteValueToJsonTest {
     @Test
@@ -38,4 +57,36 @@ class AirbyteValueToJsonTest {
 
         Assertions.assertEquals(airbyteValue, roundTripValue)
     }
+
+    /**
+     * We have some code that relies on being able to directly Jackson-serialize an AirbyteValue
+     * (for example, the [serializeToString] extension method). Verify that this behaves correctly.
+     */
+    @Test
+    fun testAllTypesSerialization() {
+        val testCases: Map<AirbyteValue, String> =
+            mapOf(
+                NullValue to "null",
+                StringValue("foo") to "\"foo\"",
+                BooleanValue(true) to "true",
+                IntegerValue(BigInteger("42")) to "42",
+                NumberValue(BigDecimal("42.1")) to "42.1",
+                DateValue(LocalDate.parse("2024-01-23")) to "\"2024-01-23\"",
+                TimestampWithTimezoneValue(OffsetDateTime.parse("2024-01-23T12:34:56.78Z")) to
+                    "\"2024-01-23T12:34:56.780Z\"",
+                TimestampWithoutTimezoneValue(LocalDateTime.parse("2024-01-23T12:34:56.78")) to
+                    "\"2024-01-23T12:34:56.780\"",
+                TimeWithTimezoneValue(OffsetTime.parse("12:34:56.78Z")) to "\"12:34:56.780Z\"",
+                TimeWithoutTimezoneValue(LocalTime.parse("12:34:56.78")) to "\"12:34:56.780\"",
+                ArrayValue(listOf(NullValue, ArrayValue(listOf(NullValue)))) to "[null,[null]]",
+                ObjectValue(linkedMapOf("foo" to ObjectValue(linkedMapOf("bar" to NullValue)))) to
+                    """{"foo":{"bar":null}}""",
+                UnknownValue(Jsons.readTree("""{"foo": "bar"}""")) to """{"foo":"bar"}"""
+            )
+        testCases.forEach { (value, expectedSerialization) ->
+            val actual =
+                assertDoesNotThrow("Failed to serialize $value") { Jsons.writeValueAsString(value) }
+            assertEquals(expectedSerialization, actual, "Incorrect serialization for $value")
+        }
+    }
 }