Destination S3 V2: copy over legacy DATs, enable tests

airbyte-cdk/bulk/core/base/src/main/resources/application.yaml

@@ -8,3 +8,5 @@ airbyte:
   flush:
     rate-ms: 900000 # 15 minutes
     window-ms: 900000 # 15 minutes
+  destination:
+    record-batch-size: ${AIRBYTE_DESTINATION_RECORD_BATCH_SIZE:209715200}

airbyte-cdk/bulk/core/load/src/integrationTest/kotlin/io/airbyte/cdk/load/mock_integration_test/MockBasicFunctionalityIntegrationTest.kt

@@ -9,7 +9,6 @@ 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
 
 class MockBasicFunctionalityIntegrationTest :
@@ -34,7 +33,6 @@ class MockBasicFunctionalityIntegrationTest :
     }
 
     @Test
-    @Disabled
     override fun testMidSyncCheckpointingStreamState() {
         super.testMidSyncCheckpointingStreamState()
     }

airbyte-cdk/bulk/core/load/src/integrationTest/kotlin/io/airbyte/cdk/load/mock_integration_test/MockDestinationBackend.kt

@@ -13,9 +13,11 @@ import io.airbyte.cdk.load.test.util.DestinationDataDumper
 import io.airbyte.cdk.load.test.util.OutputRecord
 import io.airbyte.cdk.load.test.util.RecordDiffer
 import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.ConcurrentLinkedQueue
 
 object MockDestinationBackend {
-    private val files: MutableMap<String, MutableList<OutputRecord>> = ConcurrentHashMap()
+    private val files: ConcurrentHashMap<String, ConcurrentLinkedQueue<OutputRecord>> =
+        ConcurrentHashMap()
 
     fun insert(filename: String, vararg records: OutputRecord) {
         getFile(filename).addAll(records)
@@ -96,7 +98,7 @@ object MockDestinationBackend {
     }
 
     fun readFile(filename: String): List<OutputRecord> {
-        return getFile(filename)
+        return getFile(filename).toList()
     }
 
     fun deleteOldRecords(filename: String, minGenerationId: Long) {
@@ -105,8 +107,8 @@ object MockDestinationBackend {
         }
     }
 
-    private fun getFile(filename: String): MutableList<OutputRecord> {
-        return files.getOrPut(filename) { mutableListOf() }
+    private fun getFile(filename: String): ConcurrentLinkedQueue<OutputRecord> {
+        return files.getOrPut(filename) { ConcurrentLinkedQueue() }
     }
 }
 

airbyte-cdk/bulk/core/load/src/integrationTest/kotlin/io/airbyte/cdk/load/mock_integration_test/MockDestinationConfiguration.kt

@@ -11,8 +11,8 @@ import io.micronaut.context.annotation.Factory
 import jakarta.inject.Singleton
 
 class MockDestinationConfiguration : DestinationConfiguration() {
-    // override to 10KB instead of 200MB
-    override val recordBatchSizeBytes = 10 * 1024L
+    // Micro-batch for testing.
+    override val recordBatchSizeBytes = 1L
 }
 
 @Singleton

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/command/DestinationCatalog.kt

@@ -5,6 +5,7 @@
 package io.airbyte.cdk.load.command
 
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog
+import io.github.oshai.kotlinlogging.KotlinLogging
 import io.micronaut.context.annotation.Factory
 import io.micronaut.context.annotation.Secondary
 import jakarta.inject.Singleton
@@ -14,6 +15,8 @@ import jakarta.inject.Singleton
  * for usability.
  */
 data class DestinationCatalog(val streams: List<DestinationStream> = emptyList()) {
+    private val log = KotlinLogging.logger {}
+
     private val byDescriptor: Map<DestinationStream.Descriptor, DestinationStream> =
         streams.associateBy { it.descriptor }
 
@@ -23,6 +26,7 @@ data class DestinationCatalog(val streams: List<DestinationStream> = emptyList()
                 "Catalog must have at least one stream: check that files are in the correct location."
             )
         }
+        log.info { "Destination catalog initialized: $streams" }
     }
 
     fun getStream(name: String, namespace: String?): DestinationStream {

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

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.data
+
+import java.text.Normalizer
+import java.util.regex.Pattern
+
+class Transformations {
+    companion object {
+        private const val S3_SAFE_CHARACTERS = "\\p{Alnum}/!_.*')("
+        private const val S3_SPECIAL_CHARACTERS = "&$@=;:+,?-"
+        private val S3_CHARACTER_PATTERN =
+            "[^${S3_SAFE_CHARACTERS}${Pattern.quote(S3_SPECIAL_CHARACTERS)}]"
+
+        fun toS3SafeCharacters(input: String): String {
+            return Normalizer.normalize(input, Normalizer.Form.NFKD)
+                .replace(
+                    "\\p{M}".toRegex(),
+                    "",
+                ) // P{M} matches a code point that is not a combining mark (unicode)
+                .replace(S3_CHARACTER_PATTERN.toRegex(), "_")
+        }
+    }
+}

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

@@ -318,6 +318,8 @@ data class GlobalCheckpoint(
     override val destinationStats: Stats? = null,
     val checkpoints: List<Checkpoint> = emptyList(),
     override val additionalProperties: Map<String, Any>,
+    val originalTypeField: AirbyteStateMessage.AirbyteStateType? =
+        AirbyteStateMessage.AirbyteStateType.GLOBAL,
 ) : CheckpointMessage {
     /** Convenience constructor, primarily intended for use in tests. */
     constructor(
@@ -333,7 +335,7 @@ data class GlobalCheckpoint(
     override fun asProtocolMessage(): AirbyteMessage {
         val stateMessage =
             AirbyteStateMessage()
-                .withType(AirbyteStateMessage.AirbyteStateType.GLOBAL)
+                .withType(originalTypeField)
                 .withGlobal(
                     AirbyteGlobalState()
                         .withSharedState(state)
@@ -433,7 +435,11 @@ class DestinationMessageFactory(
                         namespace = status.streamDescriptor.namespace,
                         name = status.streamDescriptor.name,
                     )
-                if (message.trace.type == AirbyteTraceMessage.Type.STREAM_STATUS) {
+                if (
+                    (message.trace.type
+                        ?: AirbyteTraceMessage.Type.STREAM_STATUS) ==
+                        AirbyteTraceMessage.Type.STREAM_STATUS
+                ) {
                     when (status.status) {
                         AirbyteStreamStatus.COMPLETE ->
                             if (fileTransferEnabled) {
@@ -444,7 +450,7 @@ class DestinationMessageFactory(
                             } else {
                                 DestinationRecordStreamComplete(
                                     stream.descriptor,
-                                    message.trace.emittedAt.toLong()
+                                    message.trace.emittedAt?.toLong() ?: 0L
                                 )
                             }
                         AirbyteStreamStatus.INCOMPLETE ->
@@ -476,6 +482,7 @@ class DestinationMessageFactory(
                                 },
                             additionalProperties = message.state.additionalProperties,
                         )
+                    null,
                     AirbyteStateMessage.AirbyteStateType.GLOBAL ->
                         GlobalCheckpoint(
                             sourceStats =
@@ -488,6 +495,7 @@ class DestinationMessageFactory(
                                     fromAirbyteStreamState(it)
                                 },
                             additionalProperties = message.state.additionalProperties,
+                            originalTypeField = message.state.type,
                         )
                     else -> // TODO: Do we still need to handle LEGACY?
                     Undefined

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

@@ -25,7 +25,8 @@ class DefaultDestinationMessageDeserializer(private val messageFactory: Destinat
             val airbyteMessage = serialized.deserializeToClass(AirbyteMessage::class.java)
             return messageFactory.fromAirbyteMessage(airbyteMessage, serialized)
         } catch (t: Throwable) {
-            throw RuntimeException("Failed to deserialize AirbyteMessage")
+            // TODO: S3V2: Remove logging the message XXX
+            throw RuntimeException("Failed to deserialize AirbyteMessage: $serialized", t)
         }
     }
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/state/CheckpointManager.kt

@@ -158,6 +158,9 @@ abstract class StreamsCheckpointManager<T> : CheckpointManager<DestinationStream
                 validateAndSendMessage(head.checkpointMessage, head.streamIndexes)
                 globalCheckpoints.poll() // don't remove until after we've successfully sent
             } else {
+                log.info {
+                    "Not flushing global checkpoint with stream indexes: ${head.streamIndexes}"
+                }
                 break
             }
         }
@@ -176,6 +179,7 @@ abstract class StreamsCheckpointManager<T> : CheckpointManager<DestinationStream
                     validateAndSendMessage(nextMessage, listOf(stream.descriptor to nextIndex))
                     streamCheckpoints.poll() // don't remove until after we've successfully sent
                 } else {
+                    log.info { "Not flushing next checkpoint for index $nextIndex" }
                     break
                 }
             }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/state/StreamManager.kt

@@ -173,6 +173,10 @@ class DefaultStreamManager(
     private fun isProcessingCompleteForState(index: Long, state: Batch.State): Boolean {
         val completeRanges = rangesState[state]!!
 
+        if (index == 0L && recordCount.get() == 0L) {
+            return true
+        }
+
         return completeRanges.encloses(Range.closedOpen(0L, index))
     }
 

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/state/SyncManager.kt

@@ -13,6 +13,7 @@ import io.micronaut.context.annotation.Secondary
 import jakarta.inject.Singleton
 import java.util.concurrent.ConcurrentHashMap
 import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.withTimeoutOrNull
 
 sealed interface SyncResult
 
@@ -79,7 +80,8 @@ class DefaultSyncManager(
         stream: DestinationStream.Descriptor
     ): StreamLoader? {
         val completable = streamLoaders[stream]
-        return completable?.let { if (it.isCompleted) it.await() else null }
+        // `.isCompleted` does not work as expected here.
+        return completable?.let { withTimeoutOrNull(1000L) { it.await() } }
     }
 
     override suspend fun awaitAllStreamsCompletedSuccessfully(): Boolean {

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

@@ -130,7 +130,7 @@ T : ScopedTask {
                 log.warn { "Stream task $innerTask was cancelled." }
                 throw e
             } catch (e: Exception) {
-                log.error { "Caught exception in sync task $innerTask: $e" }
+                log.error { "Caught exception in stream task $innerTask: $e" }
                 handleStreamFailure(stream, e)
             }
         }

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

@@ -36,10 +36,12 @@ class DefaultUpdateCheckpointsTask(
             when (it.value) {
                 is StreamCheckpointWrapped -> {
                     val (_, stream, index, message) = it.value
+                    log.info { "Updating checkpoint for stream $stream with index $index" }
                     checkpointManager.addStreamCheckpoint(stream, index, it.replace(message))
                 }
                 is GlobalCheckpointWrapped -> {
                     val (_, streamIndexes, message) = it.value
+                    log.info { "Updating global checkpoint for streams $streamIndexes" }
                     checkpointManager.addGlobalCheckpoint(streamIndexes, it.replace(message))
                 }
             }

airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/test/util/RecordDifferTest.kt

@@ -5,6 +5,7 @@
 package io.airbyte.cdk.load.test.util
 
 import java.time.OffsetDateTime
+import kotlin.test.assertEquals
 import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.Test
 
@@ -118,4 +119,118 @@ class RecordDifferTest {
             diff
         )
     }
+
+    /** Verify that the differ can sort records which are identical other than the PK */
+    @Test
+    fun testSortPk() {
+        val diff =
+            RecordDiffer(primaryKey = listOf(listOf("id")), cursor = null)
+                .diffRecords(
+                    listOf(
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "foo"),
+                            airbyteMeta = null,
+                        ),
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 2, "name" to "bar"),
+                            airbyteMeta = null,
+                        ),
+                    ),
+                    listOf(
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 2, "name" to "bar"),
+                            airbyteMeta = null,
+                        ),
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "foo"),
+                            airbyteMeta = null,
+                        ),
+                    )
+                )
+        assertEquals(null, diff)
+    }
+
+    /** Verify that the differ can sort records which are identical other than the cursor */
+    @Test
+    fun testSortCursor() {
+        val diff =
+            RecordDiffer(primaryKey = listOf(listOf("id")), cursor = listOf("updated_at"))
+                .diffRecords(
+                    listOf(
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "foo", "updated_at" to 1),
+                            airbyteMeta = null,
+                        ),
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "bar", "updated_at" to 2),
+                            airbyteMeta = null,
+                        ),
+                    ),
+                    listOf(
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "bar", "updated_at" to 2),
+                            airbyteMeta = null,
+                        ),
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "foo", "updated_at" to 1),
+                            airbyteMeta = null,
+                        ),
+                    ),
+                )
+        assertEquals(null, diff)
+    }
+
+    /** Verify that the differ can sort records which are identical other than extractedAt */
+    @Test
+    fun testSortExtractedAt() {
+        val diff =
+            RecordDiffer(primaryKey = listOf(listOf("id")), cursor = null)
+                .diffRecords(
+                    listOf(
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "foo"),
+                            airbyteMeta = null,
+                        ),
+                        OutputRecord(
+                            extractedAt = 1,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "bar"),
+                            airbyteMeta = null,
+                        ),
+                    ),
+                    listOf(
+                        OutputRecord(
+                            extractedAt = 1,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "bar"),
+                            airbyteMeta = null,
+                        ),
+                        OutputRecord(
+                            extractedAt = 0,
+                            generationId = 0,
+                            data = mapOf("id" to 1, "name" to "foo"),
+                            airbyteMeta = null,
+                        ),
+                    )
+                )
+        assertEquals(null, diff)
+    }
 }

airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/test/util/StubDestinationMessageFactory.kt

@@ -73,6 +73,7 @@ object StubDestinationMessageFactory {
             sourceStats = CheckpointMessage.Stats(recordCount),
             checkpoints = emptyList(),
             additionalProperties = emptyMap(),
+            originalTypeField = message.state.type,
         )
     }