support merge mode

mnmandahalf · mnmandahalf · commit e08956dea1a3 · 2023-05-31T11:20:05.000Z
diff --git a/README.md b/README.md
@@ -20,7 +20,9 @@ Snowflake output plugin for Embulk loads records to Snowflake.
 - **retry_limit**: max retry count for database operations (integer, default: 12). When intermediate table to create already created by another process, this plugin will retry with another table name to avoid collision.
 - **retry_wait**: initial retry wait time in milliseconds (integer, default: 1000 (1 second))
 - **max_retry_wait**: upper limit of retry wait, which will be doubled at every retry (integer, default: 1800000 (30 minutes))
-- **mode**: "insert", "insert_direct", "truncate_insert" or "replace". See below. (string, required)
+- **mode**: "insert", "insert_direct", "truncate_insert", "replace" or "merge". See below. (string, required)
+- **merge_keys**: key column names for merging records in merge mode (string array, required in merge mode if table doesn't have primary key)
+- **merge_rule**: list of column assignments for updating existing records used in merge mode, for example `"foo" = T."foo" + S."foo"` (`T` means target table and `S` means source table). (string array, default: always overwrites with new values)
 - **batch_size**: size of a single batch insert (integer, default: 16777216)
 - **default_timezone**: If input column type (embulk type) is timestamp, this plugin needs to format the timestamp into a SQL string. This default_timezone option is used to control the timezone. You can overwrite timezone for each columns using column_options option. (string, default: `UTC`)
 - **column_options**: advanced: a key-value pairs where key is a column name and value is options for the column.
@@ -49,6 +51,10 @@ Snowflake output plugin for Embulk loads records to Snowflake.
   * Behavior: This mode writes rows to an intermediate table first. If all those tasks run correctly, drops the target table and alters the name of the intermediate table into the target table name.
   * Transactional: Yes.
   * Resumable: No.
+* **merge**:
+  * Behavior: This mode writes rows to some intermediate tables first. If all those tasks run correctly, runs MERGE INTO ... WHEN MATCHED THEN UPDATE ...  WHEN NOT MATCHED THEN INSERT ... query. Namely, if merge keys of a record in the intermediate tables already exist in the target table, the target record is updated by the intermediate record, otherwise the intermediate record is inserted. If the target table doesn't exist, it is created automatically.
+  * Transactional: Yes.
+  * Resumable: No.
 
 ## Build
 
diff --git a/src/main/java/org/embulk/output/SnowflakeOutputPlugin.java b/src/main/java/org/embulk/output/SnowflakeOutputPlugin.java
@@ -68,7 +68,12 @@ protected Features getFeatures(PluginTask task) {
         .setMaxTableNameLength(127)
         .setSupportedModes(
             new HashSet<>(
-                Arrays.asList(Mode.INSERT, Mode.INSERT_DIRECT, Mode.TRUNCATE_INSERT, Mode.REPLACE)))
+                Arrays.asList(
+                    Mode.INSERT,
+                    Mode.INSERT_DIRECT,
+                    Mode.TRUNCATE_INSERT,
+                    Mode.REPLACE,
+                    Mode.MERGE)))
         .setIgnoreMergeKeys(false);
   }
 
diff --git a/src/main/java/org/embulk/output/snowflake/SnowflakeOutputConnection.java b/src/main/java/org/embulk/output/snowflake/SnowflakeOutputConnection.java
@@ -4,9 +4,12 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.List;
 import net.snowflake.client.jdbc.SnowflakeConnection;
 import org.embulk.output.jdbc.JdbcColumn;
 import org.embulk.output.jdbc.JdbcOutputConnection;
+import org.embulk.output.jdbc.JdbcSchema;
+import org.embulk.output.jdbc.MergeConfig;
 import org.embulk.output.jdbc.TableIdentifier;
 
 public class SnowflakeOutputConnection extends JdbcOutputConnection {
@@ -75,6 +78,71 @@ protected String buildColumnTypeName(JdbcColumn c) {
     }
   }
 
