Skip to content

Commit b1e562c

Browse files
mnmandahalfmnmandahalf
authored
Merge pull request #59 from trocco-io/feature/add-merge-mode
Support merge mode
2 parents e75d89d + 16b428d commit b1e562c

File tree

4 files changed

+259
-2
lines changed

4 files changed

+259
-2
lines changed

‎README.md‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@ Snowflake output plugin for Embulk loads records to Snowflake.
2020
- **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.
2121
- **retry_wait**: initial retry wait time in milliseconds (integer, default: 1000 (1 second))
2222
- **max_retry_wait**: upper limit of retry wait, which will be doubled at every retry (integer, default: 1800000 (30 minutes))
23-
- **mode**: "insert", "insert_direct", "truncate_insert" or "replace". See below. (string, required)
23+
- **mode**: "insert", "insert_direct", "truncate_insert", "replace" or "merge". See below. (string, required)
24+
- **merge_keys**: key column names for merging records in merge mode (string array, required in merge mode if table doesn't have primary key)
25+
- **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)
2426
- **batch_size**: size of a single batch insert (integer, default: 16777216)
2527
- **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`)
2628
- **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.
4951
* 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.
5052
* Transactional: Yes.
5153
* Resumable: No.
54+
* **merge**:
55+
* 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.
56+
* Transactional: Yes.
57+
* Resumable: No.
5258

5359
## Build
5460

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,12 @@ protected Features getFeatures(PluginTask task) {
6868
.setMaxTableNameLength(127)
6969
.setSupportedModes(
7070
new HashSet<>(
71-
Arrays.asList(Mode.INSERT, Mode.INSERT_DIRECT, Mode.TRUNCATE_INSERT, Mode.REPLACE)))
71+
Arrays.asList(
72+
Mode.INSERT,
73+
Mode.INSERT_DIRECT,
74+
Mode.TRUNCATE_INSERT,
75+
Mode.REPLACE,
76+
Mode.MERGE)))
7277
.setIgnoreMergeKeys(false);
7378
}
7479

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,12 @@
44
import java.sql.Connection;
55
import java.sql.SQLException;
66
import java.sql.Statement;
7+
import java.util.List;
78
import net.snowflake.client.jdbc.SnowflakeConnection;
89
import org.embulk.output.jdbc.JdbcColumn;
910
import org.embulk.output.jdbc.JdbcOutputConnection;
11+
import org.embulk.output.jdbc.JdbcSchema;
12+
import org.embulk.output.jdbc.MergeConfig;
1013
import org.embulk.output.jdbc.TableIdentifier;
1114

1215
public class SnowflakeOutputConnection extends JdbcOutputConnection {
@@ -75,6 +78,71 @@ protected String buildColumnTypeName(JdbcColumn c) {
7578
}
7679
}
7780

81+
@Override
82+
protected String buildCollectMergeSql(
83+
List<TableIdentifier> fromTables,
84+
JdbcSchema schema,
85+
TableIdentifier toTable,
86+
MergeConfig mergeConfig)
87+
throws SQLException {
88+
StringBuilder sb = new StringBuilder();
89+
90+
sb.append("MERGE INTO ");
91+
sb.append(quoteTableIdentifier(toTable));
92+
sb.append(" AS T ");
93+
sb.append(" USING (");
94+
for (int i = 0; i < fromTables.size(); i++) {
95+
if (i != 0) {
96+
sb.append(" UNION ALL ");
97+
}
98+
sb.append(" SELECT ");
99+
sb.append(buildColumns(schema, ""));
100+
sb.append(" FROM ");
101+
sb.append(quoteTableIdentifier(fromTables.get(i)));
102+
}
103+
sb.append(") AS S");
104+
sb.append(" ON (");
105+
for (int i = 0; i < mergeConfig.getMergeKeys().size(); i++) {
106+
if (i != 0) {
107+
sb.append(" AND ");
108+
}
109+
String mergeKey = quoteIdentifierString(mergeConfig.getMergeKeys().get(i));
110+
sb.append("T.");
111+
sb.append(mergeKey);
112+
sb.append(" = S.");
113+
sb.append(mergeKey);
114+
}
115+
sb.append(")");
116+
sb.append(" WHEN MATCHED THEN");
117+
sb.append(" UPDATE SET ");
118+
if (mergeConfig.getMergeRule().isPresent()) {
119+
for (int i = 0; i < mergeConfig.getMergeRule().get().size(); i++) {
120+
if (i != 0) {
121+
sb.append(", ");
122+
}
123+
sb.append(mergeConfig.getMergeRule().get().get(i));
124+
}
125+
} else {
126+
for (int i = 0; i < schema.getCount(); i++) {
127+
if (i != 0) {
128+
sb.append(", ");
129+
}
130+
String column = quoteIdentifierString(schema.getColumnName(i));
131+
sb.append(column);
132+
sb.append(" = S.");
133+
sb.append(column);
134+
}
135+
}
136+
sb.append(" WHEN NOT MATCHED THEN");
137+
sb.append(" INSERT (");
138+
sb.append(buildColumns(schema, ""));
139+
sb.append(") VALUES (");
140+
sb.append(buildColumns(schema, "S."));
141+
sb.append(");");
142+
143+
return sb.toString();
144+
}
145+
78146
protected String buildCreateStageSQL(StageIdentifier stageIdentifier) {
79147
StringBuilder sb = new StringBuilder();
80148
sb.append("CREATE STAGE IF NOT EXISTS ");
@@ -137,4 +205,16 @@ protected String quoteInternalStoragePath(
137205
sb.append(".csv.gz");
138206
return sb.toString();
139207
}
208+
209+
private String buildColumns(JdbcSchema schema, String prefix) {
210+
StringBuilder sb = new StringBuilder();
211+
for (int i = 0; i < schema.getCount(); i++) {
212+
if (i != 0) {
213+
sb.append(", ");
214+
}
215+
sb.append(prefix);
216+
sb.append(quoteIdentifierString(schema.getColumnName(i)));
217+
}
218+
return sb.toString();
219+
}
140220
}

‎src/test/java/org/embulk/output/snowflake/TestSnowflakeOutputPlugin.java‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import java.sql.SQLException;
1414
import java.sql.Statement;
1515
import java.util.ArrayList;
16+
import java.util.Arrays;
1617
import java.util.List;
1718
import java.util.Optional;
1819
import java.util.Properties;
@@ -519,6 +520,171 @@ public void testRuntimeInsertStringTable() throws IOException {
519520
}));
520521
}
521522

523+
@Test
524+
public void testRuntimeMergeTable() throws IOException {
525+
File in = testFolder.newFile(SnowflakeUtils.randomString(8) + ".csv");
526+
List<String> lines =
527+
Stream.of("c0:double,c1:double,c2:string", "1,1,aaa", "1,2,bbb", "1,3,ccc")
528+
.collect(Collectors.toList());
529+
Files.write(in.toPath(), lines);
530+
531+
String targetTableName = generateTemporaryTableName();
532+
String targetTableFullName =
533+
String.format(
534+
"\"%s\".\"%s\".\"%s\"", TEST_SNOWFLAKE_DB, TEST_SNOWFLAKE_SCHEMA, targetTableName);
535+
runQuery(
536+
String.format(
537+
"create table %s (\"c0\" NUMBER, \"c1\" NUMBER, \"c2\" STRING)", targetTableFullName),
538+
foreachResult(rs_ -> {}));
539+
runQuery(
540+
String.format("insert into %s values (1, 1, 'ddd'), (1, 4, 'eee');", targetTableFullName),
541+
foreachResult(rs_ -> {}));
542+
543+
final ConfigSource config =
544+
CONFIG_MAPPER_FACTORY
545+
.newConfigSource()
546+
.set("type", "snowflake")
547+
.set("user", TEST_SNOWFLAKE_USER)
548+
.set("password", TEST_SNOWFLAKE_PASSWORD)
549+
.set("host", TEST_SNOWFLAKE_HOST)
550+
.set("database", TEST_SNOWFLAKE_DB)
551+
.set("warehouse", TEST_SNOWFLAKE_WAREHOUSE)
552+
.set("schema", TEST_SNOWFLAKE_SCHEMA)
553+
.set("mode", "merge")
554+
.set("merge_keys", new ArrayList<String>(Arrays.asList("c0", "c1")))
555+
.set("table", targetTableName);
556+
embulk.runOutput(config, in.toPath());
557+
558+
runQuery(
559+
String.format("select count(*) from %s;", targetTableFullName),
560+
foreachResult(
561+
rs -> {
562+
assertEquals(4, rs.getInt(1));
563+
}));
564+
List<String> results = new ArrayList();
565+
runQuery(
566+
"select \"c0\",\"c1\",\"c2\" from " + targetTableFullName + " order by 1, 2",
567+
foreachResult(
568+
rs -> {
569+
results.add(rs.getString(1) + "," + rs.getString(2) + "," + rs.getString(3));
570+
}));
571+
List<String> expected =
572+
Stream.of("1,1,aaa", "1,2,bbb", "1,3,ccc", "1,4,eee").collect(Collectors.toList());
573+
for (int i = 0; i < results.size(); i++) {
574+
assertEquals(expected.get(i), results.get(i));
575+
}
576+
}
577+
578+
@Test
579+
public void testRuntimeMergeTableWithMergeRule() throws IOException {
580+
File in = testFolder.newFile(SnowflakeUtils.randomString(8) + ".csv");
581+
List<String> lines =
582+
Stream.of("c0:double,c1:string", "0.0,aaa", "0.1,bbb", "1.2,ccc")
583+
.collect(Collectors.toList());
584+
Files.write(in.toPath(), lines);
585+
586+
String targetTableName = generateTemporaryTableName();
587+
String targetTableFullName =
588+
String.format(
589+
"\"%s\".\"%s\".\"%s\"", TEST_SNOWFLAKE_DB, TEST_SNOWFLAKE_SCHEMA, targetTableName);
590+
runQuery(
591+
String.format("create table %s (\"c0\" FLOAT, \"c1\" STRING)", targetTableFullName),
592+
foreachResult(rs_ -> {}));
593+
runQuery(
594+
String.format("insert into %s values (0.0, 'ddd'), (1.3, 'eee');", targetTableFullName),
595+
foreachResult(rs_ -> {}));
596+
597+
final ConfigSource config =
598+
CONFIG_MAPPER_FACTORY
599+
.newConfigSource()
600+
.set("type", "snowflake")
601+
.set("user", TEST_SNOWFLAKE_USER)
602+
.set("password", TEST_SNOWFLAKE_PASSWORD)
603+
.set("host", TEST_SNOWFLAKE_HOST)
604+
.set("database", TEST_SNOWFLAKE_DB)
605+
.set("warehouse", TEST_SNOWFLAKE_WAREHOUSE)
606+
.set("schema", TEST_SNOWFLAKE_SCHEMA)
607+
.set("mode", "merge")
608+
.set("merge_keys", new ArrayList<String>(Arrays.asList("c0")))
609+
.set(
610+
"merge_rule", new ArrayList<String>(Arrays.asList("\"c1\" = T.\"c1\" || S.\"c1\"")))
611+
.set("table", targetTableName);
612+
embulk.runOutput(config, in.toPath());
613+
614+
runQuery(
615+
String.format("select count(*) from %s;", targetTableFullName),
616+
foreachResult(
617+
rs -> {
618+
assertEquals(4, rs.getInt(1));
619+
}));
620+
List<String> results = new ArrayList();
621+
runQuery(
622+
"select \"c0\",\"c1\" from " + targetTableFullName + " order by 1",
623+
foreachResult(
624+
rs -> {
625+
results.add(rs.getString(1) + "," + rs.getString(2));
626+
}));
627+
List<String> expected =
628+
Stream.of("0.0,dddaaa", "0.1,bbb", "1.2,ccc", "1.3,eee").collect(Collectors.toList());
629+
for (int i = 0; i < results.size(); i++) {
630+
assertEquals(expected.get(i), results.get(i));
631+
}
632+
}
633+
634+
@Test
635+
public void testRuntimeMergeTableWithoutMergeKey() throws IOException {
636+
File in = testFolder.newFile(SnowflakeUtils.randomString(8) + ".csv");
637+
List<String> lines =
638+
Stream.of("c0:double,c1:string", "0.0,aaa", "0.1,bbb", "1.2,ccc")
639+
.collect(Collectors.toList());
640+
Files.write(in.toPath(), lines);
641+
642+
String targetTableName = generateTemporaryTableName();
643+
String targetTableFullName =
644+
String.format(
645+
"\"%s\".\"%s\".\"%s\"", TEST_SNOWFLAKE_DB, TEST_SNOWFLAKE_SCHEMA, targetTableName);
646+
runQuery(
647+
String.format(
648+
"create table %s (\"c0\" FLOAT PRIMARY KEY, \"c1\" STRING)", targetTableFullName),
649+
foreachResult(rs_ -> {}));
650+
runQuery(
651+
String.format("insert into %s values (0.0, 'ddd'), (1.3, 'eee');", targetTableFullName),
652+
foreachResult(rs_ -> {}));
653+
654+
final ConfigSource config =
655+
CONFIG_MAPPER_FACTORY
656+
.newConfigSource()
657+
.set("type", "snowflake")
658+
.set("user", TEST_SNOWFLAKE_USER)
659+
.set("password", TEST_SNOWFLAKE_PASSWORD)
660+
.set("host", TEST_SNOWFLAKE_HOST)
661+
.set("database", TEST_SNOWFLAKE_DB)
662+
.set("warehouse", TEST_SNOWFLAKE_WAREHOUSE)
663+
.set("schema", TEST_SNOWFLAKE_SCHEMA)
664+
.set("mode", "merge")
665+
.set("table", targetTableName);
666+
embulk.runOutput(config, in.toPath());
667+
668+
runQuery(
669+
String.format("select count(*) from %s;", targetTableFullName),
670+
foreachResult(
671+
rs -> {
672+
assertEquals(4, rs.getInt(1));
673+
}));
674+
List<String> results = new ArrayList();
675+
runQuery(
676+
"select \"c0\",\"c1\" from " + targetTableFullName + " order by 1",
677+
foreachResult(
678+
rs -> {
679+
results.add(rs.getString(1) + "," + rs.getString(2));
680+
}));
681+
List<String> expected =
682+
Stream.of("0.0,aaa", "0.1,bbb", "1.2,ccc", "1.3,eee").collect(Collectors.toList());
683+
for (int i = 0; i < results.size(); i++) {
684+
assertEquals(expected.get(i), results.get(i));
685+
}
686+
}
687+
522688
@Ignore(
523689
"This test takes so long time because it needs to create more than 1000 tables, so ignored...")
524690
@Test(expected = Test.None.class /* no exception expected */)

0 commit comments

Comments
 (0)