Rewrite merge insert row. (#6290)

* improved sql parser bigquery support * Rewrite merge insert into row. --------- Co-authored-by: wenshao <[email protected]>
alibaba · Dec 16, 2024 · 3a68c95 · 3a68c95
1 parent eeab8df
commit 3a68c95
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 13 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLMergeStatement.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLMergeStatement.java
@@ -156,6 +156,7 @@ public boolean replace(SQLExpr expr, SQLExpr target) {
     }
 
     public static class WhenInsert extends When {
+        private boolean insertRow;
         private List<SQLExpr> columns = new ArrayList<SQLExpr>();
         private List<SQLExpr> values = new ArrayList<SQLExpr>();
 
@@ -195,6 +196,14 @@ public void setColumns(List<SQLExpr> columns) {
             this.columns = columns;
         }
 
+        public boolean isInsertRow() {
+            return insertRow;
+        }
+
+        public void setInsertRow(boolean insertRow) {
+            this.insertRow = insertRow;
+        }
+
         public List<SQLExpr> getValues() {
             return values;
         }

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/odps/visitor/OdpsOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/dialect/odps/visitor/OdpsOutputVisitor.java
@@ -1120,4 +1120,9 @@ protected void printMethodParameters(SQLMethodInvokeExpr x) {
         }
         super.printMethodParameters(x);
     }
+
+    @Override
+    public void printMergeInsertRow() {
+        print(" *");
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java
@@ -5672,13 +5672,14 @@ protected boolean parseMergeWhen(SQLMergeStatement stmt) {
             exprParser.exprList(insertClause.getColumns(), insertClause);
             accept(Token.RPAREN);
         }
-        if (lexer.nextIfIdentifier("ROW")) {
-            insertClause.getValues().add(new SQLIdentifierExpr("ROW"));
+        if (lexer.nextIfIdentifier("ROW") || lexer.nextIfIdentifier("*") || lexer.nextIf(Token.STAR)) {
+            insertClause.setInsertRow(true);
         } else {
             accept(Token.VALUES);
             accept(Token.LPAREN);
             exprParser.exprList(insertClause.getValues(), insertClause);
             accept(Token.RPAREN);
+            insertClause.setInsertRow(false);
         }
 
         if (lexer.token == Token.WHERE) {

diff --git a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
@@ -8593,17 +8593,21 @@ public boolean visit(WhenInsert x) {
             by.accept(this);
         }
         print0(ucase ? " THEN INSERT" : " then insert");
-        if (x.getColumns().size() > 0) {
-            printAndAccept(" (", ")", x.getColumns(), ", ", 5);
+        if (x.isInsertRow()) {
+            printMergeInsertRow();
+        } else {
+            if (x.getColumns().size() > 0) {
+                printAndAccept(" (", ")", x.getColumns(), ", ", 5);
+            }
+            println();
+            printAndAccept(
+                    ucase ? "VALUES (" : "values (",
+                    ")",
+                    x.getValues(),
+                    ", ",
+                    5
+            );
         }
-        println();
-        printAndAccept(
-                ucase ? "VALUES (" : "values (",
-                ")",
-                x.getValues(),
-                ", ",
-                5
-        );
         if (x.getWhere() != null) {
             this.indentCount++;
             println();
@@ -8615,6 +8619,9 @@ public boolean visit(WhenInsert x) {
         return false;
     }
 
+    public void printMergeInsertRow() {
+        print(ucase ? " ROW" : "row");
+    }
     @Override
     public boolean visit(SQLErrorLoggingClause x) {
         print0(ucase ? "LOG ERRORS " : "log errors ");

diff --git a/core/src/test/resources/bvt/parser/bigquery/0.txt b/core/src/test/resources/bvt/parser/bigquery/0.txt
@@ -1,3 +1,9 @@
+merge into a as target using b as source on a.id = b.id when not matched then insert row
+--------------------
+MERGE INTO a target
+USING b source ON a.id = b.id
+WHEN NOT MATCHED THEN INSERT ROW
+------------------------------------------------------------------------------------------------------------------------
 SELECT * EXCEPT (predicted_label),
       predicted_label AS predicted_label1
 FROM

diff --git a/core/src/test/resources/bvt/parser/odps/0.txt b/core/src/test/resources/bvt/parser/odps/0.txt
@@ -6,4 +6,12 @@ FROM a
 select named_struct(a.b, b, a.c, c, a.d, d) from a
 --------------------
 SELECT named_struct(a.b, b, a.c, c, a.d, d)
-FROM a
+FROM a
+------------------------------------------------------------------------------------------------------------------------
+MERGE INTO a target
+USING b source ON a.id = b.id
+WHEN NOT MATCHED THEN INSERT *
+--------------------
+MERGE INTO a target
+USING b source ON a.id = b.id
+WHEN NOT MATCHED THEN INSERT *