Split into bucketPath and fullOutputPath

ChristopheDuong · ChristopheDuong · commit 665e6db3ddab · 2022-04-04T11:49:48.000+02:00
diff --git a/airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/destination/record_buffer/FileBuffer.java b/airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/destination/record_buffer/FileBuffer.java
@@ -12,6 +12,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.file.Files;
+import java.util.UUID;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -46,7 +47,7 @@ public FileBuffer(final String fileExtension) {
   @Override
   public OutputStream getOutputStream() throws IOException {
     if (outputStream == null || tempFile == null) {
-      tempFile = Files.createTempFile("", fileExtension).toFile();
+      tempFile = Files.createTempFile(UUID.randomUUID().toString(), fileExtension).toFile();
       outputStream = new BufferedOutputStream(new FileOutputStream(tempFile));
     }
     return outputStream;
diff --git a/airbyte-integrations/connectors/destination-s3/src/main/java/io/airbyte/integrations/destination/s3/S3ConsumerFactory.java b/airbyte-integrations/connectors/destination-s3/src/main/java/io/airbyte/integrations/destination/s3/S3ConsumerFactory.java
@@ -80,14 +80,14 @@ private static Function<ConfiguredAirbyteStream, WriteConfig> toWriteConfig(
       final AirbyteStream abStream = stream.getStream();
       final String namespace = abStream.getNamespace();
       final String streamName = abStream.getName();
-      final String bucketPath = config.get(BUCKET_PATH_FIELD).asText();
+      final String outputBucketPath = config.get(BUCKET_PATH_FIELD).asText();
       final String customOutputFormat = String.join("/",
-          bucketPath,
+          outputBucketPath,
           config.has(PATH_FORMAT_FIELD) && !config.get(PATH_FORMAT_FIELD).asText().isBlank() ? config.get(PATH_FORMAT_FIELD).asText()
               : S3DestinationConstants.DEFAULT_PATH_FORMAT);
-      final String outputBucketPath = storageOperations.getBucketObjectPath(namespace, streamName, SYNC_DATETIME, customOutputFormat);
+      final String fullOutputPath = storageOperations.getBucketObjectPath(namespace, streamName, SYNC_DATETIME, customOutputFormat);
       final DestinationSyncMode syncMode = stream.getDestinationSyncMode();
-      final WriteConfig writeConfig = new WriteConfig(namespace, streamName, outputBucketPath, syncMode);
+      final WriteConfig writeConfig = new WriteConfig(namespace, streamName, outputBucketPath, fullOutputPath, syncMode);
       LOGGER.info("Write config: {}", writeConfig);
       return writeConfig;
     };
@@ -100,13 +100,13 @@ private OnStartFunction onStartFunction(final BlobStorageOperations storageOpera
         if (writeConfig.getSyncMode().equals(DestinationSyncMode.OVERWRITE)) {
           final String namespace = writeConfig.getNamespace();
           final String stream = writeConfig.getStreamName();
-          final String outputBucketPath = writeConfig.getOutputBucketPath();
-          LOGGER.info("Clearing storage area in destination started for namespace {} stream {} bucketObject {}", namespace, stream, outputBucketPath);
+          final String bucketPath = writeConfig.getOutputBucketPath();
+          LOGGER.info("Clearing storage area in destination started for namespace {} stream {} bucketObject {}", namespace, stream, bucketPath);
           AirbyteSentry.executeWithTracing("PrepareStreamStorage",
-              () -> storageOperations.dropBucketObject(outputBucketPath),
-              Map.of("namespace", Objects.requireNonNullElse(namespace, "null"), "stream", stream, "storage", outputBucketPath));
+              () -> storageOperations.dropBucketObject(bucketPath),
+              Map.of("namespace", Objects.requireNonNullElse(namespace, "null"), "stream", stream, "storage", bucketPath));
           LOGGER.info("Clearing storage area in destination completed for namespace {} stream {} bucketObject {}", namespace, stream,
-              outputBucketPath);
+              bucketPath);
         }
       }
       LOGGER.info("Preparing storage area in destination completed.");
@@ -139,7 +139,7 @@ private CheckedBiConsumer<AirbyteStreamNameNamespacePair, SerializableBuffer, Ex
             writer,
             writeConfig.getNamespace(),
             writeConfig.getStreamName(),
