destination-s3: add file transfer (#46302)

Co-authored-by: benmoriceau <benoit@airbyte.io>

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

@@ -8,6 +8,7 @@ import io.airbyte.cdk.load.test.util.NoopDestinationCleaner
 import io.airbyte.cdk.load.test.util.NoopExpectedRecordMapper
 import io.airbyte.cdk.load.test.util.NoopNameMapper
 import io.airbyte.cdk.load.write.BasicFunctionalityIntegrationTest
+import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
 
 class MockBasicFunctionalityIntegrationTest :
@@ -27,6 +28,7 @@ class MockBasicFunctionalityIntegrationTest :
     }
 
     @Test
+    @Disabled
     override fun testMidSyncCheckpointingStreamState() {
         super.testMidSyncCheckpointingStreamState()
     }

airbyte-cdk/java/airbyte-cdk/README.md

@@ -174,6 +174,7 @@ corresponds to that version.
 
 | Version    | Date       | Pull Request                                                | Subject                                                                                                                                                        |
 |:-----------|:-----------|:------------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| 0.48.0     | 2024-10-23 | [\#46302](https://github.com/airbytehq/airbyte/pull/46302)  | Add support for file transfer |
 | 0.47.3     | 2024-10-23 | [\#46689](https://github.com/airbytehq/airbyte/pull/46689)  | Split DestinationAcceptanceTest|
 | 0.47.2     | 2024-10-21 | [\#47216](https://github.com/airbytehq/airbyte/pull/47216)  | improve java compatibiilty|
 | 0.47.1     | 2024-09-27 | [\#45397](https://github.com/airbytehq/airbyte/pull/45397)  | Allow logical replication from Postgres 16 read-replicas|

airbyte-cdk/java/airbyte-cdk/core/src/main/kotlin/io/airbyte/cdk/integrations/destination/async/AsyncStreamConsumer.kt

@@ -54,6 +54,7 @@ constructor(
     workerPool: ExecutorService = Executors.newFixedThreadPool(5),
     private val airbyteMessageDeserializer: AirbyteMessageDeserializer =
         AirbyteMessageDeserializer(),
+    flushOnEveryMessage: Boolean = false,
 ) : SerializedAirbyteMessageConsumer {
     private val bufferEnqueue: BufferEnqueue = bufferManager.bufferEnqueue
     private val flushWorkers: FlushWorkers =
@@ -64,6 +65,7 @@ constructor(
             flushFailure,
             bufferManager.stateManager,
             workerPool,
+            flushOnEveryMessage,
         )
     private val streamNames: Set<StreamDescriptor> =
         StreamDescriptorUtils.fromConfiguredCatalog(

airbyte-cdk/java/airbyte-cdk/core/src/main/kotlin/io/airbyte/cdk/integrations/destination/async/DetectStreamToFlush.kt

@@ -28,6 +28,7 @@ internal constructor(
     private val isClosing: AtomicBoolean,
     private val flusher: DestinationFlushFunction,
     private val nowProvider: Clock,
+    private val flushOnEveryMessage: Boolean = false,
 ) {
     private val latestFlushTimeMsPerStream: ConcurrentMap<StreamDescriptor, Long> =
         ConcurrentHashMap()
@@ -37,7 +38,15 @@ internal constructor(
         runningFlushWorkers: RunningFlushWorkers,
         isClosing: AtomicBoolean,
         flusher: DestinationFlushFunction,
-    ) : this(bufferDequeue, runningFlushWorkers, isClosing, flusher, Clock.systemUTC())
+        flushOnEveryMessage: Boolean = false,
+    ) : this(
+        bufferDequeue,
+        runningFlushWorkers,
+        isClosing,
+        flusher,
+        Clock.systemUTC(),
+        flushOnEveryMessage
+    )
 
     val nextStreamToFlush: Optional<StreamDescriptor>
         /**
@@ -70,7 +79,8 @@ internal constructor(
                 bufferDequeue.totalGlobalQueueSizeBytes.toDouble() / bufferDequeue.maxQueueSizeBytes
         // when we are closing or queues are very full, flush regardless of how few items are in the
         // queue.
-        return if (isClosing.get() || isBuffer90Full) 0 else flusher.queueFlushThresholdBytes
+        return if (flushOnEveryMessage || isClosing.get() || isBuffer90Full) 0
+        else flusher.queueFlushThresholdBytes
     }
 
     // todo (cgardens) - improve prioritization by getting a better estimate of how much data
@@ -105,7 +115,7 @@ internal constructor(
                     "${isTimeTriggeredResult.second} , ${isSizeTriggeredResult.second}"
             logger.debug { "computed: $debugString" }
 
-            if (isSizeTriggeredResult.first || isTimeTriggeredResult.first) {
+            if (flushOnEveryMessage || isSizeTriggeredResult.first || isTimeTriggeredResult.first) {
                 logger.info { "flushing: $debugString" }
                 latestFlushTimeMsPerStream[stream] = nowProvider.millis()
                 return Optional.of(stream)

airbyte-cdk/java/airbyte-cdk/core/src/main/kotlin/io/airbyte/cdk/integrations/destination/async/FlushWorkers.kt

@@ -51,6 +51,7 @@ constructor(
     private val flushFailure: FlushFailure,
     private val stateManager: GlobalAsyncStateManager,
     private val workerPool: ExecutorService = Executors.newFixedThreadPool(5),
+    flushOnEveryMessage: Boolean = false,
 ) : AutoCloseable {
     private val supervisorThread: ScheduledExecutorService = Executors.newScheduledThreadPool(1)
     private val debugLoop: ScheduledExecutorService = Executors.newSingleThreadScheduledExecutor()
@@ -66,6 +67,7 @@ constructor(
                 runningFlushWorkers,
                 isClosing,
                 flusher,
+                flushOnEveryMessage,
             )
     }
 

airbyte-cdk/java/airbyte-cdk/core/src/main/kotlin/io/airbyte/cdk/integrations/destination/async/buffers/BufferDequeue.kt

@@ -63,7 +63,7 @@ class BufferDequeue(
 
                 // otherwise pull records until we hit the memory limit.
                 val newSize: Long = (memoryItem.size) + bytesRead.get()
-                if (newSize <= optimalBytesToRead) {
+                if (newSize <= optimalBytesToRead || output.isEmpty()) {
                     memoryItem.size.let { bytesRead.addAndGet(it) }
                     queue.poll()?.item?.let { output.add(it) }
                 } else {

airbyte-cdk/java/airbyte-cdk/core/src/main/kotlin/io/airbyte/cdk/integrations/destination/async/model/AirbyteRecordMessageFile.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.integrations.destination.async.model
+
+import com.fasterxml.jackson.annotation.JsonProperty
+
+class AirbyteRecordMessageFile {
+    constructor(
+        fileUrl: String? = null,
+        bytes: Long? = null,
+        fileRelativePath: String? = null,
+        modified: Long? = null,
+        sourceFileUrl: String? = null
+    ) {
+        this.fileUrl = fileUrl
+        this.bytes = bytes
+        this.fileRelativePath = fileRelativePath
+        this.modified = modified
+        this.sourceFileUrl = sourceFileUrl
+    }
+    constructor() :
+        this(
+            fileUrl = null,
+            bytes = null,
+            fileRelativePath = null,
+            modified = null,
+            sourceFileUrl = null
+        )
+
+    @get:JsonProperty("file_url")
+    @set:JsonProperty("file_url")
+    @JsonProperty("file_url")
+    var fileUrl: String? = null
+
+    @get:JsonProperty("bytes")
+    @set:JsonProperty("bytes")
+    @JsonProperty("bytes")
+    var bytes: Long? = null
+
+    @get:JsonProperty("file_relative_path")
+    @set:JsonProperty("file_relative_path")
+    @JsonProperty("file_relative_path")
+    var fileRelativePath: String? = null
+
+    @get:JsonProperty("modified")
+    @set:JsonProperty("modified")
+    @JsonProperty("modified")
+    var modified: Long? = null
+
+    @get:JsonProperty("source_file_url")
+    @set:JsonProperty("source_file_url")
+    @JsonProperty("source_file_url")
+    var sourceFileUrl: String? = null
+}

airbyte-cdk/java/airbyte-cdk/core/src/main/kotlin/io/airbyte/cdk/integrations/destination/async/model/PartialAirbyteRecordMessage.kt

@@ -5,7 +5,6 @@
 package io.airbyte.cdk.integrations.destination.async.model
 
 import com.fasterxml.jackson.annotation.JsonProperty
-import com.fasterxml.jackson.annotation.JsonPropertyDescription
 import com.fasterxml.jackson.databind.JsonNode
 import io.airbyte.protocol.models.v0.AirbyteRecordMessageMeta
 import io.airbyte.protocol.models.v0.StreamDescriptor
@@ -33,14 +32,18 @@ class PartialAirbyteRecordMessage {
     @get:JsonProperty("emitted_at")
     @set:JsonProperty("emitted_at")
     @JsonProperty("emitted_at")
-    @JsonPropertyDescription("when the data was emitted from the source. epoch in millisecond.")
     var emittedAt: Long = 0
 
     @get:JsonProperty("meta")
     @set:JsonProperty("meta")
     @JsonProperty("meta")
     var meta: AirbyteRecordMessageMeta? = null
 
+    @get:JsonProperty("file")
+    @set:JsonProperty("file")
+    @JsonProperty("file")
+    var file: AirbyteRecordMessageFile? = null
+
     fun withNamespace(namespace: String?): PartialAirbyteRecordMessage {
         this.namespace = namespace
         return this
@@ -66,6 +69,11 @@ class PartialAirbyteRecordMessage {
         return this
     }
 
+    fun withFile(file: AirbyteRecordMessageFile): PartialAirbyteRecordMessage {
+        this.file = file
+        return this
+    }
+
     override fun equals(other: Any?): Boolean {
         if (this === other) {
             return true
@@ -77,7 +85,8 @@ class PartialAirbyteRecordMessage {
         return namespace == that.namespace &&
             stream == that.stream &&
             emittedAt == that.emittedAt &&
-            meta == that.meta
+            meta == that.meta &&
+            file == that.file
     }
 
     override fun hashCode(): Int {
@@ -98,6 +107,9 @@ class PartialAirbyteRecordMessage {
             ", meta='" +
             meta +
             '\'' +
+            ", file='" +
+            file +
+            '\'' +
             '}'
     }
 

airbyte-cdk/java/airbyte-cdk/core/src/main/resources/version.properties

@@ -1 +1 @@
-version=0.47.3
+version=0.48.1

airbyte-cdk/java/airbyte-cdk/core/src/test/kotlin/io/airbyte/cdk/integrations/destination/async/AsyncStreamConsumerTest.kt

@@ -765,6 +765,9 @@ class AsyncStreamConsumerTest {
         val throwable =
             assertThrows(RuntimeException::class.java) { consumer.accept(retyped, retyped.length) }
         // Ensure that the offending data has been scrubbed from the error message
-        assertFalse(throwable.message!!.contains(offender))
+        assertFalse(
+            throwable.message!!.contains(offender),
+            "message should not contain the offender. Was ${throwable.message}"
+        )
     }
 }

airbyte-cdk/java/airbyte-cdk/core/src/test/kotlin/io/airbyte/cdk/integrations/destination/async/DetectStreamToFlushTest.kt

@@ -49,7 +49,13 @@ class DetectStreamToFlushTest {
             )
 
         val detect =
-            DetectStreamToFlush(bufferDequeue, runningFlushWorkers, AtomicBoolean(false), flusher)
+            DetectStreamToFlush(
+                bufferDequeue,
+                runningFlushWorkers,
+                AtomicBoolean(false),
+                flusher,
+                flushOnEveryMessage = false
+            )
         Assertions.assertEquals(Optional.empty<Any>(), detect.getNextStreamToFlush(0))
     }
 
@@ -66,7 +72,13 @@ class DetectStreamToFlushTest {
                 RunningFlushWorkers::class.java,
             )
         val detect =
-            DetectStreamToFlush(bufferDequeue, runningFlushWorkers, AtomicBoolean(false), flusher)
+            DetectStreamToFlush(
+                bufferDequeue,
+                runningFlushWorkers,
+                AtomicBoolean(false),
+                flusher,
+                flushOnEveryMessage = false
+            )
         // if above threshold, triggers
         Assertions.assertEquals(Optional.of(DESC1), detect.getNextStreamToFlush(0))
         // if below threshold, no trigger
@@ -94,10 +106,41 @@ class DetectStreamToFlushTest {
                 ),
             )
         val detect =
-            DetectStreamToFlush(bufferDequeue, runningFlushWorkers, AtomicBoolean(false), flusher)
+            DetectStreamToFlush(
+                bufferDequeue,
+                runningFlushWorkers,
+                AtomicBoolean(false),
+                flusher,
+                flushOnEveryMessage = false
+            )
         Assertions.assertEquals(Optional.empty<Any>(), detect.getNextStreamToFlush(0))
     }
 
+    @Test
+    internal fun testFileTransfer() {
+        val bufferDequeue =
+            Mockito.mock(
+                BufferDequeue::class.java,
+            )
+        Mockito.`when`(bufferDequeue.bufferedStreams).thenReturn(setOf(DESC1))
+        Mockito.`when`(bufferDequeue.getQueueSizeBytes(DESC1)).thenReturn(Optional.of(0L))
+        val runningFlushWorkers =
+            Mockito.mock(
+                RunningFlushWorkers::class.java,
+            )
+
+        val detect =
+            DetectStreamToFlush(
+                bufferDequeue,
+                runningFlushWorkers,
+                AtomicBoolean(false),
+                flusher,
+                flushOnEveryMessage = true
+            )
+        Assertions.assertEquals(0, detect.computeQueueThreshold())
+        Assertions.assertEquals(Optional.of(DESC1), detect.getNextStreamToFlush(0))
+    }
+
     @Test
     internal fun testGetNextPicksUpOnTimeTrigger() {
         val bufferDequeue =
@@ -127,6 +170,7 @@ class DetectStreamToFlushTest {
                 AtomicBoolean(false),
                 flusher,
                 mockedNowProvider,
+                flushOnEveryMessage = false
             )
 
         // initialize flush time

airbyte-cdk/java/airbyte-cdk/db-destinations/src/testFixtures/kotlin/io/airbyte/cdk/integrations/standardtest/destination/BaseDestinationAcceptanceTest.kt

@@ -14,6 +14,7 @@ import io.airbyte.configoss.WorkerDestinationConfig
 import io.airbyte.protocol.models.v0.AirbyteMessage
 import io.airbyte.protocol.models.v0.AirbyteStateStats
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog
+import io.airbyte.workers.exception.TestHarnessException
 import io.airbyte.workers.helper.ConnectorConfigUpdater
 import io.airbyte.workers.internal.AirbyteDestination
 import io.airbyte.workers.internal.DefaultAirbyteDestination
@@ -215,7 +216,11 @@ abstract class BaseDestinationAcceptanceTest(
             }
         }
 
-        destination.close()
+        try {
+            destination.close()
+        } catch (e: TestHarnessException) {
+            throw TestHarnessException(e.message, e, destinationOutput)
+        }
 
         return destinationOutput
     }
@@ -258,6 +263,7 @@ abstract class BaseDestinationAcceptanceTest(
                 workspaceRoot,
                 workspaceRoot.toString(),
                 localRoot.toString(),
+                fileTransferMountSource,
                 "host",
                 getConnectorEnv()
             )

airbyte-cdk/java/airbyte-cdk/db-sources/src/testFixtures/kotlin/io/airbyte/cdk/integrations/standardtest/source/AbstractSourceConnectorTest.kt

@@ -117,6 +117,7 @@ abstract class AbstractSourceConnectorTest {
                 workspaceRoot,
                 workspaceRoot.toString(),
                 localRoot.toString(),
+                fileTransferMountSource = null,
                 "host",
                 envMap
             )

airbyte-cdk/java/airbyte-cdk/dependencies/src/main/kotlin/io/airbyte/commons/features/EnvVariableFeatureFlags.kt

@@ -4,6 +4,7 @@
 package io.airbyte.commons.features
 
 import io.github.oshai.kotlinlogging.KotlinLogging
+import java.nio.file.Path
 import java.util.function.Function
 
 private val log = KotlinLogging.logger {}
@@ -46,6 +47,16 @@ class EnvVariableFeatureFlags : FeatureFlags {
         return getEnvOrDefault(DEPLOYMENT_MODE, "") { arg: String -> arg }
     }
 
+    override fun airbyteStagingDirectory(): Path? {
+        return getEnvOrDefault(AIRBYTE_STAGING_DIRECTORY_PROPERTY_NAME, null) { arg: String ->
+            Path.of(arg)
+        }
+    }
+
+    override fun useFileTransfer(): Boolean {
+        return getEnvOrDefault(USE_FILE_TRANSFER, false) { it.toBoolean() }
+    }
+
     // TODO: refactor in order to use the same method than the ones in EnvConfigs.java
     fun <T> getEnvOrDefault(key: String?, defaultValue: T, parser: Function<String, T>): T {
         val value = System.getenv(key)
@@ -73,5 +84,8 @@ class EnvVariableFeatureFlags : FeatureFlags {
         const val STRICT_COMPARISON_NORMALIZATION_TAG: String =
             "STRICT_COMPARISON_NORMALIZATION_TAG"
         const val DEPLOYMENT_MODE: String = "DEPLOYMENT_MODE"
+        val DEFAULT_AIRBYTE_STAGING_DIRECTORY: Path = Path.of("/staging/files")
+        const val AIRBYTE_STAGING_DIRECTORY_PROPERTY_NAME: String = "AIRBYTE_STAGING_DIRECTORY"
+        const val USE_FILE_TRANSFER = "USE_FILE_TRANSFER"
     }
 }