+  @Override
+  protected String buildCollectMergeSql(
+      List<TableIdentifier> fromTables,
+      JdbcSchema schema,
+      TableIdentifier toTable,
+      MergeConfig mergeConfig)
+      throws SQLException {
+    StringBuilder sb = new StringBuilder();
+
+    sb.append("MERGE INTO ");
+    sb.append(quoteTableIdentifier(toTable));
+    sb.append(" AS T ");
+    sb.append(" USING (");
+    for (int i = 0; i < fromTables.size(); i++) {
+      if (i != 0) {
+        sb.append(" UNION ALL ");
+      }
+      sb.append(" SELECT ");
+      sb.append(buildColumns(schema, ""));
+      sb.append(" FROM ");
+      sb.append(quoteTableIdentifier(fromTables.get(i)));
+    }
+    sb.append(") AS S");
+    sb.append(" ON (");
+    for (int i = 0; i < mergeConfig.getMergeKeys().size(); i++) {
+      if (i != 0) {
+        sb.append(" AND ");
+      }
+      String mergeKey = quoteIdentifierString(mergeConfig.getMergeKeys().get(i));
+      sb.append("T.");
+      sb.append(mergeKey);
+      sb.append(" = S.");
+      sb.append(mergeKey);
+    }
+    sb.append(")");
+    sb.append(" WHEN MATCHED THEN");
+    sb.append(" UPDATE SET ");
+    if (mergeConfig.getMergeRule().isPresent()) {
+      for (int i = 0; i < mergeConfig.getMergeRule().get().size(); i++) {
+        if (i != 0) {
+          sb.append(", ");
+        }
+        sb.append(mergeConfig.getMergeRule().get().get(i));
+      }
+    } else {
+      for (int i = 0; i < schema.getCount(); i++) {
+        if (i != 0) {
+          sb.append(", ");
+        }
+        String column = quoteIdentifierString(schema.getColumnName(i));
+        sb.append(column);
+        sb.append(" = S.");
+        sb.append(column);
+      }
+    }
+    sb.append(" WHEN NOT MATCHED THEN");
+    sb.append(" INSERT (");
+    sb.append(buildColumns(schema, ""));
+    sb.append(") VALUES (");
+    sb.append(buildColumns(schema, "S."));
+    sb.append(");");
+
+    return sb.toString();
+  }
+
   protected String buildCreateStageSQL(StageIdentifier stageIdentifier) {
     StringBuilder sb = new StringBuilder();
     sb.append("CREATE STAGE IF NOT EXISTS ");
@@ -137,4 +205,16 @@ protected String quoteInternalStoragePath(
     sb.append(".csv.gz");
     return sb.toString();
   }
+
+  private String buildColumns(JdbcSchema schema, String prefix) {
+    StringBuilder sb = new StringBuilder();
+    for (int i = 0; i < schema.getCount(); i++) {
+      if (i != 0) {
+        sb.append(", ");
+      }
+      sb.append(prefix);
+      sb.append(quoteIdentifierString(schema.getColumnName(i)));
+    }
+    return sb.toString();
+  }
 }
diff --git a/src/test/java/org/embulk/output/snowflake/TestSnowflakeOutputPlugin.java b/src/test/java/org/embulk/output/snowflake/TestSnowflakeOutputPlugin.java
@@ -13,6 +13,7 @@
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
 import java.util.Properties;
@@ -434,6 +435,7 @@ public void testExecuteMultiQuery() throws IOException {
     String targetTableFullName =
         String.format(
             "\"%s\".\"%s\".\"%s\"", TEST_SNOWFLAKE_DB, TEST_SNOWFLAKE_SCHEMA, targetTableName);
+    System.out.println(targetTableFullName);
     runQuery(
         String.format("create table %s (c0 FLOAT, c1 STRING)", targetTableFullName),
         foreachResult(rs_ -> {}));
@@ -519,6 +521,171 @@ public void testRuntimeInsertStringTable() throws IOException {
             }));
   }
 
