improved postgresql parser. for issue #2129

alibaba · Nov 25, 2017 · b90e175 · b90e175
1 parent 128ddaf
commit b90e175
Show file tree

Hide file tree

Showing 7 changed files with 60 additions and 7 deletions.
diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLCreateTableStatement.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLCreateTableStatement.java
@@ -65,6 +65,8 @@ public class SQLCreateTableStatement extends SQLStatementImpl implements SQLDDLS
     protected final List<SQLSelectOrderByItem> sortedBy = new ArrayList<SQLSelectOrderByItem>();
     protected int                              buckets;
 
+    protected Map<String, SQLObject> tableOptions = new LinkedHashMap<String, SQLObject>();
+
     public SQLCreateTableStatement(){
 
     }
@@ -231,6 +233,10 @@ public void setPartitioning(SQLPartitionBy partitioning) {
         this.partitioning = partitioning;
     }
 
+    public Map<String, SQLObject> getTableOptions() {
+        return tableOptions;
+    }
+
     @Override
     protected void accept0(SQLASTVisitor visitor) {
         if (visitor.visit(this)) {

diff --git a/...ain/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlCreateTableStatement.java b/...ain/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlCreateTableStatement.java
@@ -66,10 +66,6 @@ public void setTableOptions(Map<String, SQLObject> tableOptions) {
         this.tableOptions = tableOptions;
     }
 
-    public Map<String, SQLObject> getTableOptions() {
-        return tableOptions;
-    }
-
     @Deprecated
     public SQLSelect getQuery() {
         return select;

diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java
@@ -46,6 +46,7 @@
 
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 public class PGOutputVisitor extends SQLASTOutputVisitor implements PGASTVisitor, OracleASTVisitor {
@@ -2275,6 +2276,7 @@ public boolean visit(OracleCreateTableStatement x) {
             println();
             x.getSelect().accept(this);
         }
+
         return false;
     }
 

diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLCreateTableParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLCreateTableParser.java
@@ -17,6 +17,7 @@
 
 import java.util.List;
 
+import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.statement.*;
 import com.alibaba.druid.util.FnvHash;
@@ -130,7 +131,7 @@ public SQLCreateTableStatement parseCreateTable(boolean acceptCreate) {
 
             accept(Token.RPAREN);
 
-            if (lexer.identifierEquals("INHERITS")) {
+            if (lexer.identifierEquals(FnvHash.Constants.INHERITS)) {
                 lexer.nextToken();
                 accept(Token.LPAREN);
                 SQLName inherits = this.exprParser.name();
@@ -145,6 +146,29 @@ public SQLCreateTableStatement parseCreateTable(boolean acceptCreate) {
             createTable.setSelect(select);
         }
 
+        if (lexer.token == Token.WITH && JdbcConstants.POSTGRESQL.equals(dbType)) {
+            lexer.nextToken();
+            accept(Token.LPAREN);
+
+            for (;;) {
+                String name = lexer.stringVal();
+                lexer.nextToken();
+                accept(Token.EQ);
+                SQLExpr value = this.exprParser.expr();
+                value.setParent(createTable);
+
+                createTable.getTableOptions().put(name, value);
+
+                if (lexer.token == Token.COMMA) {
+                    lexer.nextToken();
+                    continue;
+                }
+
+                break;
+            }
+            accept(Token.RPAREN);
+        }
+
         return createTable;
     }
 

diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java
@@ -2431,10 +2431,14 @@ protected SQLDataType parseCharTypeRest(SQLCharacterDataType charType) {
         if (lexer.identifierEquals(FnvHash.Constants.COLLATE)) {
             lexer.nextToken();
 
-            if (lexer.token != Token.IDENTIFIER) {
+            if (lexer.token == Token.LITERAL_ALIAS) {
+                charType.setCollate(lexer.stringVal());
+            } else if (lexer.token == Token.IDENTIFIER) {
+                charType.setCollate(lexer.stringVal());
+            } else {
                 throw new ParserException();
             }
-            charType.setCollate(lexer.stringVal());
+
             lexer.nextToken();
         }
 

diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
@@ -2530,6 +2530,26 @@ protected void printTableElements(List<SQLTableElement> tableElementList) {
     public boolean visit(SQLCreateTableStatement x) {
         printCreateTable(x, true);
 
+        Map<String, SQLObject> options = x.getTableOptions();
+        if (options.size() > 0) {
+            println();
+            print0(ucase ? "WITH (" : "with (");
+            int i = 0;
+            for (Map.Entry<String, SQLObject> option : x.getTableOptions().entrySet()) {
+                if (i > 0) {
+                    print0(", ");
+                }
+                String key = option.getKey();
+                print0(key);
+
+                print0(" = ");
+
+                option.getValue().accept(this);
+                ++i;
+            }
+            print(')');
+        }
+
         return false;
     }
 

diff --git a/src/main/java/com/alibaba/druid/util/FnvHash.java b/src/main/java/com/alibaba/druid/util/FnvHash.java
@@ -620,5 +620,6 @@ public static interface Constants {
         long SAVEPOINT = fnv1a_64_lower("SAVEPOINT");
         long RELEASE = fnv1a_64_lower("RELEASE");
         long MERGE = fnv1a_64_lower("MERGE");
+        long INHERITS = fnv1a_64_lower("INHERITS");
     }
 }