增强解析支持PostgreSQL的PARTITION OF语法 #5366

core/src/main/java/com/alibaba/druid/sql/ast/SQLPartitionOf.java

@@ -0,0 +1,165 @@
+/*
+ * 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.ast;
+
+
+import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+import java.util.List;
+
+/**
+ * @author lizongbo
+ */
+public class SQLPartitionOf extends SQLObjectImpl {
+
+    protected SQLExprTableSource parentTable;
+    private boolean useDefault;
+    private SQLName columnName;
+    private SQLName constraintName;
+    private SQLExpr checkExpr;
+    private SQLExpr defaultExpr;
+    private List<SQLExpr> forValuesFrom;
+    private List<SQLExpr> forValuesTo;
+    private List<SQLExpr> forValuesIn;
+
+    private SQLExpr forValuesModulus;
+
+    private SQLExpr forValuesRemainder;
+
+    public SQLExprTableSource getParentTable() {
+        return parentTable;
+    }
+
+    public void setParentTable(SQLExprTableSource parentTable) {
+        this.parentTable = parentTable;
+    }
+
+    public boolean isUseDefault() {
+        return useDefault;
+    }
+
+    public void setUseDefault(boolean useDefault) {
+        this.useDefault = useDefault;
+    }
+
+    public List<SQLExpr> getForValuesFrom() {
+        return forValuesFrom;
+    }
+
+    public void setForValuesFrom(List<SQLExpr> forValuesFrom) {
+        this.forValuesFrom = forValuesFrom;
+    }
+
+    public List<SQLExpr> getForValuesTo() {
+        return forValuesTo;
+    }
+
+    public void setForValuesTo(List<SQLExpr> forValuesTo) {
+        this.forValuesTo = forValuesTo;
+    }
+
+    public List<SQLExpr> getForValuesIn() {
+        return forValuesIn;
+    }
+
+    public void setForValuesIn(List<SQLExpr> forValuesIn) {
+        this.forValuesIn = forValuesIn;
+    }
+
+    public SQLName getColumnName() {
+        return columnName;
+    }
+
+    public void setColumnName(SQLName columnName) {
+        this.columnName = columnName;
+    }
+
+    public SQLName getConstraintName() {
+        return constraintName;
+    }
+
+    public void setConstraintName(SQLName constraintName) {
+        this.constraintName = constraintName;
+    }
+
+    public SQLExpr getCheckExpr() {
+        return checkExpr;
+    }
+
+    public void setCheckExpr(SQLExpr checkExpr) {
+        this.checkExpr = checkExpr;
+    }
+
+    public SQLExpr getDefaultExpr() {
+        return defaultExpr;
+    }
+
+    public void setDefaultExpr(SQLExpr defaultExpr) {
+        this.defaultExpr = defaultExpr;
+    }
+
+    public SQLExpr getForValuesModulus() {
+        return forValuesModulus;
+    }
+
+    public void setForValuesModulus(SQLExpr forValuesModulus) {
+        this.forValuesModulus = forValuesModulus;
+    }
+
+    public SQLExpr getForValuesRemainder() {
+        return forValuesRemainder;
+    }
+
+    public void setForValuesRemainder(SQLExpr forValuesRemainder) {
+        this.forValuesRemainder = forValuesRemainder;
+    }
+
+    @Override
+    protected void accept0(SQLASTVisitor v) {
+        if (v.visit(this)) {
+            acceptChild(v, parentTable);
+            if (columnName != null) {
+                acceptChild(v, columnName);
+            }
+            if (constraintName != null) {
+                acceptChild(v, constraintName);
+            }
+            if (checkExpr != null) {
+                acceptChild(v, checkExpr);
+            }
+            if (defaultExpr != null) {
+                acceptChild(v, defaultExpr);
+            }
+            if (forValuesFrom != null) {
+                acceptChild(v, forValuesFrom);
+            }
+            if (forValuesTo != null) {
+                acceptChild(v, forValuesTo);
+            }
+            if (forValuesIn != null) {
+                acceptChild(v, forValuesIn);
+            }
+            if (forValuesModulus != null) {
+                acceptChild(v, forValuesModulus);
+            }
+            if (forValuesRemainder != null) {
+                acceptChild(v, forValuesRemainder);
+            }
+        }
+        v.endVisit(this);
+    }
+}

core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLCreateTableStatement.java

@@ -50,6 +50,7 @@ public class SQLCreateTableStatement extends SQLStatementImpl implements SQLDDLS
 
     protected SQLName tablespace;
     protected SQLPartitionBy partitioning;