+  @Test
+  public void testRuntimeMergeTable() throws IOException {
+    File in = testFolder.newFile(SnowflakeUtils.randomString(8) + ".csv");
+    List<String> lines =
+        Stream.of("c0:double,c1:double,c2:string", "1,1,aaa", "1,2,bbb", "1,3,ccc")
+            .collect(Collectors.toList());
+    Files.write(in.toPath(), lines);
+
+    String targetTableName = generateTemporaryTableName();
+    String targetTableFullName =
+        String.format(
+            "\"%s\".\"%s\".\"%s\"", TEST_SNOWFLAKE_DB, TEST_SNOWFLAKE_SCHEMA, targetTableName);
+    runQuery(
+        String.format(
+            "create table %s (\"c0\" NUMBER, \"c1\" NUMBER, \"c2\" STRING)", targetTableFullName),
+        foreachResult(rs_ -> {}));
+    runQuery(
+        String.format("insert into %s values (1, 1, 'ddd'), (1, 4, 'eee');", targetTableFullName),
+        foreachResult(rs_ -> {}));
+
+    final ConfigSource config =
+        CONFIG_MAPPER_FACTORY
+            .newConfigSource()
+            .set("type", "snowflake")
+            .set("user", TEST_SNOWFLAKE_USER)
+            .set("password", TEST_SNOWFLAKE_PASSWORD)
+            .set("host", TEST_SNOWFLAKE_HOST)
+            .set("database", TEST_SNOWFLAKE_DB)
+            .set("warehouse", TEST_SNOWFLAKE_WAREHOUSE)
+            .set("schema", TEST_SNOWFLAKE_SCHEMA)
+            .set("mode", "merge")
+            .set("merge_keys", new ArrayList<String>(Arrays.asList("c0", "c1")))
+            .set("table", targetTableName);
+    embulk.runOutput(config, in.toPath());
+
+    runQuery(
+        String.format("select count(*) from %s;", targetTableFullName),
+        foreachResult(
+            rs -> {
+              assertEquals(4, rs.getInt(1));
+            }));
+    List<String> results = new ArrayList();
+    runQuery(
+        "select \"c0\",\"c1\",\"c2\" from " + targetTableFullName + " order by 1, 2",
+        foreachResult(
+            rs -> {
+              results.add(rs.getString(1) + "," + rs.getString(2) + "," + rs.getString(3));
+            }));
+    List<String> expected =
+        Stream.of("1,1,aaa", "1,2,bbb", "1,3,ccc", "1,4,eee").collect(Collectors.toList());
+    for (int i = 0; i < results.size(); i++) {
+      assertEquals(expected.get(i), results.get(i));
+    }
+  }
+
+  @Test
+  public void testRuntimeMergeTableWithMergeRule() throws IOException {
+    File in = testFolder.newFile(SnowflakeUtils.randomString(8) + ".csv");
+    List<String> lines =
+        Stream.of("c0:double,c1:string", "0.0,aaa", "0.1,bbb", "1.2,ccc")
+            .collect(Collectors.toList());
+    Files.write(in.toPath(), lines);
+
+    String targetTableName = generateTemporaryTableName();
+    String targetTableFullName =
+        String.format(
+            "\"%s\".\"%s\".\"%s\"", TEST_SNOWFLAKE_DB, TEST_SNOWFLAKE_SCHEMA, targetTableName);
+    runQuery(
+        String.format("create table %s (\"c0\" FLOAT, \"c1\" STRING)", targetTableFullName),
+        foreachResult(rs_ -> {}));
+    runQuery(
+        String.format("insert into %s values (0.0, 'ddd'), (1.3, 'eee');", targetTableFullName),
+        foreachResult(rs_ -> {}));
+
+    final ConfigSource config =
+        CONFIG_MAPPER_FACTORY
+            .newConfigSource()
+            .set("type", "snowflake")
+            .set("user", TEST_SNOWFLAKE_USER)
+            .set("password", TEST_SNOWFLAKE_PASSWORD)
+            .set("host", TEST_SNOWFLAKE_HOST)
+            .set("database", TEST_SNOWFLAKE_DB)
+            .set("warehouse", TEST_SNOWFLAKE_WAREHOUSE)
+            .set("schema", TEST_SNOWFLAKE_SCHEMA)
+            .set("mode", "merge")
+            .set("merge_keys", new ArrayList<String>(Arrays.asList("c0")))
+            .set(
+                "merge_rule", new ArrayList<String>(Arrays.asList("\"c1\" = T.\"c1\" || S.\"c1\"")))
+            .set("table", targetTableName);
+    embulk.runOutput(config, in.toPath());
+
+    runQuery(
+        String.format("select count(*) from %s;", targetTableFullName),
+        foreachResult(
+            rs -> {
+              assertEquals(4, rs.getInt(1));
+            }));
+    List<String> results = new ArrayList();
+    runQuery(
+        "select \"c0\",\"c1\" from " + targetTableFullName + " order by 1",
+        foreachResult(
+            rs -> {
+              results.add(rs.getString(1) + "," + rs.getString(2));
+            }));
+    List<String> expected =
+        Stream.of("0.0,dddaaa", "0.1,bbb", "1.2,ccc", "1.3,eee").collect(Collectors.toList());
+    for (int i = 0; i < results.size(); i++) {
+      assertEquals(expected.get(i), results.get(i));
+    }
+  }
+
+  @Test
+  public void testRuntimeMergeTableWithoutMergeKey() throws IOException {
+    File in = testFolder.newFile(SnowflakeUtils.randomString(8) + ".csv");
+    List<String> lines =
+        Stream.of("c0:double,c1:string", "0.0,aaa", "0.1,bbb", "1.2,ccc")
+            .collect(Collectors.toList());
+    Files.write(in.toPath(), lines);
+
+    String targetTableName = generateTemporaryTableName();
+    String targetTableFullName =
+        String.format(
+            "\"%s\".\"%s\".\"%s\"", TEST_SNOWFLAKE_DB, TEST_SNOWFLAKE_SCHEMA, targetTableName);
+    runQuery(
+        String.format(
+            "create table %s (\"c0\" FLOAT PRIMARY KEY, \"c1\" STRING)", targetTableFullName),
+        foreachResult(rs_ -> {}));
+    runQuery(
+        String.format("insert into %s values (0.0, 'ddd'), (1.3, 'eee');", targetTableFullName),
+        foreachResult(rs_ -> {}));
+
+    final ConfigSource config =
+        CONFIG_MAPPER_FACTORY
+            .newConfigSource()
+            .set("type", "snowflake")
+            .set("user", TEST_SNOWFLAKE_USER)
+            .set("password", TEST_SNOWFLAKE_PASSWORD)
+            .set("host", TEST_SNOWFLAKE_HOST)
+            .set("database", TEST_SNOWFLAKE_DB)
+            .set("warehouse", TEST_SNOWFLAKE_WAREHOUSE)
+            .set("schema", TEST_SNOWFLAKE_SCHEMA)
+            .set("mode", "merge")
+            .set("table", targetTableName);
+    embulk.runOutput(config, in.toPath());
+
+    runQuery(
+        String.format("select count(*) from %s;", targetTableFullName),
+        foreachResult(
+            rs -> {
+              assertEquals(4, rs.getInt(1));
+            }));
+    List<String> results = new ArrayList();
+    runQuery(
+        "select \"c0\",\"c1\" from " + targetTableFullName + " order by 1",
+        foreachResult(
+            rs -> {
+              results.add(rs.getString(1) + "," + rs.getString(2));
+            }));
+    List<String> expected =
+        Stream.of("0.0,aaa", "0.1,bbb", "1.2,ccc", "1.3,eee").collect(Collectors.toList());
+    for (int i = 0; i < results.size(); i++) {
+      assertEquals(expected.get(i), results.get(i));
+    }
+  }
+
   @Ignore(
       "This test takes so long time because it needs to create more than 1000 tables, so ignored...")
   @Test(expected = Test.None.class /* no exception expected */)
