From c9f4d9d194afb5b653895d154a0af8d2a65269ab Mon Sep 17 00:00:00 2001 From: wenshao Date: Thu, 12 Dec 2024 10:29:13 +0800 Subject: [PATCH 1/2] improved sql parser bigquery support --- .../sql/ast/expr/SQLMethodInvokeExpr.java | 6 ++++ .../bigquery/parser/BigQueryExprParser.java | 2 ++ .../druid/sql/parser/SQLExprParser.java | 2 +- .../druid/sql/parser/SQLSelectParser.java | 1 + .../druid/bvt/sql/bigquery/AggregateTest.java | 30 +++++++++++++++++++ 5 files changed, 40 insertions(+), 1 deletion(-) create mode 100644 core/src/test/java/com/alibaba/druid/bvt/sql/bigquery/AggregateTest.java diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLMethodInvokeExpr.java b/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLMethodInvokeExpr.java index ba7563c348..f15adb699f 100644 --- a/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLMethodInvokeExpr.java +++ b/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLMethodInvokeExpr.java @@ -170,6 +170,12 @@ public void addArgument(SQLExpr arg) { this.arguments.add(arg); } + public void addArguments(List args) { + for (SQLExpr arg : args) { + addArgument(arg); + } + } + public SQLExpr getOwner() { return this.owner; } diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryExprParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryExprParser.java index ea8536d4b2..33cdcfd400 100644 --- a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryExprParser.java +++ b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryExprParser.java @@ -31,6 +31,8 @@ public class BigQueryExprParser extends SQLExprParser { "FIRST_VALUE", "GROUPING", "LAST_VALUE", + "LAG", + "LEAD", "LOGICAL_AND", "LOGICAL_OR", "MAX", diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java b/core/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java index 4845e11c6f..6a55461fde 100644 --- a/core/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java +++ b/core/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java @@ -1984,7 +1984,7 @@ protected SQLExpr methodRest(SQLExpr expr, boolean acceptLPAREN) { if (lexer.token == Token.OVER) { if (aggregateExpr == null) { aggregateExpr = new SQLAggregateExpr(methodName); - aggregateExpr.getArguments().addAll(methodInvokeExpr.getArguments()); + aggregateExpr.addArguments(methodInvokeExpr.getArguments()); } over(aggregateExpr); } diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java b/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java index b889aa94b8..afd615a4aa 100644 --- a/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java +++ b/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java @@ -1820,6 +1820,7 @@ public SQLTableSource parseTableSourceRest(SQLTableSource tableSource) { SQLTableSource unnestTableSource = parseUnnestTableSource(); if (unnestTableSource != null) { if (lexer.identifierEquals(FnvHash.Constants.CROSS) + || lexer.token == Token.CROSS || lexer.token == Token.LEFT || lexer.token == Token.RIGHT || lexer.token == Token.COMMA diff --git a/core/src/test/java/com/alibaba/druid/bvt/sql/bigquery/AggregateTest.java b/core/src/test/java/com/alibaba/druid/bvt/sql/bigquery/AggregateTest.java new file mode 100644 index 0000000000..2613dbefbb --- /dev/null +++ b/core/src/test/java/com/alibaba/druid/bvt/sql/bigquery/AggregateTest.java @@ -0,0 +1,30 @@ +package com.alibaba.druid.bvt.sql.bigquery; + +import com.alibaba.druid.DbType; +import com.alibaba.druid.sql.SQLUtils; +import com.alibaba.druid.sql.ast.expr.SQLAggregateExpr; +import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; +import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; +import org.junit.Test; + +import static org.junit.Assert.assertSame; + +public class AggregateTest { + @Test + public void test_agg() { + String sql = "SELECT lag(date(d,'Asia/Jakarta')) over(partition by id order by n) FROM t1"; + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseSingleStatement(sql, DbType.bigquery); + SQLSelectQueryBlock queryBlock = stmt.getSelect().getQueryBlock(); + SQLAggregateExpr expr = (SQLAggregateExpr) queryBlock.getSelectList().get(0).getExpr(); + assertSame(expr, expr.getArgument(0).getParent()); + } + + @Test + public void test_agg_1() { + String sql = "SELECT xx(date(d,'Asia/Jakarta')) over(partition by id order by n) FROM t1"; + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseSingleStatement(sql, DbType.bigquery); + SQLSelectQueryBlock queryBlock = stmt.getSelect().getQueryBlock(); + SQLAggregateExpr expr = (SQLAggregateExpr) queryBlock.getSelectList().get(0).getExpr(); + assertSame(expr, expr.getArgument(0).getParent()); + } +} From 826373292397bd4f9ae391a2e02f82bd24c4089d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E6=9E=B8?= Date: Mon, 16 Dec 2024 22:27:08 +0800 Subject: [PATCH 2/2] Rewrite merge insert into row. --- .../sql/ast/statement/SQLMergeStatement.java | 9 +++++++ .../odps/visitor/OdpsOutputVisitor.java | 5 ++++ .../druid/sql/parser/SQLStatementParser.java | 5 ++-- .../sql/visitor/SQLASTOutputVisitor.java | 27 ++++++++++++------- .../test/resources/bvt/parser/bigquery/0.txt | 6 +++++ core/src/test/resources/bvt/parser/odps/0.txt | 10 ++++++- 6 files changed, 49 insertions(+), 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 index 4059c49b11..9d8842fa52 100644 --- 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 @@ -141,6 +141,7 @@ public WhenUpdate cloneTo() { } public static class WhenInsert extends When { + private boolean insertRow; private List columns = new ArrayList(); private List values = new ArrayList(); @@ -180,6 +181,14 @@ public void setColumns(List columns) { this.columns = columns; } + public boolean isInsertRow() { + return insertRow; + } + + public void setInsertRow(boolean insertRow) { + this.insertRow = insertRow; + } + public List 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 index 967664000a..195a43533d 100644 --- 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 index 5983e0a966..689d71afa8 100644 --- 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 index fff1ee9b9e..2eed446d09 100644 --- 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 index f0b67de991..88af927d3c 100644 --- 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 index 36f2f5b99e..93341b4f52 100644 --- 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 \ No newline at end of file +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 * \ No newline at end of file