Update CDK to pass DestinationRecordRaw around (#55737)

Co-authored-by: Octavia Squidington III <[email protected]>
Co-authored-by: Edward Gao <[email protected]>

Author: 3 people

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/config/SyncBeanFactory.kt

@@ -9,7 +9,7 @@ import io.airbyte.cdk.load.command.DestinationConfiguration
 import io.airbyte.cdk.load.command.DestinationStream
 import io.airbyte.cdk.load.message.BatchEnvelope
 import io.airbyte.cdk.load.message.ChannelMessageQueue
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.MultiProducerChannel
 import io.airbyte.cdk.load.message.PartitionedQueue
 import io.airbyte.cdk.load.message.PipelineEvent
@@ -120,7 +120,7 @@ class SyncBeanFactory {
     @Named("recordQueue")
     fun recordQueue(
         loadStrategy: LoadStrategy? = null,
-    ): PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>> {
+    ): PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>> {
         return PartitionedQueue(
             Array(loadStrategy?.inputPartitions ?: 1) {
                 ChannelMessageQueue(Channel(Channel.UNLIMITED))

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/message/DestinationMessage.kt

@@ -11,6 +11,7 @@ import io.airbyte.cdk.load.command.DestinationStream
 import io.airbyte.cdk.load.data.AirbyteType
 import io.airbyte.cdk.load.data.AirbyteValue
 import io.airbyte.cdk.load.data.AirbyteValueDeepCoercingMapper
+import io.airbyte.cdk.load.data.EnrichedAirbyteValue
 import io.airbyte.cdk.load.data.IntegerValue
 import io.airbyte.cdk.load.data.StringValue
 import io.airbyte.cdk.load.data.TimestampWithTimezoneValue
@@ -143,6 +144,9 @@ data class DestinationRecord(
             serialized.length.toLong()
         )
     }
+    fun asDestinationRecordRaw(): DestinationRecordRaw {
+        return DestinationRecordRaw(stream, message, serialized)
+    }
 }
 
 /**
@@ -163,6 +167,43 @@ data class DestinationRecordAirbyteValue(
     val serializedSizeBytes: Long = 0L
 )
 
+data class EnrichedDestinationRecordAirbyteValue(
+    val stream: DestinationStream.Descriptor,
+    val declaredFields: Map<String, EnrichedAirbyteValue>,
+    val airbyteMetaFields: Map<String, EnrichedAirbyteValue>,
+    val undeclaredFields: Map<String, JsonNode>,
+    val emittedAtMs: Long,
+    val meta: Meta?,
+    val serializedSizeBytes: Long = 0L
+)
+
+data class DestinationRecordRaw(
+    val stream: DestinationStream.Descriptor,
+    private val rawData: AirbyteMessage,
+    private val serialized: String
+) {
+    fun asRawJson(): JsonNode {
+        return rawData.record.data
+    }
+
+    fun asDestinationRecordAirbyteValue(): DestinationRecordAirbyteValue {
+        return DestinationRecordAirbyteValue(
+            stream,
+            rawData.record.data.toAirbyteValue(),
+            rawData.record.emittedAt,
+            Meta(
+                rawData.record.meta?.changes?.map { Meta.Change(it.field, it.change, it.reason) }
+                    ?: emptyList()
+            ),
+            serialized.length.toLong()
+        )
+    }
+
+    fun asEnrichedDestinationRecordAirbyteValue(): EnrichedDestinationRecordAirbyteValue {
+        TODO()
+    }
+}
+
 data class DestinationFile(
     override val stream: DestinationStream.Descriptor,
     val emittedAtMs: Long,

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/pipeline/DirectLoadPipelineStep.kt

@@ -4,7 +4,7 @@
 
 package io.airbyte.cdk.load.pipeline
 
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.PartitionedQueue
 import io.airbyte.cdk.load.message.PipelineEvent
 import io.airbyte.cdk.load.message.QueueWriter
@@ -24,8 +24,7 @@ import jakarta.inject.Singleton
 class DirectLoadPipelineStep<S : DirectLoader>(
     val accumulator: DirectLoadRecordAccumulator<S, StreamKey>,
     @Named("recordQueue")
-    val inputQueue:
-        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>>,
+    val inputQueue: PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>>,
     @Named("batchStateUpdateQueue") val batchQueue: QueueWriter<BatchUpdate>,
     @Value("\${airbyte.destination.core.record-batch-size-override:null}")
     val batchSizeOverride: Long? = null,

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/pipeline/DirectLoadRecordAccumulator.kt

@@ -5,7 +5,7 @@
 package io.airbyte.cdk.load.pipeline
 
 import io.airbyte.cdk.load.message.Batch
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.WithBatchState
 import io.airbyte.cdk.load.message.WithStream
 import io.airbyte.cdk.load.write.DirectLoader
@@ -25,15 +25,12 @@ data class DirectLoadAccResult(override val state: Batch.State) : WithBatchState
 @Requires(bean = DirectLoaderFactory::class)
 class DirectLoadRecordAccumulator<S : DirectLoader, K : WithStream>(
     val directLoaderFactory: DirectLoaderFactory<S>
-) : BatchAccumulator<S, K, DestinationRecordAirbyteValue, DirectLoadAccResult> {
+) : BatchAccumulator<S, K, DestinationRecordRaw, DirectLoadAccResult> {
     override fun start(key: K, part: Int): S {
         return directLoaderFactory.create(key.stream, part)
     }
 
-    override fun accept(
-        record: DestinationRecordAirbyteValue,
-        state: S
-    ): Pair<S, DirectLoadAccResult?> {
+    override fun accept(record: DestinationRecordRaw, state: S): Pair<S, DirectLoadAccResult?> {
         state.accept(record).let {
             return when (it) {
                 is Incomplete -> Pair(state, null)

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/pipeline/InputPartitioner.kt

@@ -4,7 +4,7 @@
 
 package io.airbyte.cdk.load.pipeline
 
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.micronaut.context.annotation.Secondary
 import jakarta.inject.Singleton
 import kotlin.math.abs
@@ -14,13 +14,13 @@ import kotlin.math.abs
  * partitioned by a hash of the stream name and namespace.
  */
 interface InputPartitioner {
-    fun getPartition(record: DestinationRecordAirbyteValue, numParts: Int): Int
+    fun getPartition(record: DestinationRecordRaw, numParts: Int): Int
 }
 
 @Singleton
 @Secondary
 class ByStreamInputPartitioner : InputPartitioner {
-    override fun getPartition(record: DestinationRecordAirbyteValue, numParts: Int): Int {
+    override fun getPartition(record: DestinationRecordRaw, numParts: Int): Int {
         return abs(record.stream.hashCode()) % numParts
     }
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/task/DestinationTaskLauncher.kt

@@ -11,7 +11,7 @@ import io.airbyte.cdk.load.command.DestinationStream
 import io.airbyte.cdk.load.message.BatchEnvelope
 import io.airbyte.cdk.load.message.ChannelMessageQueue
 import io.airbyte.cdk.load.message.CheckpointMessageWrapped
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.DestinationStreamEvent
 import io.airbyte.cdk.load.message.MessageQueue
 import io.airbyte.cdk.load.message.MessageQueueSupplier
@@ -145,7 +145,7 @@ class DefaultDestinationTaskLauncher<K : WithStream>(
     // New interface shim
     @Named("recordQueue")
     private val recordQueueForPipeline:
-        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>>,
+        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>>,
     @Named("batchStateUpdateQueue") private val batchUpdateQueue: ChannelMessageQueue<BatchUpdate>,
     private val loadPipeline: LoadPipeline?,
     private val partitioner: InputPartitioner,

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/task/internal/InputConsumerTask.kt

@@ -15,7 +15,7 @@ import io.airbyte.cdk.load.message.DestinationFile
 import io.airbyte.cdk.load.message.DestinationFileStreamComplete
 import io.airbyte.cdk.load.message.DestinationFileStreamIncomplete
 import io.airbyte.cdk.load.message.DestinationRecord
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.DestinationRecordStreamComplete
 import io.airbyte.cdk.load.message.DestinationRecordStreamIncomplete
 import io.airbyte.cdk.load.message.DestinationStreamAffinedMessage
@@ -80,7 +80,7 @@ class DefaultInputConsumerTask(
     // Required by new interface
     @Named("recordQueue")
     private val recordQueueForPipeline:
-        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>>,
+        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>>,
     private val loadPipeline: LoadPipeline? = null,
     private val partitioner: InputPartitioner,
     private val openStreamQueue: QueueWriter<DestinationStream>
@@ -158,7 +158,7 @@ class DefaultInputConsumerTask(
         val manager = syncManager.getStreamManager(stream)
         when (val message = reserved.value) {
             is DestinationRecord -> {
-                val record = message.asRecordMarshaledToAirbyteValue()
+                val record = message.asDestinationRecordRaw()
                 manager.incrementReadCount()
                 val pipelineMessage =
                     PipelineMessage(
@@ -310,7 +310,7 @@ interface InputConsumerTaskFactory {
 
         // Required by new interface
         recordQueueForPipeline:
-            PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>>,
+            PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>>,
         loadPipeline: LoadPipeline?,
         partitioner: InputPartitioner,
         openStreamQueue: QueueWriter<DestinationStream>,
@@ -333,7 +333,7 @@ class DefaultInputConsumerTaskFactory(
 
         // Required by new interface
         recordQueueForPipeline:
-            PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>>,
+            PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>>,
         loadPipeline: LoadPipeline?,
         partitioner: InputPartitioner,
         openStreamQueue: QueueWriter<DestinationStream>,

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/write/DirectLoader.kt

@@ -5,7 +5,7 @@
 package io.airbyte.cdk.load.write
 
 import io.airbyte.cdk.load.command.DestinationStream
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 
 /**
  * [DirectLoader] is for the use case where records are loaded directly into the destination or
@@ -56,7 +56,7 @@ interface DirectLoader : AutoCloseable {
      * Called once per record until it returns [Complete], after which [close] is called, the loader
      * is discarded, and the records are considered processed by the platform.
      */
-    fun accept(record: DestinationRecordAirbyteValue): DirectLoadResult
+    fun accept(record: DestinationRecordRaw): DirectLoadResult
 
     /**
      * Called by the CDK to force work to finish. It will only be called if the last call to

airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/task/DestinationTaskLauncherTest.kt

@@ -13,7 +13,7 @@ import io.airbyte.cdk.load.command.MockDestinationConfiguration
 import io.airbyte.cdk.load.message.Batch
 import io.airbyte.cdk.load.message.BatchEnvelope
 import io.airbyte.cdk.load.message.CheckpointMessageWrapped
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.DestinationStreamEvent
 import io.airbyte.cdk.load.message.MessageQueue
 import io.airbyte.cdk.load.message.MessageQueueSupplier
@@ -159,7 +159,7 @@ class DestinationTaskLauncherTest {
             destinationTaskLauncher: DestinationTaskLauncher,
             fileTransferQueue: MessageQueue<FileTransferQueueMessage>,
             recordQueueForPipeline:
-                PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>>,
+                PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>>,
             loadPipeline: LoadPipeline?,
             partitioner: InputPartitioner,
             openStreamQueue: QueueWriter<DestinationStream>,

airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/task/DestinationTaskLauncherUTest.kt

@@ -11,7 +11,7 @@ import io.airbyte.cdk.load.message.Batch
 import io.airbyte.cdk.load.message.BatchEnvelope
 import io.airbyte.cdk.load.message.ChannelMessageQueue
 import io.airbyte.cdk.load.message.CheckpointMessageWrapped
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.DestinationStreamEvent
 import io.airbyte.cdk.load.message.MessageQueue
 import io.airbyte.cdk.load.message.MessageQueueSupplier
@@ -95,14 +95,14 @@ class DestinationTaskLauncherUTest {
     private val openStreamQueue: MessageQueue<DestinationStream> = mockk(relaxed = true)
 
     private val recordQueueForPipeline:
-        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>> =
+        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>> =
         mockk(relaxed = true)
     private val batchUpdateQueue: ChannelMessageQueue<BatchUpdate> = mockk(relaxed = true)
     private val partitioner: InputPartitioner = mockk(relaxed = true)
     private val updateBatchTaskFactory: UpdateBatchStateTaskFactory = mockk(relaxed = true)
 
     private fun getDefaultDestinationTaskLauncher(
-        useFileTranfer: Boolean,
+        useFileTransfer: Boolean,
         loadPipeline: LoadPipeline? = null
     ): DefaultDestinationTaskLauncher<Nothing> {
         return DefaultDestinationTaskLauncher(
@@ -124,7 +124,7 @@ class DestinationTaskLauncherUTest {
             updateCheckpointsTask,
             failStreamTaskFactory,
             failSyncTaskFactory,
-            useFileTranfer,
+            useFileTransfer,
             inputFlow,
             recordQueueSupplier,
             checkpointQueue,

airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/task/internal/InputConsumerTaskUTest.kt

@@ -10,7 +10,7 @@ import io.airbyte.cdk.load.data.ObjectTypeWithoutSchema
 import io.airbyte.cdk.load.message.CheckpointMessageWrapped
 import io.airbyte.cdk.load.message.DestinationMessage
 import io.airbyte.cdk.load.message.DestinationRecord
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.message.DestinationStreamEvent
 import io.airbyte.cdk.load.message.MessageQueue
 import io.airbyte.cdk.load.message.MessageQueueSupplier
@@ -47,7 +47,7 @@ class InputConsumerTaskUTest {
     @MockK lateinit var fileTransferQueue: MessageQueue<FileTransferQueueMessage>
     @MockK
     lateinit var recordQueueForPipeline:
-        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>>
+        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordRaw>>>
     @MockK lateinit var partitioner: InputPartitioner
     @MockK lateinit var openStreamQueue: QueueWriter<DestinationStream>
 

airbyte-integrations/connectors/destination-dev-null/build.gradle

@@ -6,7 +6,7 @@ plugins {
 airbyteBulkConnector {
     core = 'load'
     toolkits = []
-    cdk = 'local'
+    cdk = '0.344'
 }
 
 application {

airbyte-integrations/connectors/destination-dev-null/metadata.yaml

@@ -10,7 +10,7 @@ data:
     - suite: integrationTests
   connectorType: destination
   definitionId: a7bcc9d8-13b3-4e49-b80d-d020b90045e3
-  dockerImageTag: 0.7.18
+  dockerImageTag: 0.7.19
   dockerRepository: airbyte/destination-dev-null
   documentationUrl: https://docs.airbyte.com/integrations/destinations/dev-null
   githubIssueLabel: destination-dev-null

airbyte-integrations/connectors/destination-dev-null/src/main/kotlin/io/airbyte/integrations/destination/dev_null/DevNullDirectLoader.kt

@@ -5,7 +5,7 @@
 package io.airbyte.integrations.destination.dev_null
 
 import io.airbyte.cdk.load.command.DestinationStream
-import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.write.DirectLoader
 import io.airbyte.cdk.load.write.DirectLoaderFactory
 import io.github.oshai.kotlinlogging.KotlinLogging
@@ -18,9 +18,9 @@ abstract class DevNullDirectLoader(
 ) : DirectLoader {
     private var recordCount: Long = 0L
 
-    abstract fun acceptInner(record: DestinationRecordAirbyteValue)
+    abstract fun acceptInner(record: DestinationRecordRaw)
 
-    override fun accept(record: DestinationRecordAirbyteValue): DirectLoader.DirectLoadResult {
+    override fun accept(record: DestinationRecordRaw): DirectLoader.DirectLoadResult {
         acceptInner(record)
         return if (++recordCount % config.ackRatePerRecord == 0L) {
             DirectLoader.Complete
@@ -87,7 +87,7 @@ class LoggingDirectLoader(
     private var logCount: Long = 0L
     private val prng: Random = loggingConfig.seed?.let { Random(it) } ?: Random.Default
 
-    override fun acceptInner(record: DestinationRecordAirbyteValue) {
+    override fun acceptInner(record: DestinationRecordRaw) {
         if (recordCount++ % loggingConfig.logEvery == 0L) {
             if (loggingConfig.sampleRate == 1.0 || prng.nextDouble() < loggingConfig.sampleRate) {
                 if (++logCount < loggingConfig.maxEntryCount) {
@@ -101,15 +101,15 @@ class LoggingDirectLoader(
 }
 
 class SilentDirectLoader(config: DevNullConfiguration) : DevNullDirectLoader(config) {
-    override fun acceptInner(record: DestinationRecordAirbyteValue) {
+    override fun acceptInner(record: DestinationRecordRaw) {
         /* Do nothing */
     }
 }
 
 class ThrottledDirectLoader(config: DevNullConfiguration, private val millisPerRecord: Long) :
     DevNullDirectLoader(config) {
 
-    override fun acceptInner(record: DestinationRecordAirbyteValue) {
+    override fun acceptInner(record: DestinationRecordRaw) {
         Thread.sleep(millisPerRecord)
     }
 }
@@ -123,7 +123,7 @@ class FailingDirectLoader(
 
     private var messageCount: Long = 0L
 
-    override fun acceptInner(record: DestinationRecordAirbyteValue) {
+    override fun acceptInner(record: DestinationRecordRaw) {
         if (messageCount++ > numMessages) {
             val message =
                 "Failing Destination(stream=$stream, numMessages=$numMessages: failing at $record)"

airbyte-integrations/connectors/destination-s3-data-lake/build.gradle

@@ -6,7 +6,7 @@ plugins {
 airbyteBulkConnector {
     core = 'load'
     toolkits = ['load-iceberg-parquet', 'load-aws']
-    cdk = '0.343'
+    cdk = '0.344'
 }
 
 application {

airbyte-integrations/connectors/destination-s3-data-lake/metadata.yaml

@@ -26,7 +26,7 @@ data:
             alias: airbyte-connector-testing-secret-store
   connectorType: destination
   definitionId: 716ca874-520b-4902-9f80-9fad66754b89
-  dockerImageTag: 0.3.16
+  dockerImageTag: 0.3.17
   dockerRepository: airbyte/destination-s3-data-lake
   documentationUrl: https://docs.airbyte.com/integrations/destinations/s3-data-lake
   githubIssueLabel: destination-s3-data-lake