Bulk Load CDK: Fixes for various S3V2 release issues (#48733)

airbytehq · Nov 28, 2024 · 4add9ca · 4add9ca
1 parent 2337025
commit 4add9ca
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 8 deletions.
diff --git a/...k/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/json/JsonSchemaToAirbyteType.kt b/...k/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/json/JsonSchemaToAirbyteType.kt
@@ -22,6 +22,7 @@ class JsonSchemaToAirbyteType {
                 when (schema.get("type").asText()) {
                     "string" -> fromString(schema)
                     "boolean" -> BooleanType
+                    "int",
                     "integer" -> IntegerType
                     "number" -> fromNumber(schema)
                     "array" -> fromArray(schema)

diff --git a/...k/core/load/src/main/kotlin/io/airbyte/cdk/load/message/DestinationMessageDeserializer.kt b/...k/core/load/src/main/kotlin/io/airbyte/cdk/load/message/DestinationMessageDeserializer.kt
@@ -25,7 +25,28 @@ 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")
+            /**
+             * We don't want to expose client data, but we'd like to get as much info as we can
+             * about these malformed messages.
+             */
+            val type =
+                if (serialized.contains("RECORD")) {
+                    "record"
+                } else if (serialized.contains("STATE")) {
+                    "state"
+                } else if (serialized.contains("TRACE")) {
+                    if (serialized.contains("STATUS", ignoreCase = true)) {
+                        "status"
+                    } else {
+                        "trace"
+                    }
+                } else {
+                    "unknown"
+                }
+
+            throw RuntimeException(
+                "Failed to deserialize airbyte message (type=$type; length=${serialized.length}; reason=${t.javaClass})"
+            )
         }
     }
 }
diff --git a/...e/load/src/test/kotlin/io/airbyte/cdk/load/data/json/JsonSchemaToAirbyteSchemaTypeTest.kt b/...e/load/src/test/kotlin/io/airbyte/cdk/load/data/json/JsonSchemaToAirbyteSchemaTypeTest.kt
@@ -52,6 +52,14 @@ class JsonSchemaToAirbyteSchemaTypeTest {
         Assertions.assertTrue(airbyteType is IntegerType)
     }
 
