Save

Author: frifriSF59

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

@@ -147,7 +147,8 @@ data class DestinationRecord(
     fun asDestinationRecordRaw(): DestinationRecordRaw {
         return DestinationRecordRaw(
             stream,
-            message
+            message,
+            serialized
         )
     }
 }
@@ -182,8 +183,22 @@ data class EnrichedDestinationRecordAirbyteValue(
 
 data class DestinationRecordRaw(
     val stream: DestinationStream.Descriptor,
-    val rawData: AirbyteMessage
-)
+    private val rawData: AirbyteMessage,
+    private val serialized: String
+) {
+    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()
+        )
+    }
+}
 
 data class DestinationFile(
     override val stream: DestinationStream.Descriptor,

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

@@ -5,6 +5,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
@@ -25,7 +26,7 @@ class DirectLoadPipelineStep<S : DirectLoader>(
     val accumulator: DirectLoadRecordAccumulator<S, StreamKey>,
     @Named("recordQueue")
     val inputQueue:
-        PartitionedQueue<Reserved<PipelineEvent<StreamKey, DestinationRecordAirbyteValue>>>,
+        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

@@ -6,6 +6,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,13 +26,13 @@ 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,
+        record: DestinationRecordRaw,
         state: S
     ): Pair<S, DirectLoadAccResult?> {
         state.accept(record).let {

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

@@ -5,6 +5,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 +15,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

@@ -12,6 +12,7 @@ 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 +146,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

@@ -16,6 +16,7 @@ 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 +81,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>
@@ -310,7 +311,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 +334,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

@@ -6,6 +6,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 +57,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-integrations/connectors/destination-s3-data-lake/build.gradle

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

airbyte-integrations/connectors/destination-s3-data-lake/src/main/kotlin/io/airbyte/integrations/destination/s3_data_lake/S3DataLakeDirectLoader.kt

@@ -9,6 +9,7 @@ import io.airbyte.cdk.load.command.DestinationStream
 import io.airbyte.cdk.load.data.MapperPipeline
 import io.airbyte.cdk.load.data.iceberg.parquet.IcebergParquetPipelineFactory
 import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.toolkits.iceberg.parquet.io.IcebergTableWriterFactory
 import io.airbyte.cdk.load.toolkits.iceberg.parquet.io.IcebergUtil
 import io.airbyte.cdk.load.write.DirectLoader
@@ -79,17 +80,19 @@ class S3DataLakeDirectLoader(
         val commitLock: Any = Any()
     }
 
-    override fun accept(record: DestinationRecordAirbyteValue): DirectLoader.DirectLoadResult {
+    override fun accept(record: DestinationRecordRaw): DirectLoader.DirectLoadResult {
+        val recordAirbyteValue = record.asDestinationRecordAirbyteValue()
+
         val icebergRecord =
             icebergUtil.toRecord(
-                record = record,
+                record = recordAirbyteValue,
                 stream = stream,
                 tableSchema = schema,
                 pipeline = pipeline
             )
         writer.write(icebergRecord)
 
-        dataSize += record.serializedSizeBytes // TODO: use icebergRecord.size() instead?
+        dataSize += recordAirbyteValue.serializedSizeBytes // TODO: use icebergRecord.size() instead?
         if (dataSize < batchSize) {
             return DirectLoader.Incomplete
         }