Merge pull request #61 from trocco-io/feat/EMPTY_FIELD_AS_NULL

empty_field_as_null

Author: giwa

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

@@ -55,6 +55,10 @@ public interface SnowflakePluginTask extends PluginTask {
     @Config("max_upload_retries")
     @ConfigDefault("3")
     public int getMaxUploadRetries();
+
+    @Config("empty_field_as_null")
+    @ConfigDefault("true")
+    public boolean getEmtpyFieldAsNull();
   }
 
   @Override
@@ -153,12 +157,14 @@ protected BatchInsert newBatchInsert(PluginTask task, Optional<MergeConfig> merg
       this.stageIdentifier = StageIdentifierHolder.getStageIdentifier(t);
       snowflakeCon.runCreateStage(this.stageIdentifier);
     }
+    SnowflakePluginTask pluginTask = (SnowflakePluginTask) task;
 
     return new SnowflakeCopyBatchInsert(
         getConnector(task, true),
         this.stageIdentifier,
         false,
-        ((SnowflakePluginTask) task).getMaxUploadRetries());
+        pluginTask.getMaxUploadRetries(),
+        pluginTask.getEmtpyFieldAsNull());
   }
 
   @Override

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

@@ -38,12 +38,14 @@ public class SnowflakeCopyBatchInsert implements BatchInsert {
   private long totalRows;
   private int fileCount;
   private List<Future<Void>> uploadAndCopyFutures;
+  private boolean emptyFieldAsNull;
 
   public SnowflakeCopyBatchInsert(
       JdbcOutputConnector connector,
       StageIdentifier stageIdentifier,
       boolean deleteStageFile,
-      int maxUploadRetries)
+      int maxUploadRetries,
+      boolean emptyFieldAsNull)
       throws IOException {
     this.index = 0;
     openNewFile();
@@ -53,6 +55,7 @@ public SnowflakeCopyBatchInsert(
     this.deleteStageFile = deleteStageFile;
     this.uploadAndCopyFutures = new ArrayList();
     this.maxUploadRetries = maxUploadRetries;
+    this.emptyFieldAsNull = emptyFieldAsNull;
   }
 
   @Override
@@ -260,7 +263,7 @@ public void flush() throws IOException, SQLException {
     Future<Void> uploadFuture = executorService.submit(uploadTask);
     uploadAndCopyFutures.add(uploadFuture);
 
-    CopyTask copyTask = new CopyTask(uploadFuture, snowflakeStageFileName);
+    CopyTask copyTask = new CopyTask(uploadFuture, snowflakeStageFileName, emptyFieldAsNull);
     uploadAndCopyFutures.add(executorService.submit(copyTask));
 
     fileCount++;
@@ -393,10 +396,13 @@ public Void call() throws IOException, SQLException, InterruptedException {
   private class CopyTask implements Callable<Void> {
     private final Future<Void> uploadFuture;
     private final String snowflakeStageFileName;
+    private final boolean emptyFieldAsNull;
 
-    public CopyTask(Future<Void> uploadFuture, String snowflakeStageFileName) {
+    public CopyTask(
+        Future<Void> uploadFuture, String snowflakeStageFileName, boolean emptyFieldAsNull) {
       this.uploadFuture = uploadFuture;
       this.snowflakeStageFileName = snowflakeStageFileName;
+      this.emptyFieldAsNull = emptyFieldAsNull;
     }
 
     public Void call() throws SQLException, InterruptedException, ExecutionException {
@@ -408,7 +414,12 @@ public Void call() throws SQLException, InterruptedException, ExecutionException
           logger.info("Running COPY from file {}", snowflakeStageFileName);
 
           long startTime = System.currentTimeMillis();
-          con.runCopy(tableIdentifier, stageIdentifier, snowflakeStageFileName, delimiterString);
+          con.runCopy(
+              tableIdentifier,
+              stageIdentifier,
+              snowflakeStageFileName,
+              delimiterString,
+              emptyFieldAsNull);
 
           double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
 

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

@@ -21,9 +21,11 @@ public void runCopy(
       TableIdentifier tableIdentifier,
       StageIdentifier stageIdentifier,
       String filename,
-      String delimiterString)
+      String delimiterString,
+      boolean emptyFieldAsNull)
       throws SQLException {
-    String sql = buildCopySQL(tableIdentifier, stageIdentifier, filename, delimiterString);
+    String sql =
+        buildCopySQL(tableIdentifier, stageIdentifier, filename, delimiterString, emptyFieldAsNull);
     runUpdate(sql);
   }
 
@@ -171,15 +173,20 @@ protected String buildCopySQL(
       TableIdentifier tableIdentifier,
       StageIdentifier stageIdentifier,
       String snowflakeStageFileName,
-      String delimiterString) {
+      String delimiterString,
+      boolean emptyFieldAsNull) {
     StringBuilder sb = new StringBuilder();
     sb.append("COPY INTO ");
     quoteTableIdentifier(sb, tableIdentifier);
     sb.append(" FROM ");
     quoteInternalStoragePath(sb, stageIdentifier, snowflakeStageFileName);
     sb.append(" FILE_FORMAT = ( TYPE = CSV FIELD_DELIMITER = '");
     sb.append(delimiterString);
-    sb.append("');");
+    sb.append("') ");
+    if (!emptyFieldAsNull) {
+      sb.append("EMPTY_FIELD_AS_NULL = FALSE");
+    }
+    sb.append(";");
     return sb.toString();
   }