-            writeConfig.getOutputBucketPath()));
+            writeConfig.getFullOutputPath()));
       } catch (final Exception e) {
         LOGGER.error("Failed to flush and upload buffer to storage:", e);
         throw new RuntimeException("Failed to upload buffer to storage", e);
@@ -153,7 +153,7 @@ private OnCloseFunction onCloseFunction(final BlobStorageOperations storageOpera
       if (hasFailed) {
         LOGGER.info("Cleaning up destination started for {} streams", writeConfigs.size());
         for (final WriteConfig writeConfig : writeConfigs) {
-          storageOperations.cleanUpBucketObject(writeConfig.getOutputBucketPath(), writeConfig.getStoredFiles());
+          storageOperations.cleanUpBucketObject(writeConfig.getFullOutputPath(), writeConfig.getStoredFiles());
           writeConfig.clearStoredFiles();
         }
         LOGGER.info("Cleaning up destination completed.");
diff --git a/airbyte-integrations/connectors/destination-s3/src/main/java/io/airbyte/integrations/destination/s3/S3DestinationConstants.java b/airbyte-integrations/connectors/destination-s3/src/main/java/io/airbyte/integrations/destination/s3/S3DestinationConstants.java
@@ -12,7 +12,7 @@ public final class S3DestinationConstants {
   public static final S3NameTransformer NAME_TRANSFORMER = new S3NameTransformer();
   public static final String PART_SIZE_MB_ARG_NAME = "part_size_mb";
   public static final int DEFAULT_PART_SIZE_MB = 10;
-  public static final String DEFAULT_PATH_FORMAT = "${NAMESPACE}/${STREAM_NAME}/${YEAR}_${MONTH}_${DAY}_${EPOCH}_${PART_ID}";
+  public static final String DEFAULT_PATH_FORMAT = "${NAMESPACE}/${STREAM_NAME}/${YEAR}_${MONTH}_${DAY}_${EPOCH}_";
 
   private S3DestinationConstants() {}
 
diff --git a/airbyte-integrations/connectors/destination-s3/src/main/java/io/airbyte/integrations/destination/s3/S3StorageOperations.java b/airbyte-integrations/connectors/destination-s3/src/main/java/io/airbyte/integrations/destination/s3/S3StorageOperations.java
@@ -17,7 +17,6 @@
 import io.airbyte.integrations.destination.NamingConventionTransformer;
 import io.airbyte.integrations.destination.record_buffer.SerializableBuffer;
 import io.airbyte.integrations.destination.s3.util.StreamTransferManagerHelper;
-import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.List;
@@ -50,7 +49,7 @@ public S3StorageOperations(final NamingConventionTransformer nameTransformer, fi
   public String getBucketObjectPath(final String namespace, final String streamName, final DateTime writeDatetime, final String customPathFormat) {
     final String namespaceStr = nameTransformer.getNamespace(isNotBlank(namespace) ? namespace : "");
     final String streamNameStr = nameTransformer.getIdentifier(streamName);
-    final String objectPath = nameTransformer.applyDefaultCase(
+    return nameTransformer.applyDefaultCase(
         customPathFormat
             .replaceAll(Pattern.quote("${NAMESPACE}"), namespaceStr)
             .replaceAll(Pattern.quote("${STREAM_NAME}"), streamNameStr)
@@ -64,22 +63,6 @@ public String getBucketObjectPath(final String namespace, final String streamNam
             .replaceAll(Pattern.quote("${EPOCH}"), String.format("%d", writeDatetime.getMillis()))
             .replaceAll(Pattern.quote("${UUID}"), String.format("%s", UUID.randomUUID()))
             .replaceAll("/+", "/"));
-    if (customPathFormat.contains("${PART_ID}")) {
-      return objectPath.replaceAll(Pattern.quote("${PART_ID}"), String.format("%s", getPartId(objectPath)));
-    } else {
-      return objectPath;
-    }
-  }
-
-  private String getPartId(final String objectPath) {
-    final String bucket = s3Config.getBucketName();
-    final ObjectListing objects = s3Client.listObjects(bucket, objectPath);
-    if (objects.isTruncated()) {
-      // bucket contains too many objects, use an uuid instead
-      return UUID.randomUUID().toString();
-    } else {
-      return Integer.toString(objects.getObjectSummaries().size());
-    }
   }
 
   @Override
@@ -122,12 +105,12 @@ public String uploadRecordsToBucket(final SerializableBuffer recordsData, final
     return recordsData.getFilename();
   }
 
-  private void loadDataIntoBucket(final String objectPath, final SerializableBuffer recordsData) throws IOException {
+  private void loadDataIntoBucket(final String objectPath, final SerializableBuffer recordsData) {
     final long partSize = s3Config.getFormatConfig() != null ? s3Config.getFormatConfig().getPartSize() : DEFAULT_PART_SIZE;
     final String bucket = s3Config.getBucketName();
-    final String objectKey = String.format("%s%s", objectPath, recordsData.getFilename());
+    final String objectKeyWithPartId = String.format("%s%s", objectPath, getPartId(objectPath));
     final StreamTransferManager uploadManager = StreamTransferManagerHelper
-        .getDefault(bucket, objectKey, s3Client, partSize)
+        .getDefault(bucket, objectKeyWithPartId, s3Client, partSize)
         .checkIntegrity(true)
         .numUploadThreads(DEFAULT_UPLOAD_THREADS)
         .queueCapacity(DEFAULT_QUEUE_CAPACITY);
@@ -146,12 +129,23 @@ private void loadDataIntoBucket(final String objectPath, final SerializableBuffe
         uploadManager.complete();
       }
     }
-    if (!s3Client.doesObjectExist(bucket, objectKey)) {
-      LOGGER.error("Failed to upload data into storage, object {} not found", objectKey);
+    if (!s3Client.doesObjectExist(bucket, objectKeyWithPartId)) {
+      LOGGER.error("Failed to upload data into storage, object {} not found", objectKeyWithPartId);
       throw new RuntimeException("Upload failed");
     }
   }
 
+  private String getPartId(final String objectPath) {
+    final String bucket = s3Config.getBucketName();
+    final ObjectListing objects = s3Client.listObjects(bucket, objectPath);
+    if (objects.isTruncated()) {
+      // bucket contains too many objects, use an uuid instead
+      return UUID.randomUUID().toString();
+    } else {
+      return Integer.toString(objects.getObjectSummaries().size());
+    }
+  }
+
   @Override
   public void dropBucketObject(final String objectPath) {
     cleanUpBucketObject(objectPath, List.of());
diff --git a/airbyte-integrations/connectors/destination-s3/src/main/java/io/airbyte/integrations/destination/s3/WriteConfig.java b/airbyte-integrations/connectors/destination-s3/src/main/java/io/airbyte/integrations/destination/s3/WriteConfig.java
@@ -16,16 +16,19 @@ public class WriteConfig {
   private final String namespace;
   private final String streamName;
   private final String outputBucketPath;
+  private final String fullOutputPath;
   private final DestinationSyncMode syncMode;
   private final List<String> storedFiles;
 
   public WriteConfig(final String namespace,
                      final String streamName,
                      final String outputBucketPath,
+                     final String fullOutputPath,
                      final DestinationSyncMode syncMode) {
     this.namespace = namespace;
     this.streamName = streamName;
     this.outputBucketPath = outputBucketPath;
+    this.fullOutputPath = fullOutputPath;
     this.syncMode = syncMode;
     this.storedFiles = new ArrayList<>();
   }
@@ -42,6 +45,10 @@ public String getOutputBucketPath() {
     return outputBucketPath;
   }
 
+  public String getFullOutputPath() {
+    return fullOutputPath;
+  }
+
   public DestinationSyncMode getSyncMode() {
     return syncMode;
   }
@@ -64,6 +71,7 @@ public String toString() {
         "streamName=" + streamName +
         ", namespace=" + namespace +
         ", outputBucketPath=" + outputBucketPath +
+        ", fullOutputPath=" + fullOutputPath +
         ", syncMode=" + syncMode +
         '}';
   }
diff --git a/airbyte-integrations/connectors/destination-s3/src/test-integration/java/io/airbyte/integrations/destination/s3/S3DestinationAcceptanceTest.java b/airbyte-integrations/connectors/destination-s3/src/test-integration/java/io/airbyte/integrations/destination/s3/S3DestinationAcceptanceTest.java
@@ -100,8 +100,9 @@ protected List<S3ObjectSummary> getAllSyncedObjects(final String streamName, fin
         streamNameStr,
         DateTime.now(DateTimeZone.UTC),
         S3DestinationConstants.DEFAULT_PATH_FORMAT);
+    final String parentFolder = outputPrefix.substring(0, outputPrefix.lastIndexOf("/") + 1);
     final List<S3ObjectSummary> objectSummaries = s3Client
-        .listObjects(config.getBucketName(), outputPrefix)
+        .listObjects(config.getBucketName(), parentFolder)
         .getObjectSummaries()
         .stream()
         .filter(o -> o.getKey().contains(streamNameStr + "/"))
diff --git a/airbyte-integrations/connectors/destination-s3/src/test/java/io/airbyte/integrations/destination/s3/avro/AvroSerializedBufferTest.java b/airbyte-integrations/connectors/destination-s3/src/test/java/io/airbyte/integrations/destination/s3/avro/AvroSerializedBufferTest.java
@@ -55,7 +55,7 @@ public class AvroSerializedBufferTest {
   public void testSnappyAvroWriter() throws Exception {
     final S3AvroFormatConfig config = new S3AvroFormatConfig(Jsons.jsonNode(Map.of("compression_codec", Map.of(
         "codec", "snappy"))));
-    runTest(new InMemoryBuffer(AVRO_FILE_EXTENSION), 970L, 980L, config, getExpectedString());
+    runTest(new InMemoryBuffer(AVRO_FILE_EXTENSION), 965L, 980L, config, getExpectedString());
   }
 
   @Test

Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,7 @@ public class AvroSerializedBufferTest {`
`55`	`55`	`public void testSnappyAvroWriter() throws Exception {`
`56`	`56`	`final S3AvroFormatConfig config = new S3AvroFormatConfig(Jsons.jsonNode(Map.of("compression_codec", Map.of(`
`57`	`57`	`"codec", "snappy"))));`
`58`		`- runTest(new InMemoryBuffer(AVRO_FILE_EXTENSION), 970L, 980L, config, getExpectedString());`
	`58`	`+ runTest(new InMemoryBuffer(AVRO_FILE_EXTENSION), 965L, 980L, config, getExpectedString());`
`59`	`59`	`}`
`60`	`60`
`61`	`61`	`@Test`