Support doris/sr insert overwrite statement.

alibaba · Dec 23, 2024 · 14d444a · 14d444a
1 parent f7266bc
commit 14d444a
Show file tree

Hide file tree

Showing 9 changed files with 86 additions and 34 deletions.
diff --git a/.../src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryStatementParser.java b/.../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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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<SQLCommentHint> headHints = x.getHeadHintsDirect();
         if (headHints != null) {
@@ -3787,15 +3790,20 @@ 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 ");
             }
         }
 
         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
@@ -1,3 +1,13 @@
 DROP TABLE IF EXISTS `biz_jxfw_class_attend_data_v1`;
 --------------------
-DROP TABLE IF EXISTS `biz_jxfw_class_attend_data_v1`;
+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
diff --git 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
+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