Bulk Load CDK: S3V2: CSV Printer config matches old cdk; await state … (

#48527)
airbytehq · Nov 18, 2024 · 8901591 · 8901591
1 parent d705fe9
commit 8901591
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 13 deletions.
diff --git a/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/file/StreamProcessor.kt b/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/file/StreamProcessor.kt
@@ -4,6 +4,7 @@
 
 package io.airbyte.cdk.load.file
 
+import java.io.BufferedOutputStream
 import java.io.ByteArrayOutputStream
 import java.io.OutputStream
 import java.util.zip.GZIPOutputStream
@@ -20,8 +21,10 @@ data object NoopProcessor : StreamProcessor<ByteArrayOutputStream> {
     override val extension: String? = null
 }
 
-data object GZIPProcessor : StreamProcessor<GZIPOutputStream> {
-    override val wrapper: (ByteArrayOutputStream) -> GZIPOutputStream = { GZIPOutputStream(it) }
-    override val partFinisher: GZIPOutputStream.() -> Unit = { finish() }
+data object GZIPProcessor : StreamProcessor<BufferedOutputStream> {
+    override val wrapper: (ByteArrayOutputStream) -> BufferedOutputStream = {
+        BufferedOutputStream(GZIPOutputStream(it))
+    }
+    override val partFinisher: BufferedOutputStream.() -> Unit = { close() }
     override val extension: String = "gz"
 }
diff --git a/...e-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/task/implementor/TeardownTask.kt b/...e-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/task/implementor/TeardownTask.kt
@@ -31,14 +31,14 @@ class DefaultTeardownTask(
 
     override suspend fun execute() {
         syncManager.awaitInputProcessingComplete()
-        checkpointManager.awaitAllCheckpointsFlushed()
 
         log.info { "Teardown task awaiting stream completion" }
         if (!syncManager.awaitAllStreamsCompletedSuccessfully()) {
             log.info { "Streams failed to complete successfully, doing nothing." }
             return
         }
 
+        checkpointManager.awaitAllCheckpointsFlushed()
         log.info { "Starting teardown task" }
         destination.teardown()
         log.info { "Teardown task complete, marking sync succeeded." }

diff --git a/airbyte-cdk/bulk/toolkits/load-csv/src/main/kotlin/io/airbyte/cdk/load/file/csv/CSVWriter.kt b/airbyte-cdk/bulk/toolkits/load-csv/src/main/kotlin/io/airbyte/cdk/load/file/csv/CSVWriter.kt
@@ -7,12 +7,15 @@ package io.airbyte.cdk.load.file.csv
 import io.airbyte.cdk.load.data.ObjectType
 import io.airbyte.cdk.load.data.csv.toCsvHeader
 import java.io.OutputStream
+import java.io.PrintWriter
+import java.nio.charset.StandardCharsets
 import org.apache.commons.csv.CSVFormat
 import org.apache.commons.csv.CSVPrinter
+import org.apache.commons.csv.QuoteMode
 
-fun ObjectType.toCsvPrinterWithHeader(outputStream: OutputStream): CSVPrinter =
-    CSVFormat.Builder.create()
-        .setHeader(*toCsvHeader())
-        .setAutoFlush(true)
-        .build()
-        .print(outputStream.writer(charset = Charsets.UTF_8))
+@Suppress("DEPRECATION")
+fun ObjectType.toCsvPrinterWithHeader(outputStream: OutputStream): CSVPrinter {
+    val csvSettings =
+        CSVFormat.DEFAULT.withQuoteMode(QuoteMode.NON_NUMERIC).withHeader(*toCsvHeader())
+    return CSVPrinter(PrintWriter(outputStream, true, StandardCharsets.UTF_8), csvSettings)
+}
diff --git a/...e/src/test/kotlin/io/airbyte/cdk/load/file/object_storage/ObjectStoragePathFactoryTest.kt b/...e/src/test/kotlin/io/airbyte/cdk/load/file/object_storage/ObjectStoragePathFactoryTest.kt
@@ -20,10 +20,10 @@ import io.micronaut.context.annotation.Property
 import io.micronaut.context.annotation.Requires
 import io.micronaut.test.extensions.junit5.annotation.MicronautTest
 import jakarta.inject.Singleton
+import java.io.BufferedOutputStream
 import java.time.LocalDateTime
 import java.time.ZoneId
 import java.time.format.DateTimeFormatter
-import java.util.zip.GZIPOutputStream
 import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.Nested
 import org.junit.jupiter.api.Test
@@ -84,9 +84,9 @@ class ObjectStoragePathFactoryTest {
     @Primary
     @Requires(env = ["ObjectStoragePathFactoryTest"])
     class MockCompressionConfigProvider :
-        ObjectStorageCompressionConfigurationProvider<GZIPOutputStream> {
+        ObjectStorageCompressionConfigurationProvider<BufferedOutputStream> {
         override val objectStorageCompressionConfiguration:
-            ObjectStorageCompressionConfiguration<GZIPOutputStream> =
+            ObjectStorageCompressionConfiguration<BufferedOutputStream> =
             ObjectStorageCompressionConfiguration(compressor = GZIPProcessor)
     }