support odps udtf expr

src/main/java/com/alibaba/druid/sql/dialect/odps/ast/OdpsUDTFSQLSelectItem.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 1999-2011 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.odps.ast;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
+
+public class OdpsUDTFSQLSelectItem extends SQLSelectItem {
+
+    private final List<String> aliasList = new ArrayList<String>();
+
+    public String getAlias() {
+        throw new UnsupportedOperationException();
+    }
+
+    public void setAlias(String alias) {
+        throw new UnsupportedOperationException();
+    }
+
+    public List<String> getAliasList() {
+        return aliasList;
+    }
+
+}

src/main/java/com/alibaba/druid/sql/dialect/odps/parser/OdpsSelectParser.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright 1999-2011 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.odps.parser;
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
+import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
+import com.alibaba.druid.sql.dialect.odps.ast.OdpsUDTFSQLSelectItem;
+import com.alibaba.druid.sql.parser.SQLExprParser;
+import com.alibaba.druid.sql.parser.SQLSelectParser;
+import com.alibaba.druid.sql.parser.Token;
+
+
+public class OdpsSelectParser extends SQLSelectParser {
+    public OdpsSelectParser(SQLExprParser exprParser){
+        super(exprParser.getLexer());
+        this.exprParser = exprParser;
+    }
+    
+    protected SQLSelectItem parseSelectItem() {
+        SQLExpr expr;
+        if (lexer.token() == Token.IDENTIFIER) {
+            expr = new SQLIdentifierExpr(lexer.stringVal());
+            lexer.nextTokenComma();
+
+            if (lexer.token() != Token.COMMA) {
+                expr = this.exprParser.primaryRest(expr);
+                expr = this.exprParser.exprRest(expr);
+            }
+        } else {
+            expr = expr();
+        }
+        
+        if (lexer.token() == Token.AS) {
+            lexer.nextToken();
+            
+            if (lexer.token() == Token.LPAREN) {
+                lexer.nextToken();
+                
+                OdpsUDTFSQLSelectItem selectItem = new OdpsUDTFSQLSelectItem();
+                
+                selectItem.setExpr(expr);
+                
+                for (;;) {
+                    String alias = lexer.stringVal();
+                    lexer.nextToken();
+                    
+                    selectItem.getAliasList().add(alias);
+                    
+                    if (lexer.token() == Token.COMMA) {
+                        lexer.nextToken();
+                        continue;
+                    }
+                    break;
+                }
+                
+                accept(Token.RPAREN);
+                
+                return selectItem;
+            }
+        }
+        
+        final String alias = as();
+
+        return new SQLSelectItem(expr, alias);
+    }
+}

src/main/java/com/alibaba/druid/sql/dialect/odps/visitor/OdpsASTVisitor.java

@@ -18,6 +18,7 @@ package com.alibaba.druid.sql.dialect.odps.visitor;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsCreateTableStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsInsert;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsInsertStatement;
+import com.alibaba.druid.sql.dialect.odps.ast.OdpsUDTFSQLSelectItem;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 public interface OdpsASTVisitor extends SQLASTVisitor {
@@ -33,4 +34,8 @@ public interface OdpsASTVisitor extends SQLASTVisitor {
     void endVisit(OdpsInsert x);
     
     boolean visit(OdpsInsert x);
+    
+    void endVisit(OdpsUDTFSQLSelectItem x);
+    
+    boolean visit(OdpsUDTFSQLSelectItem x);
 }

src/main/java/com/alibaba/druid/sql/dialect/odps/visitor/OdpsOutputVisitor.java

@@ -29,6 +29,7 @@ import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsCreateTableStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsInsert;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsInsertStatement;
+import com.alibaba.druid.sql.dialect.odps.ast.OdpsUDTFSQLSelectItem;
 import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor;
 
 public class OdpsOutputVisitor extends SQLASTOutputVisitor implements OdpsASTVisitor {
@@ -275,4 +276,27 @@ public class OdpsOutputVisitor extends SQLASTOutputVisitor implements OdpsASTVis
 
         return false;
     }
+
+    @Override
+    public void endVisit(OdpsUDTFSQLSelectItem x) {
+        
+    }
+
+    @Override
+    public boolean visit(OdpsUDTFSQLSelectItem x) {
+        x.getExpr().accept(this);
+
+        print(" AS (");
+        
+        for (int i = 0; i < x.getAliasList().size(); ++i) {
+            if (i != 0) {
+                print(", ");
+            }
+            print(x.getAliasList().get(i));
+        }
+        
+        print(")");
+
+        return false;
+    }
 }

src/main/java/com/alibaba/druid/sql/dialect/odps/visitor/OdpsSchemaStatVisitor.java

