feat: add support for snowflake merge statements (#1887)

* feat: support snowflake merge statements Adds support for the ON clause of a MERGE statement to be without enclosing parens, as is the convention in Snowflake. * feat: add support for and predicate in merge update * feat: add support for and predicate in merge insert
JSQLParser · Nov 10, 2023 · 36b806d · 36b806d
1 parent 97e9229
commit 36b806d
Show file tree

Hide file tree

Showing 6 changed files with 110 additions and 18 deletions.
diff --git a/src/main/java/net/sf/jsqlparser/statement/merge/Merge.java b/src/main/java/net/sf/jsqlparser/statement/merge/Merge.java
@@ -186,9 +186,8 @@ public String toString() {
         b.append(table);
         b.append(" USING ");
         b.append(fromItem);
-        b.append(" ON (");
+        b.append(" ON ");
         b.append(onCondition);
-        b.append(")");
 
         if (insertFirst && mergeInsert != null) {
             b.append(mergeInsert);

diff --git a/src/main/java/net/sf/jsqlparser/statement/merge/MergeInsert.java b/src/main/java/net/sf/jsqlparser/statement/merge/MergeInsert.java
@@ -20,10 +20,19 @@
 
 public class MergeInsert implements Serializable {
 
+    private Expression andPredicate;
     private ExpressionList<Column> columns;
     private ExpressionList<Expression> values;
     private Expression whereCondition;
 
+    public Expression getAndPredicate() {
+        return andPredicate;
+    }
+
+    public void setAndPredicate(Expression andPredicate) {
+        this.andPredicate = andPredicate;
+    }
+
     public ExpressionList<Column> getColumns() {
         return columns;
     }
@@ -50,12 +59,25 @@ public void setWhereCondition(Expression whereCondition) {
 
     @Override
     public String toString() {
-        return " WHEN NOT MATCHED THEN INSERT "
-                + (columns != null ? columns.toString() : "")
-                + " VALUES " + values.toString()
-                + (whereCondition != null
-                        ? " WHERE " + whereCondition
-                        : "");
+        StringBuilder b = new StringBuilder();
+        b.append(" WHEN NOT MATCHED");
+        if (andPredicate != null) {
+            b.append(" AND ").append(andPredicate.toString());
+        }
+        b.append(" THEN INSERT ");
+        if (columns != null) {
+            b.append(columns.toString());
+        }
+        b.append(" VALUES ").append(values.toString());
+        if (whereCondition != null) {
+            b.append(" WHERE ").append(whereCondition.toString());
+        }
+        return b.toString();
+    }
+
+    public MergeInsert withAndPredicate(Expression andPredicate) {
+        this.setAndPredicate(andPredicate);
+        return this;
     }
 
     public MergeInsert withColumns(ExpressionList<Column> columns) {
@@ -95,6 +117,10 @@ public MergeInsert withWhereCondition(Expression whereCondition) {
         return this;
     }
 
+    public <E extends Expression> E getAndPredicate(Class<E> type) {
+        return type.cast(getAndPredicate());
+    }
+
     public <E extends Expression> E getWhereCondition(Class<E> type) {
         return type.cast(getWhereCondition());
     }

diff --git a/src/main/java/net/sf/jsqlparser/statement/merge/MergeUpdate.java b/src/main/java/net/sf/jsqlparser/statement/merge/MergeUpdate.java
@@ -18,9 +18,13 @@
 public class MergeUpdate implements Serializable {
 
     private List<UpdateSet> updateSets;
+    private Expression andPredicate;
     private Expression whereCondition;
     private Expression deleteWhereCondition;
 
+    public MergeUpdate() {
+    }
+
     public MergeUpdate(List<UpdateSet> updateSets) {
         this.updateSets = updateSets;
     }
@@ -34,6 +38,14 @@ public MergeUpdate setUpdateSets(List<UpdateSet> updateSets) {
         return this;
     }
 
+    public Expression getAndPredicate() {
+        return andPredicate;
+    }
+
+    public void setAndPredicate(Expression andPredicate) {
+        this.andPredicate = andPredicate;
+    }
+
     public Expression getWhereCondition() {
         return whereCondition;
     }
@@ -53,7 +65,11 @@ public void setDeleteWhereCondition(Expression deleteWhereCondition) {
     @Override
     public String toString() {
         StringBuilder b = new StringBuilder();
-        b.append(" WHEN MATCHED THEN UPDATE SET ");
+        b.append(" WHEN MATCHED");
+        if (andPredicate != null) {
+            b.append(" AND ").append(andPredicate.toString());
+        }
+        b.append(" THEN UPDATE SET ");
         UpdateSet.appendUpdateSetsTo(b, updateSets);
 
         if (whereCondition != null) {
@@ -65,6 +81,11 @@ public String toString() {
         return b.toString();
     }
 
+    public MergeUpdate withAndPredicate(Expression andPredicate) {
+        this.setAndPredicate(andPredicate);
+        return this;
+    }
+
     public MergeUpdate withWhereCondition(Expression whereCondition) {
         this.setWhereCondition(whereCondition);
         return this;
@@ -75,6 +96,10 @@ public MergeUpdate withDeleteWhereCondition(Expression deleteWhereCondition) {
         return this;
     }
 
+    public <E extends Expression> E getAndPredicate(Class<E> type) {
+        return type.cast(getAndPredicate());
+}
+
     public <E extends Expression> E getWhereCondition(Class<E> type) {
         return type.cast(getWhereCondition());
     }

diff --git a/src/main/java/net/sf/jsqlparser/util/deparser/StatementDeParser.java b/src/main/java/net/sf/jsqlparser/util/deparser/StatementDeParser.java
@@ -216,9 +216,8 @@ public void visit(Merge merge) {
         buffer.append(" USING ");
         merge.getFromItem().accept(selectDeParser);
 
-        buffer.append(" ON (");
+        buffer.append(" ON ");
         merge.getOnCondition().accept(expressionDeParser);
-        buffer.append(")");
 
         MergeInsert mergeInsert = merge.getMergeInsert();
         MergeUpdate mergeUpdate = merge.getMergeUpdate();
@@ -227,7 +226,12 @@ public void visit(Merge merge) {
         }
 
         if (mergeUpdate != null) {
-            buffer.append(" WHEN MATCHED THEN UPDATE SET ");
+            buffer.append(" WHEN MATCHED");
+            if (mergeUpdate.getAndPredicate() != null) {
+                buffer.append(" AND ");
+                mergeUpdate.getAndPredicate().accept(expressionDeParser);
+            }
+            buffer.append(" THEN UPDATE SET ");
             deparseUpdateSets(mergeUpdate.getUpdateSets(), buffer, expressionDeParser);
 
             if (mergeUpdate.getWhereCondition() != null) {
@@ -251,7 +255,12 @@ public void visit(Merge merge) {
     }
 
     private void deparseMergeInsert(MergeInsert mergeInsert) {
-        buffer.append(" WHEN NOT MATCHED THEN INSERT ");
+        buffer.append(" WHEN NOT MATCHED");
+        if (mergeInsert.getAndPredicate() != null) {
+            buffer.append(" AND ");
+            mergeInsert.getAndPredicate().accept(expressionDeParser);
+        }
+        buffer.append(" THEN INSERT ");
         if (mergeInsert.getColumns() != null) {
             mergeInsert.getColumns().accept(expressionDeParser);
         }

diff --git a/src/main/jjtree/net/sf/jsqlparser/parser/JSqlParserCC.jjt b/src/main/jjtree/net/sf/jsqlparser/parser/JSqlParserCC.jjt
@@ -1659,7 +1659,7 @@ Statement Merge( List<WithItem> with ) : {
 {
     <K_MERGE> { merge.setOracleHint(getOracleHint()); } <K_INTO> table=TableWithAlias() { merge.setTable(table); }
     <K_USING> fromItem = FromItem() { merge.setFromItem(fromItem); }
-    <K_ON> "(" condition = Expression() { merge.setOnCondition(condition); } ")"
+    <K_ON> condition = Expression() { merge.setOnCondition(condition); }
 
     [
         ( LOOKAHEAD(2) update = MergeUpdateClause() { merge.setMergeUpdate(update); }
@@ -1675,14 +1675,17 @@ Statement Merge( List<WithItem> with ) : {
 }
 
 MergeUpdate MergeUpdateClause() : {
-    MergeUpdate mu;
+    MergeUpdate mu = new MergeUpdate();
     List<UpdateSet> updateSets;
+    Expression predicate;
     Expression condition;
 }
 {
-    <K_WHEN> <K_MATCHED> <K_THEN>  <K_UPDATE>
+    <K_WHEN> <K_MATCHED>
+    [ <K_AND> predicate = Expression() { mu.setAndPredicate(predicate); } ]
+    <K_THEN> <K_UPDATE>
     <K_SET>
-    updateSets = UpdateSets() { mu = new MergeUpdate(updateSets); }
+    updateSets = UpdateSets() { mu.setUpdateSets(updateSets); }
 
     [ <K_WHERE> condition = Expression() { mu.setWhereCondition(condition); }]
     [ <K_DELETE> <K_WHERE> condition = Expression() { mu.setDeleteWhereCondition(condition); } ]
@@ -1692,12 +1695,15 @@ MergeUpdate MergeUpdateClause() : {
 
 MergeInsert MergeInsertClause() : {
     MergeInsert mi = new MergeInsert();
+    Expression predicate;
     ExpressionList<Column> columns;
     ExpressionList expList;
     Expression condition;
 }
 {
-    <K_WHEN> <K_NOT> <K_MATCHED> <K_THEN>
+    <K_WHEN> <K_NOT> <K_MATCHED>
+    [ <K_AND> predicate = Expression() { mi.setAndPredicate(predicate); } ]
+    <K_THEN>
         <K_INSERT>
         [ "(" columns = ColumnList() ")"
             {

diff --git a/src/test/java/net/sf/jsqlparser/statement/merge/MergeTest.java b/src/test/java/net/sf/jsqlparser/statement/merge/MergeTest.java
@@ -236,4 +236,31 @@ public void testOutputClause() throws JSQLParserException {
                 + "        TAB_MergeActions_RoomLocation";
         assertSqlCanBeParsedAndDeparsed(sqlStr, true);
     }
+
+    @Test
+    public void testSnowflakeMergeStatementSimple() throws JSQLParserException {
+        String sql = "MERGE INTO target\n" +
+                "  USING src ON target.k = src.k\n" +
+                "  WHEN MATCHED THEN UPDATE SET target.v = src.v";
+
+        assertSqlCanBeParsedAndDeparsed(sql, true);
+    }
+
+    @Test
+    public void testSnowflakeMergeStatementWithMatchedAndPredicate() throws JSQLParserException {
+        String sql = "MERGE INTO target\n" +
+                "  USING src ON target.k = src.k\n" +
+                "  WHEN MATCHED AND src.v = 11 THEN UPDATE SET target.v = src.v";
+
+        assertSqlCanBeParsedAndDeparsed(sql, true);
+    }
+
+    @Test
+    void testSnowflakeMergeStatementWithNotMatchedAndPredicate() throws JSQLParserException {
+        String sql = "MERGE INTO target USING (select k, max(v) as v from src group by k) AS b ON target.k = b.k\n" +
+                "  WHEN MATCHED THEN UPDATE SET target.v = b.v\n" +
+                "  WHEN NOT MATCHED AND b.v != 11 THEN INSERT (k, v) VALUES (b.k, b.v)";
+
+        assertSqlCanBeParsedAndDeparsed(sql, true);
+    }
 }