增加TiDB的Split Region语法解析支持 #5219

增加TiDB的Split Region语法解析支持 #5219

core/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/TidbSplitTableStatement.java

@@ -0,0 +1,125 @@
+/*
+ * Copyright 1999-2017 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.druid.sql.dialect.mysql.ast.statement;
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLName;
+import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
+import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitor;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * @author lizongbo
+ * @see  <a href="https://docs.pingcap.com/zh/tidb/stable/sql-statement-split-region">...</a>
+ */
+public class TidbSplitTableStatement extends MySqlStatementImpl {
+    //region for
+    private boolean splitSyntaxOptionRegionFor;
+    //partition
+    private boolean splitSyntaxOptionPartition;
+
+    private SQLExprTableSource tableName;
+
+    private List<SQLExpr> partitionNameListOptions = new ArrayList<>();
+
+    private SQLName indexName;
+
+    private List<List<SQLExpr>> splitOptionBys = new ArrayList<>();
+    private List<SQLExpr> splitOptionBetween;
+    private List<SQLExpr> splitOptionAnd;
+    private long splitOptionRegions;
+
+    public boolean isSplitSyntaxOptionRegionFor() {
+        return splitSyntaxOptionRegionFor;
+    }
+
+    public void setSplitSyntaxOptionRegionFor(boolean splitSyntaxOptionRegionFor) {
+        this.splitSyntaxOptionRegionFor = splitSyntaxOptionRegionFor;
+    }
+
+    public boolean isSplitSyntaxOptionPartition() {
+        return splitSyntaxOptionPartition;
+    }
+
+    public void setSplitSyntaxOptionPartition(boolean splitSyntaxOptionPartition) {
+        this.splitSyntaxOptionPartition = splitSyntaxOptionPartition;
+    }
+
+    public SQLExprTableSource getTableName() {
+        return tableName;
+    }
+
+    public void setTableName(SQLExprTableSource tableName) {
+        this.tableName = tableName;
+    }
+
+    public List<SQLExpr> getPartitionNameListOptions() {
+        return partitionNameListOptions;
+    }
+
+    public void setPartitionNameListOptions(List<SQLExpr> partitionNameListOptions) {
+        this.partitionNameListOptions = partitionNameListOptions;
+    }
+
+    public SQLName getIndexName() {
+        return indexName;
+    }
+
+    public void setIndexName(SQLName indexName) {
+        this.indexName = indexName;
+    }
+
+    public List<List<SQLExpr>> getSplitOptionBys() {
+        return splitOptionBys;
+    }
+
+    public void setSplitOptionBys(List<List<SQLExpr>> splitOptionBys) {
+        this.splitOptionBys = splitOptionBys;
+    }
+
+    public List<SQLExpr> getSplitOptionBetween() {
+        return splitOptionBetween;
+    }
+
+    public void setSplitOptionBetween(List<SQLExpr> splitOptionBetween) {
+        this.splitOptionBetween = splitOptionBetween;
+    }
+
+    public List<SQLExpr> getSplitOptionAnd() {
+        return splitOptionAnd;
+    }
+
+    public void setSplitOptionAnd(List<SQLExpr> splitOptionAnd) {
+        this.splitOptionAnd = splitOptionAnd;
+    }
+
+    public long getSplitOptionRegions() {
+        return splitOptionRegions;
+    }
+
+    public void setSplitOptionRegions(long splitOptionRegions) {
+        this.splitOptionRegions = splitOptionRegions;
+    }
+
+    public void accept0(MySqlASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            this.getTableName().accept(visitor);
+        }
+        visitor.endVisit(this);
+    }
+}

core/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java

@@ -1805,10 +1805,89 @@ public boolean parseStatementListDialect(List<SQLStatement> statementList) {
             }
             return true;
         }
+        String strVal = lexer.stringVal();
+        if (strVal.equalsIgnoreCase("SPLIT")) {
+            TidbSplitTableStatement stmt = this.parseTiDBSplitTableStatement();
+            statementList.add(stmt);
+            return true;
+        }
 
         return false;
     }
 
