Merge pull request #718 from wenshao/master

improved sql parser for odps

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

@@ -0,0 +1,59 @@
+package com.alibaba.druid.sql.dialect.odps.ast;
+
+import java.util.ArrayList;
+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.SQLTableSource;
+import com.alibaba.druid.sql.dialect.odps.visitor.OdpsASTVisitor;
+
+public class OdpsSetLabelStatement extends OdpsStatementImpl {
+
+    private String         label;
+
+    private SQLExpr        user;
+
+    private SQLTableSource table;
+
+    private List<SQLName>  columns = new ArrayList<SQLName>();
+
+    @Override
+    protected void accept0(OdpsASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            acceptChild(visitor, user);
+        }
+        visitor.endVisit(this);
+    }
+
+    public String getLabel() {
+        return label;
+    }
+
+    public void setLabel(String label) {
+        this.label = label;
+    }
+
+    public SQLExpr getUser() {
+        return user;
+    }
+
+    public void setUser(SQLExpr user) {
+        this.user = user;
+        user.setParent(this);
+    }
+
+    public SQLTableSource getTable() {
+        return table;
+    }
+
+    public void setTable(SQLTableSource table) {
+        this.table = table;
+        table.setParent(this);
+    }
+
+    public List<SQLName> getColumns() {
+        return columns;
+    }
+
+}

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

@@ -0,0 +1,20 @@
+package com.alibaba.druid.sql.dialect.odps.ast;
+
+import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatementImpl;
+import com.alibaba.druid.sql.dialect.odps.visitor.OdpsASTVisitor;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+
+public abstract class OdpsStatementImpl extends SQLStatementImpl {
+    @Override
+    protected void accept0(SQLASTVisitor visitor) {
+        accept0((OdpsASTVisitor) visitor);
+    }
+    
+    protected abstract void accept0(OdpsASTVisitor visitor);
+    
+    public String toString() {
+        return SQLUtils.toOdpsString(this);
+    }
+}

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

@@ -15,6 +15,8 @@
  */
 package com.alibaba.druid.sql.dialect.odps.parser;
 
+import static com.alibaba.druid.sql.parser.CharTypes.isFirstIdentifierChar;
+import static com.alibaba.druid.sql.parser.CharTypes.isIdentifierChar;
 import static com.alibaba.druid.sql.parser.LayoutCharacters.EOI;
 
 import java.util.HashMap;
@@ -23,6 +25,7 @@ import java.util.Map;
 import com.alibaba.druid.sql.parser.Keywords;
 import com.alibaba.druid.sql.parser.Lexer;
 import com.alibaba.druid.sql.parser.NotAllowCommentException;
+import com.alibaba.druid.sql.parser.ParserException;
 import com.alibaba.druid.sql.parser.Token;
 
 public class OdpsLexer extends Lexer {
@@ -142,4 +145,52 @@ public class OdpsLexer extends Lexer {
         }
     }
 
+    public void scanIdentifier() {
+        final char first = ch;
+
+        final boolean firstFlag = isFirstIdentifierChar(first);
+        if (!firstFlag) {
+            throw new ParserException("illegal identifier");
+        }
+
+        mark = pos;
+        bufPos = 1;
+        char ch;
+        for (;;) {
+            ch = charAt(++pos);
+
+            if (!isIdentifierChar(ch)) {
+                break;
+            }
+
+            bufPos++;
+            continue;
+        }
+
+        this.ch = charAt(pos);
+        
+        if (ch == '@') { // for user identifier, like email, xx@alibaba-inc.com
+            bufPos++;
+            for (;;) {
+                ch = charAt(++pos);
+
+                if (ch != '-' && ch != '.' && !isIdentifierChar(ch)) {
+                    break;
+                }
+
+                bufPos++;
+                continue;
+            }
+        }
+        this.ch = charAt(pos);
+
+        stringVal = addSymbol();
+        Token tok = keywods.getKeyword(stringVal);
+        if (tok != null) {
+            token = tok;
+        } else {
+            token = Token.IDENTIFIER;
+        }
+    }
+
 }

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

@@ -18,14 +18,17 @@ package com.alibaba.druid.sql.dialect.odps.parser;
 import java.util.List;
 
 import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.ast.statement.SQLAssignItem;
 import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement;
 import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLSelect;
+import com.alibaba.druid.sql.ast.statement.SQLSetStatement;
 import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource;
 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.OdpsSetLabelStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsShowPartitionsStmt;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsShowStatisticStmt;
 import com.alibaba.druid.sql.parser.ParserException;