+    /** Note: this is nonstandard, but some sources apparently use it. */
+    @Test
+    fun testInt() {
+        val integerType = ofType("int")
+        val airbyteType = JsonSchemaToAirbyteType().convert(integerType)
+        Assertions.assertTrue(airbyteType is IntegerType)
+    }
+
     @Test
     fun testNumber() {
         val numberType = ofType("number")

diff --git a/...orage/src/main/kotlin/io/airbyte/cdk/load/file/object_storage/ObjectStoragePathFactory.kt b/...orage/src/main/kotlin/io/airbyte/cdk/load/file/object_storage/ObjectStoragePathFactory.kt
@@ -223,6 +223,9 @@ class ObjectStoragePathFactory(
                 FileVariable("date", """\d{4}_\d{2}_\d{2}""") {
                     DATE_FORMATTER.format(it.syncTime)
                 },
+                FileVariable("date:yyyy_MM", """\d{4}_\d{2}""") {
+                    DATE_FORMATTER.format(it.syncTime).substring(0, 7)
+                },
                 FileVariable("timestamp", """\d+""") {
                     // NOTE: We use a constant time for the path but wall time for the files
                     it.currentTimeProvider.currentTimeMillis().toString()
@@ -381,7 +384,7 @@ class ObjectStoragePathFactory(
         val replacedForFile =
             buildPattern(
                 filePatternResolved,
-                """\{(\w+)}""",
+                """\{([\w\:]+)}""",
                 pathVariableToPattern,
                 variableToIndex
             )

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
@@ -60,7 +60,8 @@ class ObjectStoragePathFactoryTest {
                 stagingPrefix = "staging/prefix",
                 pathSuffixPattern =
                     "\${NAMESPACE}/\${STREAM_NAME}/\${YEAR}/\${MONTH}/\${DAY}/\${HOUR}/\${MINUTE}/\${SECOND}/\${MILLISECOND}/\${EPOCH}/",
-                fileNamePattern = "{date}-{timestamp}-{part_number}-{sync_id}{format_extension}",
+                fileNamePattern =
+                    "{date}-{date:yyyy_MM}-{timestamp}-{part_number}-{sync_id}{format_extension}",
                 usesStagingDirectory = true
             )
     }
@@ -125,7 +126,7 @@ class ObjectStoragePathFactoryTest {
             val stream1 = MockDestinationCatalogFactory.stream1
             val (namespace, name) = stream1.descriptor
             val prefixOnly = "prefix/$namespace/$name/2020/01/02/03/04/05/0678/$syncTime/"
-            val fileName = "2020_01_02-$wallTime-173-42.jsonl.gz"
+            val fileName = "2020_01_02-2020_01-$wallTime-173-42.jsonl.gz"
             Assertions.assertEquals(
                 "staging/$prefixOnly",
                 pathFactory.getStagingDirectory(stream1).toString(),
@@ -153,7 +154,7 @@ class ObjectStoragePathFactoryTest {
             val stream1 = MockDestinationCatalogFactory.stream1
             val (namespace, name) = stream1.descriptor
             val expectedToMatch =
-                "prefix/$namespace/$name/2020/01/02/03/04/05/0678/$syncTime/2020_01_02-1577934245678-173-42.jsonl.gz"
+                "prefix/$namespace/$name/2020/01/02/03/04/05/0678/$syncTime/2020_01_02-2020_01-1577934245678-173-42.jsonl.gz"
             val match = pathFactory.getPathMatcher(stream1).match(expectedToMatch)
             Assertions.assertTrue(match != null)
             Assertions.assertTrue(match?.partNumber == 173L)
@@ -170,7 +171,7 @@ class ObjectStoragePathFactoryTest {
             val emptyNamespaceStream =
                 stream1.copy(descriptor = stream1.descriptor.copy(namespace = null))
             val expectedToMatch =
-                "prefix/$name/2020/01/02/03/04/05/0678/$epochMilli/2020_01_02-1577934245678-173-42.jsonl.gz"
+                "prefix/$name/2020/01/02/03/04/05/0678/$epochMilli/2020_01_02-2020_01-1577934245678-173-42.jsonl.gz"
             val match = pathFactory.getPathMatcher(emptyNamespaceStream).match(expectedToMatch)
             Assertions.assertTrue(match != null)
             Assertions.assertTrue(match?.partNumber == 173L)
@@ -231,7 +232,7 @@ class ObjectStoragePathFactoryTest {
             val stream1 = MockDestinationCatalogFactory.stream1
             val (namespace, name) = stream1.descriptor
             val prefixOnly = "prefix/$namespace/$name/2020/01/02/03/04/05/0678/$syncTime/"
-            val fileName = "2020_01_02-$wallTime-173-42.jsonl.gz"
+            val fileName = "2020_01_02-2020_01-$wallTime-173-42.jsonl.gz"
             Assertions.assertEquals(
                 prefixOnly,
                 pathFactory.getFinalDirectory(stream1),
@@ -270,7 +271,7 @@ class ObjectStoragePathFactoryTest {
             val stream1 = MockDestinationCatalogFactory.stream1
             val (namespace, name) = stream1.descriptor
             val prefixOnly = "prefix/$namespace/$name/2020/01/02/03/04/05/0678/${syncTime}_"
-            val fileName = "2020_01_02-$wallTime-173-42.jsonl.gz"
+            val fileName = "2020_01_02-2020_01-$wallTime-173-42.jsonl.gz"
             Assertions.assertEquals(
                 prefixOnly,
                 pathFactory.getFinalDirectory(stream1),