+    private TidbSplitTableStatement parseTiDBSplitTableStatement() {
+        TidbSplitTableStatement stmt = new TidbSplitTableStatement();
+        lexer.nextToken();
+        Lexer.SavePoint mark = lexer.mark();
+        String strVal = lexer.stringVal();
+        if (lexer.token() == PARTITION) {
+            accept(PARTITION);
+            stmt.setSplitSyntaxOptionPartition(true);
+        } else if (strVal.equalsIgnoreCase("REGION")) {
+            acceptIdentifier("REGION");
+            accept(FOR);
+            stmt.setSplitSyntaxOptionRegionFor(true);
+        } else {
+            lexer.reset(mark);
+        }
+        strVal = lexer.stringVal();
+        accept(TABLE);
+        SQLName tableNameTmp = this.exprParser.name();
+        SQLExprTableSource sqlExprTableSource = new SQLExprTableSource(tableNameTmp);
+        stmt.setTableName(sqlExprTableSource);
+        strVal = lexer.stringVal();
+        if (lexer.token() == INDEX) {
+            accept(INDEX);
+            SQLName indexName = this.exprParser.name();
+            stmt.setIndexName(indexName);
+        } else if (lexer.token() == PARTITION) {
+            accept(PARTITION);
+            accept(LPAREN);
+            this.exprParser.exprList(stmt.getPartitionNameListOptions(), stmt);
+            accept(RPAREN);
+            if (lexer.token() == INDEX) {
+                accept(INDEX);
+                SQLName indexName = this.exprParser.name();
+                stmt.setIndexName(indexName);
+            }
+        }
+        if (lexer.token() == BETWEEN) {
+            accept(BETWEEN);
+            accept(LPAREN);
+            List<SQLExpr> sqlExprBetweens = new ArrayList<>();
+            this.exprParser.exprList(sqlExprBetweens, stmt);
+            stmt.setSplitOptionBetween(sqlExprBetweens);
+            accept(RPAREN);
+            accept(AND);
+            accept(LPAREN);
+            List<SQLExpr> sqlExprAnds = new ArrayList<>();
+            this.exprParser.exprList(sqlExprAnds, stmt);
+            stmt.setSplitOptionAnd(sqlExprAnds);
+            accept(RPAREN);
+            acceptIdentifier("REGIONS");
+            SQLIntegerExpr num = this.exprParser.integerExpr();
+            stmt.setSplitOptionRegions(num.getNumber().longValue());
+        } else if (lexer.token() == BY) {
+            accept(BY);
+            accept(LPAREN);
+            List<SQLExpr> byItems = new ArrayList<>();
+            this.exprParser.exprList(byItems, stmt);
+            accept(RPAREN);
+            stmt.getSplitOptionBys().add(byItems);
+            while (lexer.token() == COMMA) {
+                accept(COMMA);
+                accept(LPAREN);
+                byItems = new ArrayList<>();
+                this.exprParser.exprList(byItems, stmt);
+                accept(RPAREN);
+                stmt.getSplitOptionBys().add(byItems);
+            }
+
+        } else if (lexer.token() == PARTITION) {
+            accept(PARTITION);
+        }
+        return stmt;
+    }
     private SQLStatement parseArchive() {
         lexer.nextToken();
         accept(Token.TABLE);

core/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitor.java

@@ -1274,4 +1274,11 @@ default boolean visit(MySqlXAStatement x) {
 
     default void endVisit(MySqlXAStatement x) {
     }
+    default boolean visit(TidbSplitTableStatement x) {
+        return true;
+    }
+
+    default void endVisit(TidbSplitTableStatement x) {
+    }
+
 } //

core/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java

@@ -5442,6 +5442,56 @@ public boolean visit(MySqlJSONTableExpr x) {
         return false;
     }
 
+    public boolean visit(TidbSplitTableStatement x) {
+        print0(ucase ? "SPLIT " : "split ");
+        if (x.isSplitSyntaxOptionRegionFor()) {
+            print0(ucase ? "REGION FOR " : "region for ");
+        }
+        if (x.isSplitSyntaxOptionPartition()) {
+            print0(ucase ? "PARTITION " : "partition ");
+        }
+        print0(ucase ? "TABLE " : "table ");
+        x.getTableName().accept(this);
+        print(' ');
+
+        if (!x.getPartitionNameListOptions().isEmpty()) {
+            print0(ucase ? "PARTITION (" : "partition (");
+            printAndAccept(x.getPartitionNameListOptions(), ",");
+            print(") ");
+        }
+        if (x.getIndexName() != null) {
+            print0(ucase ? "INDEX " : "index ");
+            x.getIndexName().accept(this);
+            print(' ');
+        }
+        if (!x.getSplitOptionBys().isEmpty()) {
+            print0(ucase ? "BY " : "by ");
+            boolean needCommon = false;
+            for (List<SQLExpr> list : x.getSplitOptionBys()) {
+                if (!needCommon) {
+                    needCommon = true;
+                } else {
+                    print0(", ");
+                }
+                print0("(");
+                printlnAndAccept(list, ", ");
+                print0(")");
+            }
+        }
+        if (x.getSplitOptionBetween() != null) {
+            print0(ucase ? "BETWEEN (" : " between (");
+            printAndAccept(x.getSplitOptionBetween(), ", ");
+            print0(ucase ? ") AND (" : ") and (");
+            printAndAccept(x.getSplitOptionAnd(), ", ");
+            print0(") ");
+            print0(ucase ? "REGIONS " : "regions ");
+            print(x.getSplitOptionRegions());
+        }
+
+        return false;
+    }
+
+
     public boolean visit(MySqlJSONTableExpr.Column x) {
         x.getName().accept(this);
 

core/src/test/java/com/alibaba/druid/bvt/sql/mysql/issues/Issue5219.java

@@ -0,0 +1,68 @@
+package com.alibaba.druid.bvt.sql.mysql.issues;
+
+import java.util.Map;
+
+import com.alibaba.druid.DbType;
+import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.parser.SQLParserUtils;
+import com.alibaba.druid.sql.parser.SQLStatementParser;
+import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
+import com.alibaba.druid.stat.TableStat;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author lizongbo
+ * @see <a href="https://docs.pingcap.com/zh/tidb/stable/sql-statement-split-region">Split Region 使用文档</a>
+ * @see <a href="https://github.com/alibaba/druid/issues/5219">Issue来源</a>
+ */
+public class Issue5219 {
+
+    @Test
+    public void test_split_table() throws Exception {
+     //   for (DbType dbType : new DbType[]{DbType.mysql, DbType.tidb}) {
+        for (DbType dbType : new DbType[]{DbType.tidb}) {
+            for (String sql : new String[]{
+                "split TABLE t BETWEEN (-9223372036854775808) AND (9223372036854775807) REGIONS 16;",
+                "split partition TABLE t BETWEEN (-9223372036854775808) AND (9223372036854775807) REGIONS 16;",
+                "split region for TABLE t BETWEEN (-9223372036854775808) AND (9223372036854775807) REGIONS 16;",
+                "SPLIT TABLE t INDEX idx BETWEEN (-9223372036854775808) AND (9223372036854775807) REGIONS 16;",
+                "SPLIT TABLE t1 INDEX idx4 BY (\"a\", \"2000-01-01 00:00:01\"), (\"b\", \"2019-04-17 14:26:19\"), (\"c\", \"\");",
+                "split partition table t partition (p1,p3) between (0) and (10000) regions 2;",
+                "split TABLE t BETWEEN (0) AND (1000000000) REGIONS 16;",
+                "split TABLE t BY (10000), (90000);",
+                "SPLIT TABLE t INDEX idx1 BETWEEN (\"a\") AND (\"z\") REGIONS 25;",
+                "SPLIT TABLE t INDEX idx1 BETWEEN (\"a\") AND (\"{\") REGIONS 26;",
+                "SPLIT TABLE t INDEX idx2 BETWEEN (\"2010-01-01 00:00:00\") AND (\"2020-01-01 00:00:00\") REGIONS 10;",
+                "SPLIT TABLE t INDEX idx2 BETWEEN (\"2020-06-01 00:00:00\") AND (\"2020-07-01 00:00:00\") REGIONS 30;",
+                "SPLIT TABLE t INDEX idx3 BETWEEN (\"2010-01-01 00:00:00\") AND (\"2020-01-01 00:00:00\") REGIONS 10;",
+                "SPLIT TABLE t INDEX idx3 BETWEEN (\"2010-01-01 00:00:00\", \"a\") AND (\"2010-01-01 00:00:00\", \"z\") REGIONS 10;",
+                "SPLIT TABLE t INDEX `PRIMARY` BETWEEN (-9223372036854775808) AND (9223372036854775807) REGIONS 16;",
+                "SPLIT TABLE t1 INDEX idx4 BY (\"a\", \"2000-01-01 00:00:01\"), (\"b\", \"2019-04-17 14:26:19\"), (\"c\", \"\");",
+                "split partition table t between (0) and (10000) regions 4;",
+                "split region for table t index idx between (1000) and (10000) regions 2;",
+                "split partition table t index idx between (1000) and (10000) regions 2;",
+                "split partition table t partition (p1) between (0) and (10000) regions 2;",
+                "split partition table t partition (p2) between (10000) and (20000) regions 2;",
+                "split partition table t partition (p1,p2) index idx between (0) and (20000) regions 2;",
+            }) {
+                SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
+                SQLStatement statement = parser.parseStatement();
+                System.out.println("原始的sql===" + sql);
+                String newSql=statement.toString().replace("\n","").replace('\'','"')+";";
+                System.out.println("生成的sql===" + newSql);
+                assertTrue(newSql.equalsIgnoreCase(sql));
+                SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType);
+                statement.accept(visitor);
+                System.out.println("getTables==" + visitor.getTables());
+                Map<TableStat.Name, TableStat> tableMap = visitor.getTables();
+                assertFalse(tableMap.isEmpty());
+
+            }
+        }
+    }
+}