Fix Presto Parser

alibaba · Aug 23, 2024 · 5fc0140 · 5fc0140
1 parent 56b82fb
commit 5fc0140
Show file tree

Hide file tree

Showing 24 changed files with 284 additions and 80 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/SQLUtils.java b/core/src/main/java/com/alibaba/druid/sql/SQLUtils.java
@@ -59,7 +59,7 @@
 import com.alibaba.druid.sql.dialect.oscar.visitor.OscarOutputVisitor;
 import com.alibaba.druid.sql.dialect.postgresql.visitor.PGOutputVisitor;
 import com.alibaba.druid.sql.dialect.postgresql.visitor.PGSchemaStatVisitor;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoOutputVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoOutputASTVisitor;
 import com.alibaba.druid.sql.dialect.redshift.visitor.RedshiftOutputVisitor;
 import com.alibaba.druid.sql.dialect.spark.visitor.SparkOutputVisitor;
 import com.alibaba.druid.sql.dialect.spark.visitor.SparkSchemaStatVisitor;
@@ -551,7 +551,7 @@ public static SQLASTOutputVisitor createFormatOutputVisitor(
                 return new SparkOutputVisitor(out);
             case presto:
             case trino:
-                return new PrestoOutputVisitor(out);
+                return new PrestoOutputASTVisitor(out);
             case clickhouse:
                 return new CKOutputVisitor(out);
             case oscar:

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/PrestoColumnConstraint.java b/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/PrestoColumnConstraint.java
@@ -0,0 +1,26 @@
+package com.alibaba.druid.sql.dialect.presto.ast;
+
+import com.alibaba.druid.DbType;
+import com.alibaba.druid.sql.ast.statement.SQLColumnConstraint;
+import com.alibaba.druid.sql.ast.statement.SQLConstraintImpl;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+public class PrestoColumnConstraint extends SQLConstraintImpl implements SQLColumnConstraint {
+    public PrestoColumnConstraint() {
+        dbType = DbType.presto;
+    }
+    @Override
+    protected void accept0(SQLASTVisitor v) {}
+
+    @Override
+    public PrestoColumnConstraint clone() {
+        PrestoColumnConstraint prestoColumnConstraint = new PrestoColumnConstraint();
+        cloneTo(prestoColumnConstraint);
+        return prestoColumnConstraint;
+    }
+
+    public void cloneTo(PrestoColumnConstraint x) {
+        super.cloneTo(x);
+        x.dbType = dbType;
+    }
+}
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/PrestoColumnWith.java b/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/PrestoColumnWith.java
@@ -0,0 +1,41 @@
+package com.alibaba.druid.sql.dialect.presto.ast;
+
+import com.alibaba.druid.sql.ast.statement.SQLAssignItem;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class PrestoColumnWith extends PrestoColumnConstraint implements PrestoObject {
+    private final List<SQLAssignItem> properties;
+
+    public PrestoColumnWith() {
+        properties = new ArrayList<>();
+    }
+
+    public void addProperty(SQLAssignItem sqlAssignItem) {
+        if (sqlAssignItem != null) {
+            sqlAssignItem.setParent(this);
+            properties.add(sqlAssignItem);
+        }
+    }
+
+    public List<SQLAssignItem> getProperties() {
+        return properties;
+    }
+
+    @Override
+    public void accept0(PrestoASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            acceptChild(visitor, properties);
+        }
+    }
+
+    @Override
+    public void accept0(SQLASTVisitor visitor) {
+        if (visitor instanceof PrestoASTVisitor) {
+            accept0((PrestoASTVisitor) visitor);
+        }
+    }
+}
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/PrestoObject.java b/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/PrestoObject.java
@@ -1,8 +1,14 @@
 package com.alibaba.druid.sql.dialect.presto.ast;
 
 import com.alibaba.druid.sql.ast.SQLObject;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 public interface PrestoObject extends SQLObject {
-    void accept0(PrestoVisitor visitor);
+    void accept0(PrestoASTVisitor visitor);
+    default void accept0(SQLASTVisitor visitor) {
+        if (visitor instanceof PrestoASTVisitor) {
+            accept0((PrestoASTVisitor) visitor);
+        }
+    }
 }
