Destination Snowflake Execute COPY in parallel

airbyte-integrations/connectors/destination-snowflake/src/main/java/io/airbyte/integrations/destination/snowflake/SnowflakeGcsStreamCopier.java

@@ -19,40 +19,55 @@
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.StringJoiner;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class SnowflakeGcsStreamCopier extends GcsStreamCopier {
+public class SnowflakeGcsStreamCopier extends GcsStreamCopier implements SnowflakeParallelCopyStreamCopier {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeGcsStreamCopier.class);
+  private final ExecutorService executorService = Executors.newFixedThreadPool(5);
 
   public SnowflakeGcsStreamCopier(final String stagingFolder,
-                                  final DestinationSyncMode destSyncMode,
-                                  final String schema,
-                                  final String streamName,
-                                  final Storage storageClient,
-                                  final JdbcDatabase db,
-                                  final GcsConfig gcsConfig,
-                                  final ExtendedNameTransformer nameTransformer,
-                                  final SqlOperations sqlOperations,
-                                  final StagingFilenameGenerator stagingFilenameGenerator) {
+      final DestinationSyncMode destSyncMode,
+      final String schema,
+      final String streamName,
+      final Storage storageClient,
+      final JdbcDatabase db,
+      final GcsConfig gcsConfig,
+      final ExtendedNameTransformer nameTransformer,
+      final SqlOperations sqlOperations,
+      final StagingFilenameGenerator stagingFilenameGenerator) {
     super(stagingFolder, destSyncMode, schema, streamName, storageClient, db, gcsConfig, nameTransformer, sqlOperations);
     this.filenameGenerator = stagingFilenameGenerator;
   }
 
   @Override
   public void copyStagingFileToTemporaryTable() throws Exception {
-    List<List<String>> partition = Lists.partition(new ArrayList<>(gcsStagingFiles), MAX_FILES_PER_COPY);
-    for (int i = 0; i < partition.size(); i++) {
-      LOGGER.info("Starting copy chunk {} to tmp table: {} in destination for stream: {}, schema: {}. Chunks count {}", i, tmpTableName, streamName,
-          schemaName, partition.size());
-      executeCopy(partition.get(i));
-    }
+    List<List<String>> partitions = Lists.partition(new ArrayList<>(gcsStagingFiles), MAX_FILES_PER_COPY);
+    LOGGER.info("Starting parallel imt copy to tmp table: {} in destination for stream: {}, schema: {}. Chunks count {}", tmpTableName, streamName,
+        schemaName, partitions.size());
+
+    copyFilesInParallel(partitions);
     LOGGER.info("Copy to tmp table {} in destination for stream {} complete.", tmpTableName, streamName);
   }
 
-  private void executeCopy(List<String> files) {
+  private void copyFilesInParallel(List<List<String>> partitions) {
+    partitions.forEach(files -> {
+      try {
+        executorService.execute(() -> copyIntoStage(files));
+      } catch (Exception e) {
+        LOGGER.error("Failed to copy files from stage to tmp table {}", tmpTableName, e);
+        throw new RuntimeException(e);
+      } finally {
+        executorService.shutdown();
+      }
+    });
+  }
+
+  @Override
+  public void copyIntoStage(List<String> files) {
 
     final var copyQuery = String.format(
         "COPY INTO %s.%s FROM '%s' storage_integration = gcs_airbyte_integration "
@@ -65,22 +80,17 @@ private void executeCopy(List<String> files) {
     Exceptions.toRuntime(() -> db.execute(copyQuery));
   }
 
+  @Override
   public String generateBucketPath() {
     return "gcs://" + gcsConfig.getBucketName() + "/" + stagingFolder + "/" + schemaName + "/";
   }
 
-  public String generateFilesList(List<String> files) {
-    StringJoiner joiner = new StringJoiner(",");
-    files.forEach(filename -> joiner.add("'" + filename.substring(filename.lastIndexOf("/") + 1) + "'"));
-    return joiner.toString();
-  }
-
   @Override
   public void copyGcsCsvFileIntoTable(final JdbcDatabase database,
-                                      final String gcsFileLocation,
-                                      final String schema,
-                                      final String tableName,
-                                      final GcsConfig gcsConfig)
+      final String gcsFileLocation,
+      final String schema,
+      final String tableName,
+      final GcsConfig gcsConfig)
       throws SQLException {
     throw new RuntimeException("Snowflake GCS Stream Copier should not copy individual files without use of a parallel copy");
 

airbyte-integrations/connectors/destination-snowflake/src/main/java/io/airbyte/integrations/destination/snowflake/SnowflakeParallelCopyStreamCopier.java

@@ -0,0 +1,17 @@
+package io.airbyte.integrations.destination.snowflake;
+
+import java.util.List;
+import java.util.StringJoiner;
+
+interface SnowflakeParallelCopyStreamCopier {
+
+  default String generateFilesList(List<String> files) {
+    StringJoiner joiner = new StringJoiner(",");
+    files.forEach(filename -> joiner.add("'" + filename.substring(filename.lastIndexOf("/") + 1) + "'"));
+    return joiner.toString();
+  }
+
+  void copyIntoStage(List<String> files);
+
+  String generateBucketPath();
+}

airbyte-integrations/connectors/destination-snowflake/src/main/java/io/airbyte/integrations/destination/snowflake/SnowflakeS3StreamCopier.java

@@ -21,11 +21,12 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.StringJoiner;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class SnowflakeS3StreamCopier extends S3StreamCopier {
+public class SnowflakeS3StreamCopier extends S3StreamCopier implements SnowflakeParallelCopyStreamCopier {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeS3StreamCopier.class);
 
@@ -34,15 +35,16 @@ public class SnowflakeS3StreamCopier extends S3StreamCopier {
   // compression"
   public static final int MAX_PARTS_PER_FILE = 4;
   public static final int MAX_FILES_PER_COPY = 1000;
+  private final ExecutorService executorService = Executors.newFixedThreadPool(5);
 
   public SnowflakeS3StreamCopier(final String stagingFolder,
-                                 final String schema,
-                                 final AmazonS3 client,
-                                 final JdbcDatabase db,
-                                 final S3CopyConfig config,
-                                 final ExtendedNameTransformer nameTransformer,
-                                 final SqlOperations sqlOperations,
-                                 final ConfiguredAirbyteStream configuredAirbyteStream) {
+      final String schema,
+      final AmazonS3 client,
+      final JdbcDatabase db,
+      final S3CopyConfig config,
+      final ExtendedNameTransformer nameTransformer,
+      final SqlOperations sqlOperations,
+      final ConfiguredAirbyteStream configuredAirbyteStream) {
     this(
         stagingFolder,
         schema,
@@ -57,14 +59,14 @@ public SnowflakeS3StreamCopier(final String stagingFolder,
 
   @VisibleForTesting
   SnowflakeS3StreamCopier(final String stagingFolder,
-                          final String schema,
-                          final AmazonS3 client,
-                          final JdbcDatabase db,
-                          final S3CopyConfig config,
-                          final ExtendedNameTransformer nameTransformer,
-                          final SqlOperations sqlOperations,
-                          final Timestamp uploadTime,
-                          final ConfiguredAirbyteStream configuredAirbyteStream) {
+      final String schema,
+      final AmazonS3 client,
+      final JdbcDatabase db,
+      final S3CopyConfig config,
+      final ExtendedNameTransformer nameTransformer,
+      final SqlOperations sqlOperations,
+      final Timestamp uploadTime,
+      final ConfiguredAirbyteStream configuredAirbyteStream) {
 
     super(stagingFolder,
         schema,
@@ -80,16 +82,28 @@ public SnowflakeS3StreamCopier(final String stagingFolder,
 
   @Override
   public void copyStagingFileToTemporaryTable() throws Exception {
-    List<List<String>> partition = Lists.partition(new ArrayList<>(stagingWritersByFile.keySet()), MAX_FILES_PER_COPY);
-    for (int i = 0; i < partition.size(); i++) {
-      LOGGER.info("Starting copy chunk {} to tmp table: {} in destination for stream: {}, schema: {}. Chunks count {}", i, tmpTableName, streamName,
-          schemaName, partition.size());
-      executeCopy(partition.get(i));
-    }
+    List<List<String>> partitions = Lists.partition(new ArrayList<>(stagingWritersByFile.keySet()), MAX_FILES_PER_COPY);
+    LOGGER.info("Starting parallel int copy to tmp table: {} in destination for stream: {}, schema: {}. Chunks count {}", tmpTableName, streamName,
+        schemaName, partitions.size());
+
+    copyFilesInParallel(partitions);
     LOGGER.info("Copy to tmp table {} in destination for stream {} complete.", tmpTableName, streamName);
   }
 
-  private void executeCopy(List<String> files) {
+    private void copyFilesInParallel(List<List<String>> partitions) {
+    partitions.forEach(files -> {
+      try {
+        executorService.execute(() -> copyIntoStage(files));
+      } catch (Exception e) {
+        LOGGER.error("Failed to copy files from stage to tmp table {}", tmpTableName, e);
+        throw new RuntimeException(e);
+      } finally {
+        executorService.shutdown();
+      }
+    });
+  }
+  @Override
+  public void copyIntoStage(List<String> files) {
     final var copyQuery = String.format(
         "COPY INTO %s.%s FROM '%s' "
             + "CREDENTIALS=(aws_key_id='%s' aws_secret_key='%s') "
@@ -104,23 +118,18 @@ private void executeCopy(List<String> files) {
     Exceptions.toRuntime(() -> db.execute(copyQuery));
   }
 
+  @Override
   public String generateBucketPath() {
     return "s3://" + s3Config.getBucketName() + "/"
         + S3OutputPathHelper.getOutputPrefix(s3Config.getBucketPath(), configuredAirbyteStream.getStream()) + "/";
   }
 
-  public String generateFilesList(List<String> files) {
-    StringJoiner joiner = new StringJoiner(",");
-    files.forEach(filename -> joiner.add("'" + filename.substring(filename.lastIndexOf("/") + 1) + "'"));
-    return joiner.toString();
-  }
-
   @Override
   public void copyS3CsvFileIntoTable(final JdbcDatabase database,
-                                     final String s3FileLocation,
-                                     final String schema,
-                                     final String tableName,
-                                     final S3DestinationConfig s3Config)
+      final String s3FileLocation,
+      final String schema,
+      final String tableName,
+      final S3DestinationConfig s3Config)
       throws SQLException {
     throw new RuntimeException("Snowflake Stream Copier should not copy individual files without use of a parallel copy");