+    protected SQLPartitionOf partitionOf;
     protected SQLPartitionBy localPartitioning;
     protected SQLExpr storedAs;
     protected SQLExpr location;
@@ -362,6 +363,16 @@ public void setPartitioning(SQLPartitionBy partitioning) {
         this.partitioning = partitioning;
     }
 
+    public SQLPartitionOf getPartitionOf() {
+        return partitionOf;
+    }
+
+    public void setPartitionOf(SQLPartitionOf partitionOf) {
+        if (partitionOf != null) {
+            partitionOf.setParent(this);
+        }
+        this.partitionOf = partitionOf;
+    }
     public void setLocalPartitioning(SQLPartitionBy localPartitioning) {
         if (localPartitioning != null) {
             localPartitioning.setParent(this);

core/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGCreateTableParser.java

@@ -1,9 +1,22 @@
 package com.alibaba.druid.sql.dialect.postgresql.parser;
 
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLPartitionBy;
 import com.alibaba.druid.sql.ast.SQLPartitionByHash;
 import com.alibaba.druid.sql.ast.SQLPartitionByList;
+import com.alibaba.druid.sql.ast.SQLPartitionByRange;
+import com.alibaba.druid.sql.ast.SQLPartitionOf;
+import com.alibaba.druid.sql.ast.expr.SQLBetweenExpr;
+import com.alibaba.druid.sql.ast.expr.SQLIntegerExpr;
+import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.parser.*;
+import com.alibaba.druid.util.FnvHash;
+
+import java.util.ArrayList;
+import java.util.List;
+
 
 public class PGCreateTableParser extends SQLCreateTableParser {
     public PGCreateTableParser(Lexer lexer) {
@@ -63,8 +76,121 @@ public SQLPartitionBy parsePartitionBy() {
             this.exprParser.exprList(hash.getColumns(), hash);
             accept(Token.RPAREN);
             return hash;
+        } else if (lexer.token() == Token.IDENTIFIER) {
+            SQLPartitionByRange clause = partitionByRange();
+            return clause;
         }
 
         throw new ParserException("TODO " + lexer.info());
     }
+    protected SQLPartitionByRange partitionByRange() {
+        SQLPartitionByRange clause = new SQLPartitionByRange();
+        if (lexer.identifierEquals(FnvHash.Constants.RANGE)) {
+            lexer.nextToken();
+
+            if (lexer.token() == Token.LPAREN) {
+                lexer.nextToken();
+                clause.addColumn(this.exprParser.expr());
+                accept(Token.RPAREN);
+            } else {
+                acceptIdentifier("COLUMNS");
+                accept(Token.LPAREN);
+                for (; ; ) {
+                    clause.addColumn(this.exprParser.name());
+                    if (lexer.token() == Token.COMMA) {
+                        lexer.nextToken();
+                        continue;
+                    }
+                    break;
+                }
+                accept(Token.RPAREN);
+            }
+        } else {
+            SQLExpr expr = this.exprParser.expr();
+            if (lexer.identifierEquals(FnvHash.Constants.STARTWITH)) {
+                lexer.nextToken();
+                SQLExpr start = this.exprParser.primary();
+                acceptIdentifier("ENDWITH");
+                SQLExpr end = this.exprParser.primary();
+                expr = new SQLBetweenExpr(expr, start, end);
+            }
+            clause.setInterval(expr);
+        }
+
+        return clause;
+    }
+
+    public SQLPartitionOf parsePartitionOf() {
+        lexer.nextToken();
+        accept(Token.OF);
+        SQLPartitionOf partitionOf = new SQLPartitionOf();
+        SQLName tableNameTmp = this.exprParser.name();
+        SQLExprTableSource sqlExprTableSource = new SQLExprTableSource(tableNameTmp);
+        partitionOf.setParentTable(sqlExprTableSource);
+        if (lexer.token() == Token.LPAREN) {
+            lexer.nextToken();
+            if (lexer.token() == Token.CONSTRAINT) {
+                lexer.nextToken();
+                SQLName constraintName = this.exprParser.name();
+                partitionOf.setConstraintName(constraintName);
+                accept(Token.CHECK);
+                SQLExpr checkExpr = this.exprParser.expr();
+                partitionOf.setCheckExpr(checkExpr);
+            } else {
+                SQLName columnName = this.exprParser.name();
+                partitionOf.setColumnName(columnName);
+                if (lexer.token() == Token.DEFAULT) {
+                    accept(Token.DEFAULT);
+                    SQLExpr defaultExpr = this.exprParser.primary();
+                    partitionOf.setDefaultExpr(defaultExpr);
+                }
+            }
+            accept(Token.RPAREN);
+        }
+        if (lexer.token() == Token.DEFAULT) {
+            accept(Token.DEFAULT);
+            partitionOf.setUseDefault(true);
+            return partitionOf;
+        }
+        accept(Token.FOR);
+        accept(Token.VALUES);
+        if (lexer.token() == Token.FROM) {
+            accept(Token.FROM);
+            accept(Token.LPAREN);
+            List<SQLExpr> sqlExprBetweens = new ArrayList<>();
+            this.exprParser.exprList(sqlExprBetweens, partitionOf);
+            partitionOf.setForValuesFrom(sqlExprBetweens);
+            accept(Token.RPAREN);
+            accept(Token.TO);
+            accept(Token.LPAREN);
+            List<SQLExpr> sqlExprAnds = new ArrayList<>();
+            this.exprParser.exprList(sqlExprAnds, partitionOf);
+            partitionOf.setForValuesTo(sqlExprAnds);
+            accept(Token.RPAREN);
+            return partitionOf;
+        } else if (lexer.token() == Token.IN) {
+            accept(Token.IN);
+            accept(Token.LPAREN);
+            List<SQLExpr> sqlExprBetweens = new ArrayList<>();
+            this.exprParser.exprList(sqlExprBetweens, partitionOf);
+            partitionOf.setForValuesIn(sqlExprBetweens);
+            accept(Token.RPAREN);
+            return partitionOf;
+        }
+        if (lexer.token() == Token.WITH) {
+            accept(Token.WITH);
+            accept(Token.LPAREN);
+            acceptIdentifier("MODULUS");
+            SQLExpr modulus = this.exprParser.primary();
+            partitionOf.setForValuesModulus(modulus);
+            accept(Token.COMMA);
+            acceptIdentifier("REMAINDER");
+            SQLExpr remainder = (SQLIntegerExpr) this.exprParser.primary();
+            partitionOf.setForValuesRemainder(remainder);
+            accept(Token.RPAREN);
+            return partitionOf;
+        } else {
+            throw new ParserException("TODO " + lexer.info());
+        }
+    }
 }

core/src/main/java/com/alibaba/druid/sql/parser/SQLCreateTableParser.java

@@ -18,6 +18,7 @@
 import com.alibaba.druid.DbType;
 import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLPartitionBy;
+import com.alibaba.druid.sql.ast.SQLPartitionOf;
 import com.alibaba.druid.sql.ast.statement.*;
 import com.alibaba.druid.sql.dialect.oracle.parser.OracleSelectParser;
 import com.alibaba.druid.util.FnvHash;
@@ -176,8 +177,31 @@ public SQLCreateTableStatement parseCreateTable(boolean acceptCreate) {
         }
 
         if (lexer.token() == Token.PARTITION) {
-            SQLPartitionBy partitionClause = parsePartitionBy();
-            createTable.setPartitioning(partitionClause);
+            Lexer.SavePoint mark = lexer.mark();
+            lexer.nextToken();
+            if (Token.OF.equals(lexer.token())) {
+                lexer.reset(mark);
+                SQLPartitionOf partitionOf = parsePartitionOf();
+                createTable.setPartitionOf(partitionOf);
+            } else if (Token.BY.equals(lexer.token())) {
+                lexer.reset(mark);
+                SQLPartitionBy partitionClause = parsePartitionBy();
+                createTable.setPartitioning(partitionClause);
+            }
+        }
+
+        if (lexer.token() == Token.PARTITION) {
+            Lexer.SavePoint mark = lexer.mark();
+            lexer.nextToken();
+            if (Token.OF.equals(lexer.token())) {
+                lexer.reset(mark);
+                SQLPartitionOf partitionOf = parsePartitionOf();
+                createTable.setPartitionOf(partitionOf);
+            } else if (Token.BY.equals(lexer.token())) {
+                lexer.reset(mark);
+                SQLPartitionBy partitionClause = parsePartitionBy();
+                createTable.setPartitioning(partitionClause);
+            }
         }
 
         parseCreateTableRest(createTable);
@@ -192,6 +216,10 @@ public SQLPartitionBy parsePartitionBy() {
         return null;
     }
 
+    public SQLPartitionOf parsePartitionOf() {
+        return null;
+    }
+
     protected SQLTableElement parseCreateTableSupplementalLogingProps() {
         throw new ParserException("TODO " + lexer.info());
     }