diff --git a/...main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoAlterFunctionStatement.java b/...main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoAlterFunctionStatement.java
@@ -16,7 +16,7 @@
 package com.alibaba.druid.sql.dialect.presto.ast.stmt;
 
 import com.alibaba.druid.sql.ast.statement.SQLAlterFunctionStatement;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 public class PrestoAlterFunctionStatement extends SQLAlterFunctionStatement implements PrestoSQLStatement {
@@ -43,16 +43,16 @@ public void setCalledOnNullInput(boolean calledOnNullInput) {
     }
 
     @Override
-    protected void accept0(SQLASTVisitor v) {
-        if (v instanceof PrestoVisitor) {
-            this.accept0((PrestoVisitor) v);
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof PrestoASTVisitor) {
+            this.accept0((PrestoASTVisitor) v);
         } else {
             super.accept0(v);
         }
     }
 
     @Override
-    public void accept0(PrestoVisitor visitor) {
+    public void accept0(PrestoASTVisitor visitor) {
         visitor.visit(this);
     }
 }
diff --git a/...c/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoAlterSchemaStatement.java b/...c/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoAlterSchemaStatement.java
@@ -19,7 +19,7 @@
 import com.alibaba.druid.sql.ast.SQLStatementImpl;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
 import com.alibaba.druid.sql.ast.statement.SQLAlterStatement;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 public class PrestoAlterSchemaStatement extends SQLStatementImpl implements PrestoSQLStatement, SQLAlterStatement {
@@ -28,16 +28,16 @@ public class PrestoAlterSchemaStatement extends SQLStatementImpl implements Pres
     private SQLIdentifierExpr newName;
 
     @Override
-    protected void accept0(SQLASTVisitor v) {
-        if (v instanceof PrestoVisitor) {
-            this.accept0((PrestoVisitor) v);
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof PrestoASTVisitor) {
+            this.accept0((PrestoASTVisitor) v);
         } else {
             super.accept0(v);
         }
     }
 
     @Override
-    public void accept0(PrestoVisitor visitor) {
+    public void accept0(PrestoASTVisitor visitor) {
         visitor.visit(this);
     }
 

diff --git a/...c/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoCreateTableStatement.java b/...c/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoCreateTableStatement.java
@@ -17,7 +17,7 @@
 
 import com.alibaba.druid.DbType;
 import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 public class PrestoCreateTableStatement extends SQLCreateTableStatement implements PrestoSQLStatement {
@@ -30,16 +30,16 @@ public PrestoCreateTableStatement(DbType dbType) {
     }
 
     @Override
-    protected void accept0(SQLASTVisitor v) {
-        if (v instanceof PrestoVisitor) {
-            this.accept0((PrestoVisitor) v);
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof PrestoASTVisitor) {
+            this.accept0((PrestoASTVisitor) v);
         }
         super.accept0(v);
     }
 
     @Override
-    public void accept0(PrestoVisitor visitor) {
-        visitor.visit(this);
+    public void accept0(PrestoASTVisitor visitor) {
+        visitor.visit();
     }
 
     protected void acceptChild(SQLASTVisitor v) {

diff --git a/.../java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoDeallocatePrepareStatement.java b/.../java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoDeallocatePrepareStatement.java
@@ -18,7 +18,7 @@
 import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLObject;
 import com.alibaba.druid.sql.ast.SQLStatementImpl;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 import java.util.Collections;
@@ -36,16 +36,16 @@ public void setStatementName(SQLName statementName) {
     }
 
     @Override
-    protected void accept0(SQLASTVisitor v) {
-        if (v instanceof PrestoVisitor) {
-            this.accept0((PrestoVisitor) v);
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof PrestoASTVisitor) {
+            this.accept0((PrestoASTVisitor) v);
         } else {
             super.accept0(v);
         }
     }
 
     @Override
-    public void accept0(PrestoVisitor visitor) {
+    public void accept0(PrestoASTVisitor visitor) {
         if (visitor.visit(this)) {
             acceptChild(visitor, statementName);
         }

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoExecuteStatement.java b/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoExecuteStatement.java
@@ -20,7 +20,7 @@
 import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLObject;
 import com.alibaba.druid.sql.ast.SQLStatementImpl;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 import java.util.ArrayList;
@@ -47,16 +47,16 @@ public List<SQLExpr> getParameters() {
     }
 
     @Override
-    protected void accept0(SQLASTVisitor v) {
-        if (v instanceof PrestoVisitor) {
-            this.accept0((PrestoVisitor) v);
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof PrestoASTVisitor) {
+            this.accept0((PrestoASTVisitor) v);
         } else {
             super.accept0(v);
         }
     }
 
     @Override
-    public void accept0(PrestoVisitor visitor) {
+    public void accept0(PrestoASTVisitor visitor) {
         if (visitor.visit(this)) {
             acceptChild(visitor, statementName);
             acceptChild(visitor, parameters);

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoPrepareStatement.java b/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoPrepareStatement.java
@@ -21,7 +21,7 @@
 import com.alibaba.druid.sql.ast.SQLStatementImpl;
 import com.alibaba.druid.sql.ast.statement.SQLSelect;
 import com.alibaba.druid.sql.dialect.hive.ast.HiveInsertStatement;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 import java.util.ArrayList;
@@ -65,16 +65,16 @@ public void setInsert(HiveInsertStatement insert) {
     }
 
     @Override
-    protected void accept0(SQLASTVisitor v) {
-        if (v instanceof PrestoVisitor) {
-            this.accept0((PrestoVisitor) v);
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof PrestoASTVisitor) {
+            this.accept0((PrestoASTVisitor) v);
         } else {
             super.accept0(v);
         }
     }
 
     @Override
-    public void accept0(PrestoVisitor visitor) {
+    public void accept0(PrestoASTVisitor visitor) {
         if (visitor.visit(this)) {
             acceptChild(visitor, name);
             if (this.select != null) {

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoSelectStatement.java b/core/src/main/java/com/alibaba/druid/sql/dialect/presto/ast/stmt/PrestoSelectStatement.java
@@ -18,7 +18,7 @@
 import com.alibaba.druid.DbType;
 import com.alibaba.druid.sql.ast.statement.SQLSelect;
 import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
-import com.alibaba.druid.sql.dialect.presto.visitor.PrestoVisitor;
+import com.alibaba.druid.sql.dialect.presto.visitor.PrestoASTVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 /**
@@ -37,15 +37,15 @@ public PrestoSelectStatement(SQLSelect select) {
     }
 
     @Override
-    protected void accept0(SQLASTVisitor visitor) {
-        if (visitor instanceof PrestoVisitor) {
-            this.accept0((PrestoVisitor) visitor);
+    public void accept0(SQLASTVisitor visitor) {
+        if (visitor instanceof PrestoASTVisitor) {
+            this.accept0((PrestoASTVisitor) visitor);
         } else {
             super.accept0(visitor);
         }
     }
 
-    public void accept0(PrestoVisitor visitor) {
+    public void accept0(PrestoASTVisitor visitor) {
         super.accept0(visitor);
     }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/presto/parser/PrestoCreateTableParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/presto/parser/PrestoCreateTableParser.java
@@ -17,6 +17,7 @@
 
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.statement.*;
+import com.alibaba.druid.sql.dialect.presto.ast.PrestoColumnWith;
 import com.alibaba.druid.sql.dialect.presto.ast.stmt.PrestoCreateTableStatement;
 import com.alibaba.druid.sql.parser.SQLCreateTableParser;
 import com.alibaba.druid.sql.parser.SQLExprParser;
@@ -76,7 +77,27 @@ public SQLCreateTableStatement parseCreateTable() {
 
             for (; ; ) {
                 Token token = lexer.token();
-                if (lexer.token() == Token.LIKE) {
+                if (token == Token.IDENTIFIER
+                        || token == Token.LITERAL_ALIAS) {
+                    SQLColumnDefinition column = this.exprParser.parseColumn(createTable);
+                    column.setParent(createTable);
+                    createTable.getTableElementList().add(column);
+
+                    if (lexer.token() == Token.COMMENT) {
+                        lexer.nextToken();
+                        SQLExpr comment = this.exprParser.expr();
+                        column.setComment(comment);
+                    }
+
+                    if (lexer.token() == Token.WITH) {
+                        lexer.nextToken();
+                        PrestoColumnWith prestoColumnWith = new PrestoColumnWith();
+                        accept(Token.LPAREN);
+                        parseAssignItems(prestoColumnWith.getProperties(), prestoColumnWith, false);
+                        accept(Token.RPAREN);
+                        column.addConstraint(prestoColumnWith);
+                    }
+                } else if (lexer.token() == Token.LIKE) {
                     lexer.nextToken();
                     SQLTableLike tableLike = new SQLTableLike();
                     tableLike.setTable(new SQLExprTableSource(this.exprParser.name()));

diff --git a/...dialect/presto/visitor/PrestoVisitor.java → ...lect/presto/visitor/PrestoASTVisitor.java b/...dialect/presto/visitor/PrestoVisitor.java → ...lect/presto/visitor/PrestoASTVisitor.java
@@ -1,5 +1,6 @@
 package com.alibaba.druid.sql.dialect.presto.visitor;
 
+import com.alibaba.druid.sql.dialect.presto.ast.PrestoColumnWith;
 import com.alibaba.druid.sql.dialect.presto.ast.stmt.PrestoAlterFunctionStatement;
 import com.alibaba.druid.sql.dialect.presto.ast.stmt.PrestoAlterSchemaStatement;
 import com.alibaba.druid.sql.dialect.presto.ast.stmt.PrestoCreateTableStatement;
@@ -8,8 +9,8 @@
 import com.alibaba.druid.sql.dialect.presto.ast.stmt.PrestoPrepareStatement;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
-public interface PrestoVisitor extends SQLASTVisitor {
-    default boolean visit(PrestoCreateTableStatement x) {
+public interface PrestoASTVisitor extends SQLASTVisitor {
+    default boolean visit() {
         return true;
     }
 
@@ -50,4 +51,10 @@ default boolean visit(PrestoDeallocatePrepareStatement x) {
 
     default void endVisit(PrestoDeallocatePrepareStatement x) {
     }
+
+    default boolean visit(PrestoColumnWith x) { return true; }
+
+    default void endVisit(PrestoColumnWith x) {
+    }
+
 }