@@ -69,7 +72,7 @@ public class OdpsStatementParser extends SQLStatementParser {
         if (lexer.token() == Token.FROM) {
             lexer.nextToken();
             accept(Token.LPAREN);
-            
+
             SQLSelectParser selectParser = createSQLSelectParser();
             SQLSelect select = selectParser.select();
 
@@ -82,11 +85,11 @@ public class OdpsStatementParser extends SQLStatementParser {
 
             stmt.setFrom(from);
         }
-        
+
         for (;;) {
             OdpsInsert insert = parseOdpsInsert();
             stmt.getItems().add(insert);
-            
+
             if (lexer.token() != Token.INSERT) {
                 break;
             }
@@ -140,29 +143,65 @@ public class OdpsStatementParser extends SQLStatementParser {
 
     public SQLStatement parseShow() {
         accept(Token.SHOW);
-        
+
         if (identifierEquals("PARTITIONS")) {
             lexer.nextToken();
-            
+
             OdpsShowPartitionsStmt stmt = new OdpsShowPartitionsStmt();
-            
+
             SQLExpr expr = this.exprParser.expr();
             stmt.setTableSource(new SQLExprTableSource(expr));
-            
+
             return stmt;
         }
-        
+
         if (identifierEquals("STATISTIC")) {
             lexer.nextToken();
-            
+
             OdpsShowStatisticStmt stmt = new OdpsShowStatisticStmt();
-            
+
             SQLExpr expr = this.exprParser.expr();
             stmt.setTableSource(new SQLExprTableSource(expr));
-            
+
             return stmt;
         }
         throw new ParserException("TODO " + lexer.token() + " " + lexer.stringVal());
     }
 
+    public SQLStatement parseSet() {
+        accept(Token.SET);
+
+        if (identifierEquals("LABEL")) {
+            lexer.nextToken();
+            OdpsSetLabelStatement stmt = new OdpsSetLabelStatement();
+            stmt.setLabel(lexer.stringVal());
+            lexer.nextToken();
+            accept(Token.TO);
+            if (lexer.token() == Token.USER) {
+                lexer.nextToken();
+                
+                SQLName name = this.exprParser.name();
+                stmt.setUser(name);
+                return stmt;
+            }
+            accept(Token.TABLE);
+            SQLExpr expr = this.exprParser.name();
+            stmt.setTable(new SQLExprTableSource(expr));
+            
+            if (lexer.token() == Token.LPAREN) {
+                lexer.nextToken();
+                this.exprParser.names(stmt.getColumns(), stmt);
+                accept(Token.RPAREN);
+            }
+            
+            return stmt;
+        } else {
+            SQLSetStatement stmt = new SQLSetStatement(getDbType());
+
+            parseAssignItems(stmt.getItems(), stmt);
+
+            return stmt;
+        }
+    }
+
 }

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.OdpsSetLabelStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsShowPartitionsStmt;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsShowStatisticStmt;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsUDTFSQLSelectItem;
@@ -48,4 +49,8 @@ public interface OdpsASTVisitor extends SQLASTVisitor {
     void endVisit(OdpsShowStatisticStmt x);
     
     boolean visit(OdpsShowStatisticStmt x);
+    
+    void endVisit(OdpsSetLabelStatement x);
+    
+    boolean visit(OdpsSetLabelStatement 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.OdpsSetLabelStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsShowPartitionsStmt;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsShowStatisticStmt;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsUDTFSQLSelectItem;
@@ -331,4 +332,31 @@ public class OdpsOutputVisitor extends SQLASTOutputVisitor implements OdpsASTVis
         x.getTableSource().accept(this);
         return false;
     }
+
+    @Override
+    public void endVisit(OdpsSetLabelStatement x) {
+        
+    }
+
+    @Override
+    public boolean visit(OdpsSetLabelStatement x) {
+        print("SET LABEL ");
+        print(x.getLabel());
+        print(" TO ");
+        
+        if (x.getUser() != null) {
+            print("USER ");
+            x.getUser().accept(this);
+        } else if (x.getTable() != null) {
+            print("TABLE ");
+            x.getTable().accept(this);
+            if (x.getColumns().size() > 0) {
+                print("(");
+                printAndAccept(x.getColumns(), ", ");
+                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.OdpsSetLabelStatement;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsShowPartitionsStmt;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsShowStatisticStmt;
 import com.alibaba.druid.sql.dialect.odps.ast.OdpsUDTFSQLSelectItem;
@@ -86,6 +87,19 @@ public class OdpsSchemaStatVisitor extends SchemaStatVisitor implements OdpsASTV
     public boolean visit(OdpsShowStatisticStmt x) {
         return true;
     }
+    
+    @Override
+    public void endVisit(OdpsSetLabelStatement x) {
+        
+    }
+    
+    @Override
+    public boolean visit(OdpsSetLabelStatement x) {
+        if (x.getTable() != null) {
+            x.getTable().accept(this);
+        }
+        return false;
+    }
 
     
 }

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

@@ -0,0 +1,22 @@
+package com.alibaba.druid.bvt.sql.odps;
+
+import junit.framework.TestCase;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.dialect.odps.parser.OdpsStatementParser;
+import com.alibaba.druid.sql.parser.Token;
+
+public class OdpsSetLabelTest extends TestCase {
+
+    public void test_odps() throws Exception {
+        String sql = "SET LABEL S3 TO USER aliyun$abc@alibaba-inc.com";
+        OdpsStatementParser parser = new OdpsStatementParser(sql);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+        parser.match(Token.EOF);
+        String output = SQLUtils.toOdpsString(stmt);
+        Assert.assertEquals("SET LABEL S3 TO USER aliyun$abc@alibaba-inc.com", output);
+    }
+}

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

@@ -0,0 +1,22 @@
+package com.alibaba.druid.bvt.sql.odps;
+
+import junit.framework.TestCase;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.dialect.odps.parser.OdpsStatementParser;
+import com.alibaba.druid.sql.parser.Token;
+
+public class OdpsSetLabelTest2 extends TestCase {
+
+    public void test_odps() throws Exception {
+        String sql = "SET LABEL S3 TO TABLE xx(f1,f2)";
+        OdpsStatementParser parser = new OdpsStatementParser(sql);
+        SQLStatement stmt = parser.parseStatementList().get(0);
+        parser.match(Token.EOF);
+        String output = SQLUtils.toOdpsString(stmt);
+        Assert.assertEquals("SET LABEL S3 TO TABLE xx(f1, f2)", output);
+    }
+}