@@ -19,6 +19,7 @@ import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsCreateTableStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsInsert;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsInsertStatement;
+import com.alibaba.druid.sql.dialect.odps.ast.OdpsUDTFSQLSelectItem;
 import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
 
 
@@ -54,5 +55,15 @@ public class OdpsSchemaStatVisitor extends SchemaStatVisitor implements OdpsASTV
         return false;
     }
 
+    @Override
+    public void endVisit(OdpsUDTFSQLSelectItem x) {
+        
+    }
+
+    @Override
+    public boolean visit(OdpsUDTFSQLSelectItem x) {
+        return true;
+    }
+
     
 }

src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java

@@ -270,22 +270,9 @@ public class SQLSelectParser extends SQLParser {
     protected final void parseSelectList(SQLSelectQueryBlock queryBlock) {
         final List<SQLSelectItem> selectList = queryBlock.getSelectList();
         for (;;) {
-            SQLExpr expr;
-            if (lexer.token() == Token.IDENTIFIER) {
-                expr = new SQLIdentifierExpr(lexer.stringVal());
-                lexer.nextTokenComma();
-
-                if (lexer.token() != Token.COMMA) {
-                    expr = this.exprParser.primaryRest(expr);
-                    expr = this.exprParser.exprRest(expr);
-                }
-            } else {
-                expr = expr();
-            }
-            final String alias = as();
-
-            final SQLSelectItem selectItem = new SQLSelectItem(expr, alias);
+            final SQLSelectItem selectItem = parseSelectItem();
             selectList.add(selectItem);
+            
             if (lexer.token() != Token.COMMA) {
                 break;
             }
@@ -294,6 +281,24 @@ public class SQLSelectParser extends SQLParser {
         }
     }
 
+    protected SQLSelectItem parseSelectItem() {
+        SQLExpr expr;
+        if (lexer.token() == Token.IDENTIFIER) {
+            expr = new SQLIdentifierExpr(lexer.stringVal());
+            lexer.nextTokenComma();
+
+            if (lexer.token() != Token.COMMA) {
+                expr = this.exprParser.primaryRest(expr);
+                expr = this.exprParser.exprRest(expr);
+            }
+        } else {
+            expr = expr();
+        }
+        final String alias = as();
+
+        return new SQLSelectItem(expr, alias);
+    }
+
     public void parseFrom(SQLSelectQueryBlock queryBlock) {
         if (lexer.token() != Token.FROM) {
             return;

src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java

@@ -79,6 +79,7 @@ import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement;
 import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem;
 import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
 import com.alibaba.druid.sql.ast.statement.SQLUseStatement;
+import com.alibaba.druid.sql.dialect.odps.parser.OdpsSelectParser;
 
 public class SQLStatementParser extends SQLParser {
 
@@ -1357,7 +1358,7 @@ public class SQLStatementParser extends SQLParser {
     }
 
     public SQLSelectParser createSQLSelectParser() {
-        return new SQLSelectParser(this.exprParser);
+        return new OdpsSelectParser(this.exprParser);
     }
 
     public SQLUpdateStatement parseUpdateStatement() {

src/test/java/com/alibaba/druid/bvt/sql/odps/OdpsUDTFTest.java

@@ -0,0 +1,42 @@
+package com.alibaba.druid.bvt.sql.odps;
+
+import java.util.List;
+
+import junit.framework.TestCase;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.dialect.odps.parser.OdpsStatementParser;
+import com.alibaba.druid.sql.dialect.odps.visitor.OdpsSchemaStatVisitor;
+import com.alibaba.druid.stat.TableStat;
+import com.alibaba.druid.stat.TableStat.Column;
+
+
+public class OdpsUDTFTest extends TestCase {
+    public void test_0() throws Exception {
+        String sql = "SELECT my_udtf(1,2,3) as (a, b, c) from employee t1 WHERE t1.name = 'aaa';";
+
+        OdpsStatementParser parser = new OdpsStatementParser(sql);
+        List<SQLStatement> statementList = parser.parseStatementList();
+        SQLStatement stmt = statementList.get(0);
+
+        Assert.assertEquals(1, statementList.size());
+
+        OdpsSchemaStatVisitor visitor = new OdpsSchemaStatVisitor();
+        stmt.accept(visitor);
+
+        System.out.println("Tables : " + visitor.getTables());
+        System.out.println("fields : " + visitor.getColumns());
+        System.out.println("coditions : " + visitor.getConditions());
+        System.out.println("orderBy : " + visitor.getOrderByColumns());
+
+        Assert.assertEquals(1, visitor.getTables().size());
+        Assert.assertEquals(1, visitor.getColumns().size());
+        Assert.assertEquals(1, visitor.getConditions().size());
+
+        Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("employee")));
+
+        Assert.assertTrue(visitor.getColumns().contains(new Column("employee", "name")));
+    }
+}