Merge pull request #54 from trocco-io/retry-upload-files

Add file upload retry handling

Author: u110

build.gradle

@@ -51,7 +51,7 @@ dependencies {
     testImplementation "org.embulk:embulk-parser-csv:0.10.31"
 
     compile "org.embulk:embulk-output-jdbc:0.10.2"
-    compile "net.snowflake:snowflake-jdbc:3.13.14"
+    compile "net.snowflake:snowflake-jdbc:3.13.26"
 }
 embulkPlugin {
     mainClass = "org.embulk.output.SnowflakeOutputPlugin"

gradle/dependency-locks/embulkPluginRuntime.lockfile

@@ -6,7 +6,7 @@ com.fasterxml.jackson.core:jackson-core:2.6.7
 com.fasterxml.jackson.core:jackson-databind:2.6.7
 com.fasterxml.jackson.datatype:jackson-datatype-jdk8:2.6.7
 javax.validation:validation-api:1.1.0.Final
-net.snowflake:snowflake-jdbc:3.13.14
+net.snowflake:snowflake-jdbc:3.13.26
 org.embulk:embulk-output-jdbc:0.10.2
 org.embulk:embulk-util-config:0.3.0
 org.embulk:embulk-util-json:0.1.1

makefile

@@ -0,0 +1,9 @@
+build: gradlew/build
+gem: gradlew/gem
+check: gradlew/check
+lint-auto: gradlew/spotlessApply
+test: gradlew/test
+update-dependencies:
+	./gradlew dependencies --write-locks
+gradlew/%:
+	./gradlew $(@F)

src/main/java/org/embulk/output/SnowflakeOutputPlugin.java

@@ -51,6 +51,10 @@ public interface SnowflakePluginTask extends PluginTask {
     @Config("delete_stage")
     @ConfigDefault("false")
     public boolean getDeleteStage();
+
+    @Config("max_upload_retries")
+    @ConfigDefault("3")
+    public int getMaxUploadRetries();
   }
 
   @Override
@@ -145,7 +149,11 @@ protected BatchInsert newBatchInsert(PluginTask task, Optional<MergeConfig> merg
       snowflakeCon.runCreateStage(this.stageIdentifier);
     }
 
-    return new SnowflakeCopyBatchInsert(getConnector(task, true), this.stageIdentifier, false);
+    return new SnowflakeCopyBatchInsert(
+        getConnector(task, true),
+        this.stageIdentifier,
+        false,
+        ((SnowflakePluginTask) task).getMaxUploadRetries());
   }
 
   @Override

src/main/java/org/embulk/output/snowflake/SnowflakeCopyBatchInsert.java

@@ -26,6 +26,7 @@ public class SnowflakeCopyBatchInsert implements BatchInsert {
   protected static final String nullString = "\\N";
   protected static final String newLineString = "\n";
   protected static final String delimiterString = "\t";
+  private final int maxUploadRetries;
 
   private SnowflakeOutputConnection connection = null;
   private TableIdentifier tableIdentifier = null;
@@ -39,7 +40,10 @@ public class SnowflakeCopyBatchInsert implements BatchInsert {
   private List<Future<Void>> uploadAndCopyFutures;
 
   public SnowflakeCopyBatchInsert(
-      JdbcOutputConnector connector, StageIdentifier stageIdentifier, boolean deleteStageFile)
+      JdbcOutputConnector connector,
+      StageIdentifier stageIdentifier,
+      boolean deleteStageFile,
+      int maxUploadRetries)
       throws IOException {
     this.index = 0;
     openNewFile();
@@ -48,6 +52,7 @@ public SnowflakeCopyBatchInsert(
     this.executorService = Executors.newCachedThreadPool();
     this.deleteStageFile = deleteStageFile;
     this.uploadAndCopyFutures = new ArrayList();
+    this.maxUploadRetries = maxUploadRetries;
   }
 
   @Override
@@ -251,7 +256,7 @@ public void flush() throws IOException, SQLException {
     String snowflakeStageFileName = "embulk_snowflake_" + SnowflakeUtils.randomString(8);
 
     UploadTask uploadTask =
-        new UploadTask(file, batchRows, stageIdentifier, snowflakeStageFileName);
+        new UploadTask(file, batchRows, stageIdentifier, snowflakeStageFileName, maxUploadRetries);
     Future<Void> uploadFuture = executorService.submit(uploadTask);
     uploadAndCopyFutures.add(uploadFuture);
 
@@ -330,28 +335,49 @@ private class UploadTask implements Callable<Void> {
     private final int batchRows;
     private final String snowflakeStageFileName;
     private final StageIdentifier stageIdentifier;
+    private final int maxUploadRetries;
 
     public UploadTask(
-        File file, int batchRows, StageIdentifier stageIdentifier, String snowflakeStageFileName) {
+        File file,
+        int batchRows,
+        StageIdentifier stageIdentifier,
+        String snowflakeStageFileName,
+        int maxUploadRetries) {
       this.file = file;
       this.batchRows = batchRows;
       this.snowflakeStageFileName = snowflakeStageFileName;
       this.stageIdentifier = stageIdentifier;
+      this.maxUploadRetries = maxUploadRetries;
     }
 
-    public Void call() throws IOException, SQLException {
+    public Void call() throws IOException, SQLException, InterruptedException {
       logger.info(
           String.format(
               "Uploading file id %s to Snowflake (%,d bytes %,d rows)",
               snowflakeStageFileName, file.length(), batchRows));
 
+      int retries = 0;
       try {
         long startTime = System.currentTimeMillis();
         // put file to snowflake internal storage
         SnowflakeOutputConnection con = (SnowflakeOutputConnection) connector.connect(true);
 
         FileInputStream fileInputStream = new FileInputStream(file);
-        con.runUploadFile(stageIdentifier, snowflakeStageFileName, fileInputStream);
+        do {
+          try {
+            con.runUploadFile(stageIdentifier, snowflakeStageFileName, fileInputStream);
+          } catch (SQLException e) {
+            retries++;
+            if (retries > this.maxUploadRetries) {
+              throw e;
+            }
+            logger.warn(
+                String.format(
+                    "Upload error %s file %s retries: %d", e, snowflakeStageFileName, retries));
+            Thread.sleep(retries * retries * 1000);
+          }
+          break;
+        } while (retries < this.maxUploadRetries);
 
         double seconds = (System.currentTimeMillis() - startTime) / 1000.0;