source-postgres: field type mapping (#65983)

Co-authored-by: Rodi Reich Zilberman <[email protected]>

Author: mwbayley

airbyte-cdk/bulk/core/base/src/main/kotlin/io/airbyte/cdk/data/JsonCodec.kt

@@ -336,18 +336,23 @@ data class ArrayEncoder<T>(
     override fun encode(decoded: List<T>): JsonNode =
         Jsons.arrayNode().apply {
             for (e in decoded) {
-                add(elementEncoder.encode(e))
+                // Note: in generics, T can be nullable!
+                if (e == null) add(NullCodec.encode(e)) else add(elementEncoder.encode(e))
             }
         }
 }
 
 data class ArrayDecoder<T>(
     val elementDecoder: JsonDecoder<T>,
-) : JsonDecoder<List<T>> {
-    override fun decode(encoded: JsonNode): List<T> {
+) : JsonDecoder<List<T?>> {
+    override fun decode(encoded: JsonNode): List<T?> {
         if (!encoded.isArray) {
             throw IllegalArgumentException("invalid array value $encoded")
         }
-        return encoded.elements().asSequence().map { elementDecoder.decode(it) }.toList()
+        return encoded
+            .elements()
+            .asSequence()
+            .map { if (it == null) null else elementDecoder.decode(it) }
+            .toList()
     }
 }

airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/discover/Field.kt

@@ -19,7 +19,8 @@ sealed interface DataOrMetaField {
 
 @Deprecated(
     message = "Use `DataOrMetaField` directly instead.",
-    replaceWith = ReplaceWith("DataOrMetaField"))
+    replaceWith = ReplaceWith("DataOrMetaField")
+)
 typealias FieldOrMetaField = DataOrMetaField
 /**
  * Root of our own type hierarchy for Airbyte record fields.
@@ -43,11 +44,12 @@ interface LosslessFieldType : FieldType {
     val jsonDecoder: JsonDecoder<*>
 }
 
-interface DataField: DataOrMetaField
+interface DataField : DataOrMetaField
 
 @Deprecated(
     message = "Use `EmittedField` directly instead.",
-    replaceWith = ReplaceWith("EmittedField"))
+    replaceWith = ReplaceWith("EmittedField")
+)
 typealias Field = EmittedField
 /**
  * Internal equivalent of [io.airbyte.protocol.models.Field] for values which come from the source

airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/output/NativeRecord.kt

@@ -81,7 +81,7 @@ fun <T> JsonEncoder<T>.toProtobufEncoder(): ProtoEncoder<*> {
         is LocalTimeCodec, -> localTimeProtoEncoder
         is LocalDateTimeCodec, -> localDateTimeProtoEncoder
         is OffsetTimeCodec, -> offsetTimeProtoEncoder
-        is ArrayEncoder<*>, -> anyProtoEncoder
+        is ArrayEncoder<*>, -> arrayProtoEncoder
         else -> anyProtoEncoder
     }
 }
@@ -162,33 +162,35 @@ val floatProtoEncoder =
 
 val nullProtoEncoder = generateProtoEncoder<Any?> { builder, _ -> builder.setIsNull(true) }
 val anyProtoEncoder = textProtoEncoder
-// typealias AnyProtoEncoder = TextProtoEncoder
+
+// For now arrays are encoded in protobuf as json strings
+val arrayProtoEncoder = textProtoEncoder
 
 fun NativeRecordPayload.toProtobuf(
     schema: Set<FieldOrMetaField>,
     recordMessageBuilder: AirbyteRecordMessageProtobuf.Builder,
     valueBuilder: AirbyteRecordMessage.AirbyteValueProtobuf.Builder
 ): AirbyteRecordMessageProtobuf.Builder {
     return recordMessageBuilder.apply {
-        schema
-            .sortedBy { it.id }
-            .forEachIndexed { index, field ->
-                // Protobuf does not have field names, so we use a sorted order of fields
-                // So for destination to know which fields it is, we order the fields alphabetically
-                // to make sure that the order is consistent.
-                this@toProtobuf[field.id]?.let { value ->
-                    @Suppress("UNCHECKED_CAST")
-                    setData(
-                        index,
-                        value.fieldValue?.let {
-                            (value.jsonEncoder.toProtobufEncoder() as ProtoEncoder<Any>).encode(
-                                valueBuilder.clear(),
-                                value.fieldValue
-                            )
+        // We use toSortedMap() to ensure that the order is consistent
+        // Since protobuf has no field name the contract with destination is that
+        // field are alphabetically ordered.
+        this@toProtobuf.toSortedMap().onEachIndexed { index, entry ->
+            @Suppress("UNCHECKED_CAST")
+            setData(
+                index,
+                entry.value.fieldValue?.let {
+                    (entry.value.jsonEncoder.toProtobufEncoder() as ProtoEncoder<Any>).encode(
+                        valueBuilder.clear(),
+                        when (entry.value.jsonEncoder) {
+                            // For arrays we use the value of its json string.
+                            is ArrayEncoder<*> -> entry.value.encode().toString()
+                            else -> entry.value.fieldValue!!
                         }
-                            ?: nullProtoEncoder.encode(valueBuilder.clear(), null)
+
                     )
                 }
-            }
+            )
+        }
     }
 }

airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/read/Feed.kt

@@ -2,8 +2,8 @@
 package io.airbyte.cdk.read
 
 import io.airbyte.cdk.StreamIdentifier
-import io.airbyte.cdk.discover.Field
 import io.airbyte.cdk.discover.DataOrMetaField
+import io.airbyte.cdk.discover.Field
 
 /**
  * [Feed] identifies part of the data consumed during a READ operation.

airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/read/StateManagerFactory.kt

@@ -14,8 +14,8 @@ import io.airbyte.cdk.command.StreamInputState
 import io.airbyte.cdk.data.AirbyteSchemaType
 import io.airbyte.cdk.data.ArrayAirbyteSchemaType
 import io.airbyte.cdk.data.LeafAirbyteSchemaType
-import io.airbyte.cdk.discover.Field
 import io.airbyte.cdk.discover.DataOrMetaField
+import io.airbyte.cdk.discover.Field
 import io.airbyte.cdk.discover.MetaField
 import io.airbyte.cdk.discover.MetaFieldDecorator
 import io.airbyte.cdk.discover.MetadataQuerier

airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/discover/JdbcMetadataQuerier.kt

@@ -130,7 +130,7 @@ class JdbcMetadataQuerier(
                     }
                 }
             val clause = if (constants.includePseudoColumns) " and pseudo-column(s)" else ""
-            log.info { "Discovered ${results.size} column(s)${clause}."}
+            log.info { "Discovered ${results.size} column(s)${clause}." }
         } catch (e: Exception) {
             throw RuntimeException("Column name discovery query failed: ${e.message}", e)
         }

airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/read/DefaultJdbcPartition.kt

@@ -5,7 +5,6 @@
 package io.airbyte.cdk.read
 
 import com.fasterxml.jackson.databind.JsonNode
-import com.fasterxml.jackson.databind.node.ObjectNode
 import io.airbyte.cdk.command.OpaqueStateValue
 import io.airbyte.cdk.discover.Field
 import io.airbyte.cdk.output.sockets.toJson
@@ -197,7 +196,8 @@ class DefaultJdbcSplittableSnapshotPartition(
     override fun incompleteState(lastRecord: SelectQuerier.ResultRow): OpaqueStateValue =
         DefaultJdbcStreamStateValue.snapshotCheckpoint(
             primaryKey = checkpointColumns,
-            primaryKeyCheckpoint = checkpointColumns.map { lastRecord.data.toJson()[it.id] ?: Jsons.nullNode() },
+            primaryKeyCheckpoint =
+                checkpointColumns.map { lastRecord.data.toJson()[it.id] ?: Jsons.nullNode() },
         )
 }
 
@@ -264,7 +264,8 @@ class DefaultJdbcSplittableSnapshotWithCursorPartition(
     override fun incompleteState(lastRecord: SelectQuerier.ResultRow): OpaqueStateValue =
         DefaultJdbcStreamStateValue.snapshotWithCursorCheckpoint(
             primaryKey = checkpointColumns,
-            primaryKeyCheckpoint = checkpointColumns.map { lastRecord.data.toJson()[it.id] ?: Jsons.nullNode() },
+            primaryKeyCheckpoint =
+                checkpointColumns.map { lastRecord.data.toJson()[it.id] ?: Jsons.nullNode() },
             cursor,
             cursorUpperBound,
         )

airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/read/DefaultJdbcPartitionFactory.kt

@@ -9,8 +9,8 @@ import io.airbyte.cdk.ConfigErrorException
 import io.airbyte.cdk.StreamIdentifier
 import io.airbyte.cdk.command.JdbcSourceConfiguration
 import io.airbyte.cdk.command.OpaqueStateValue
-import io.airbyte.cdk.discover.Field
 import io.airbyte.cdk.discover.DataOrMetaField
+import io.airbyte.cdk.discover.Field
 import io.airbyte.cdk.output.CatalogValidationFailureHandler
 import io.airbyte.cdk.output.InvalidCursor
 import io.airbyte.cdk.output.InvalidPrimaryKey

airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/read/JdbcPartition.kt

@@ -4,7 +4,6 @@
 
 package io.airbyte.cdk.read
 
-import com.fasterxml.jackson.databind.node.ObjectNode
 import io.airbyte.cdk.command.OpaqueStateValue
 
 /**
@@ -46,7 +45,6 @@ interface JdbcSplittablePartition<S : JdbcStreamState<*>> : JdbcPartition<S> {
 
     /** State value to emit when the partition is read up to (and including) [lastRecord]. */
     fun incompleteState(lastRecord: SelectQuerier.ResultRow): OpaqueStateValue
-
 }
 
 /** A [JdbcPartition] which allows cursor-based incremental reads. */

airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/read/JdbcPartitionReader.kt

@@ -1,16 +1,13 @@
 /* Copyright (c) 2024 Airbyte, Inc., all rights reserved. */
 package io.airbyte.cdk.read
 
-import com.fasterxml.jackson.databind.node.ObjectNode
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings
 import io.airbyte.cdk.TransientErrorException
 import io.airbyte.cdk.command.OpaqueStateValue
 import io.airbyte.cdk.discover.Field
 import io.airbyte.cdk.output.DataChannelMedium.*
 import io.airbyte.cdk.output.OutputMessageRouter
 import io.airbyte.cdk.output.sockets.NativeRecordPayload
-import io.airbyte.cdk.output.sockets.toJson
-import io.airbyte.cdk.util.Jsons
 import io.airbyte.protocol.models.v0.AirbyteStateMessage
 import io.airbyte.protocol.models.v0.AirbyteStreamStatusTraceMessage
 import java.time.Duration
@@ -196,7 +193,7 @@ class JdbcResumablePartitionReader<P : JdbcSplittablePartition<*>>(
             .use { result: SelectQuerier.Result ->
                 for (row in result) {
                     out(row)
-//                    lastRecord.set(row.data.toJson(Jsons.objectNode()))
+                    //                    lastRecord.set(row.data.toJson(Jsons.objectNode()))
                     lastRecord.set(row)
                     // Check activity periodically to handle timeout.
                     if (numRecords.incrementAndGet() % fetchSize == 0L) {