From 14d444acfbfa0ec6d14abd9896ca015b5bc2a9b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E6=9E=B8?= Date: Mon, 23 Dec 2024 17:02:45 +0800 Subject: [PATCH] Support doris/sr insert overwrite statement. --- .../parser/BigQueryStatementParser.java | 20 -------- .../sql/dialect/doris/parser/DorisLexer.java | 1 + .../doris/parser/DorisStatementParser.java | 8 ++++ .../doris/visitor/DorisOutputVisitor.java | 6 +++ .../starrocks/parser/StarRocksLexer.java | 1 + .../druid/sql/parser/SQLStatementParser.java | 48 +++++++++++++++---- .../sql/visitor/SQLASTOutputVisitor.java | 12 ++++- .../src/test/resources/bvt/parser/doris/2.txt | 12 ++++- .../test/resources/bvt/parser/starrocks/2.txt | 12 ++++- 9 files changed, 86 insertions(+), 34 deletions(-) diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryStatementParser.java index 915444d984..6d37fa2d06 100644 --- a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryStatementParser.java +++ b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryStatementParser.java @@ -236,26 +236,6 @@ public SQLStatement parseBlock() { return block; } - protected void parseInsert0(SQLInsertInto insertStatement, boolean acceptSubQuery) { - if (lexer.token() == Token.IDENTIFIER) { - SQLName tableName = this.exprParser.name(); - insertStatement.setTableName(tableName); - - if (lexer.token() == Token.LITERAL_ALIAS) { - insertStatement.setAlias(tableAlias()); - } - - parseInsert0Hints(insertStatement, false); - - if (lexer.token() == Token.IDENTIFIER) { - insertStatement.setAlias(lexer.stringVal()); - lexer.nextToken(); - } - } - - super.parseInsert0(insertStatement, acceptSubQuery); - } - protected void createViewAs(SQLCreateViewStatement createView) { if (lexer.nextIfIdentifier(FnvHash.Constants.OPTIONS)) { exprParser.parseAssignItem(createView.getOptions(), createView); diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/doris/parser/DorisLexer.java b/core/src/main/java/com/alibaba/druid/sql/dialect/doris/parser/DorisLexer.java index d59998f6e6..05f3c9f6ae 100644 --- a/core/src/main/java/com/alibaba/druid/sql/dialect/doris/parser/DorisLexer.java +++ b/core/src/main/java/com/alibaba/druid/sql/dialect/doris/parser/DorisLexer.java @@ -41,6 +41,7 @@ protected Keywords loadKeywords() { map.put("ROW", Token.ROW); map.put("MOD", Token.MOD); map.put("RLIKE", Token.RLIKE); + map.put("OVERWRITE", Token.OVERWRITE); return new Keywords(map); } diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/doris/parser/DorisStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/doris/parser/DorisStatementParser.java index 18c14ca5d7..6c0a63448a 100644 --- a/core/src/main/java/com/alibaba/druid/sql/dialect/doris/parser/DorisStatementParser.java +++ b/core/src/main/java/com/alibaba/druid/sql/dialect/doris/parser/DorisStatementParser.java @@ -1,9 +1,11 @@ package com.alibaba.druid.sql.dialect.doris.parser; import com.alibaba.druid.DbType; +import com.alibaba.druid.sql.ast.statement.SQLInsertInto; import com.alibaba.druid.sql.dialect.starrocks.parser.StarRocksStatementParser; import com.alibaba.druid.sql.parser.SQLExprParser; import com.alibaba.druid.sql.parser.SQLParserFeature; +import com.alibaba.druid.sql.parser.Token; public class DorisStatementParser extends StarRocksStatementParser { @@ -24,4 +26,10 @@ public DorisSelectParser createSQLSelectParser() { public DorisCreateTableParser getSQLCreateTableParser() { return new DorisCreateTableParser(this.exprParser); } + + @Override + protected void parseInsertOverwrite(SQLInsertInto insertStatement) { + insertStatement.setOverwrite(true); + accept(Token.TABLE); + } } diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/doris/visitor/DorisOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/dialect/doris/visitor/DorisOutputVisitor.java index 5adc3e8a47..73c73feb68 100644 --- a/core/src/main/java/com/alibaba/druid/sql/dialect/doris/visitor/DorisOutputVisitor.java +++ b/core/src/main/java/com/alibaba/druid/sql/dialect/doris/visitor/DorisOutputVisitor.java @@ -6,6 +6,7 @@ import com.alibaba.druid.sql.ast.SQLPivot; import com.alibaba.druid.sql.ast.SQLSetQuantifier; import com.alibaba.druid.sql.ast.SQLUnpivot; +import com.alibaba.druid.sql.ast.statement.SQLInsertStatement; import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; import com.alibaba.druid.sql.ast.statement.SQLTableSampling; import com.alibaba.druid.sql.dialect.doris.ast.DorisExprTableSource; @@ -135,4 +136,9 @@ public boolean visit(SQLTableSampling x) { print(')'); return false; } + + @Override + public void printInsertOverWrite(SQLInsertStatement x) { + print0(ucase ? "INSERT OVERWRITE TABLE " : "insert overwrite table "); + } } diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/starrocks/parser/StarRocksLexer.java b/core/src/main/java/com/alibaba/druid/sql/dialect/starrocks/parser/StarRocksLexer.java index 304caa5f5f..64babc7514 100644 --- a/core/src/main/java/com/alibaba/druid/sql/dialect/starrocks/parser/StarRocksLexer.java +++ b/core/src/main/java/com/alibaba/druid/sql/dialect/starrocks/parser/StarRocksLexer.java @@ -34,6 +34,7 @@ protected Keywords loadKeywords() { map.put("MOD", Token.MOD); map.put("RLIKE", Token.RLIKE); map.put("USING", Token.USING); + map.put("OVERWRITE", Token.OVERWRITE); return new Keywords(map); } 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 689d71afa8..a5c5f34562 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 @@ -3783,21 +3783,49 @@ protected void parseInsert0(SQLInsertInto insertStatement) { protected void parseInsert0Hints(SQLInsertInto insertStatement, boolean isInsert) { } + protected void parseInsertOverwrite(SQLInsertInto insertStatement) { + insertStatement.setOverwrite(true); + } protected void parseInsert0(SQLInsertInto insertStatement, boolean acceptSubQuery) { - if (lexer.nextIf(INTO)) { - SQLName tableName = this.exprParser.name(); - insertStatement.setTableName(tableName); + if (lexer.nextIf(OVERWRITE) || lexer.nextIfIdentifier(Constants.OVERWRITE)) { + parseInsertOverwrite(insertStatement); + } else if (lexer.nextIf(INTO)) { + insertStatement.setOverwrite(false); + } - if (lexer.token == Token.LITERAL_ALIAS) { - insertStatement.setAlias(tableAlias()); - } + SQLName tableName = this.exprParser.name(); + insertStatement.setTableName(tableName); + + if (lexer.token == Token.LITERAL_ALIAS) { + insertStatement.setAlias(tableAlias()); + } - parseInsert0Hints(insertStatement, false); + parseInsert0Hints(insertStatement, false); - if (lexer.token == Token.IDENTIFIER) { - insertStatement.setAlias(lexer.stringVal()); - lexer.nextToken(); + if (lexer.token == Token.IDENTIFIER) { + insertStatement.setAlias(lexer.stringVal()); + lexer.nextToken(); + } + + if (lexer.token == Token.PARTITION) { + lexer.nextToken(); + accept(Token.LPAREN); + for (; ; ) { + SQLAssignItem ptExpr = new SQLAssignItem(); + ptExpr.setTarget(this.exprParser.name()); + if (lexer.token == Token.EQ || lexer.token == Token.EQEQ) { + lexer.nextTokenValue(); + SQLExpr ptValue = this.exprParser.expr(); + ptExpr.setValue(ptValue); + } + insertStatement.addPartition(ptExpr); + if (!(lexer.token == (Token.COMMA))) { + break; + } else { + lexer.nextToken(); + } } + accept(Token.RPAREN); } if (lexer.token == (Token.LPAREN)) { 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 9675077009..dfeb96867c 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 @@ -3765,6 +3765,9 @@ public boolean visit(SQLCurrentOfCursorExpr x) { return false; } + public void printInsertOverWrite(SQLInsertStatement x) { + print0(ucase ? "INSERT OVERWRITE " : "insert overwrite "); + } public boolean visit(SQLInsertStatement x) { List headHints = x.getHeadHintsDirect(); if (headHints != null) { @@ -3787,8 +3790,8 @@ public boolean visit(SQLInsertStatement x) { if (x.isUpsert()) { print0(ucase ? "UPSERT INTO " : "upsert into "); } else { - if (x.isOverwrite() && dbType == DbType.odps) { - print0(ucase ? "INSERT OVERWRITE " : "insert overwrite "); + if (x.isOverwrite()) { + printInsertOverWrite(x); } else { print0(ucase ? "INSERT INTO " : "insert into "); } @@ -3796,6 +3799,11 @@ public boolean visit(SQLInsertStatement x) { x.getTableSource().accept(this); + if (x.getPartitions() != null && !x.getPartitions().isEmpty()) { + print0(ucase ? " PARTITION (" : " partition ("); + printAndAccept(x.getPartitions(), ", "); + print(')'); + } String columnsString = x.getColumnsString(); if (columnsString != null) { print0(columnsString); diff --git a/core/src/test/resources/bvt/parser/doris/2.txt b/core/src/test/resources/bvt/parser/doris/2.txt index f8f82c0924..adabdd11e8 100644 --- a/core/src/test/resources/bvt/parser/doris/2.txt +++ b/core/src/test/resources/bvt/parser/doris/2.txt @@ -1,3 +1,13 @@ DROP TABLE IF EXISTS `biz_jxfw_class_attend_data_v1`; -------------------- -DROP TABLE IF EXISTS `biz_jxfw_class_attend_data_v1`; \ No newline at end of file +DROP TABLE IF EXISTS `biz_jxfw_class_attend_data_v1`; +------------------------------------------------------------------------------------------------------------------------ +insert overwrite table a.b partition (p20210101) +select +c, +d +from e +-------------------- +INSERT OVERWRITE TABLE a.b PARTITION (p20210101) +SELECT c, d +FROM e \ No newline at end of file diff --git a/core/src/test/resources/bvt/parser/starrocks/2.txt b/core/src/test/resources/bvt/parser/starrocks/2.txt index 5b34873a0d..8b7ebabf16 100644 --- a/core/src/test/resources/bvt/parser/starrocks/2.txt +++ b/core/src/test/resources/bvt/parser/starrocks/2.txt @@ -86,4 +86,14 @@ CREATE TABLE example_db.table_hash ( k3 CHAR(10) COMMENT 'string column', k4 INT NOT NULL DEFAULT "1" COMMENT 'int column' ) -DISTRIBUTED BY HASH (k1) BUCKETS 32 \ No newline at end of file +DISTRIBUTED BY HASH (k1) BUCKETS 32 +------------------------------------------------------------------------------------------------------------------------ +insert overwrite a.b partition (p20210101) +select +c, +d +from e +-------------------- +INSERT OVERWRITE a.b PARTITION (p20210101) +SELECT c, d +FROM e \ No newline at end of file