improved odps sql formatter

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

@@ -74,8 +74,6 @@ public class Lexer {
 
     protected CommentHandler commentHandler;
 
-    protected boolean        hasComment   = false;
-
     protected boolean        endOfComment = false;
     
     protected boolean        keepComments = false;
@@ -289,8 +287,9 @@ public class Lexer {
 
     public final void nextToken() {
         bufPos = 0;
-        hasComment = false;
-        comments = null;
+        if (comments != null) {
+            comments = null;
+        }
 
         for (;;) {
             if (isWhitespace(ch)) {
@@ -804,7 +803,6 @@ public class Lexer {
 
         stringVal = subString(mark, bufPos);
         token = Token.MULTI_LINE_COMMENT;
-        hasComment = true;
         if (keepComments) {
             addComment(stringVal);
         }
@@ -852,7 +850,6 @@ public class Lexer {
 
         stringVal = subString(mark, bufPos);
         token = Token.LINE_COMMENT;
-        hasComment = true;
         if (keepComments) {
             addComment(stringVal);
         }
@@ -1020,23 +1017,9 @@ public class Lexer {
         return stringVal;
     }
     
-    public final String readAndResetCommentValue() {
-        if (comments.size() == 0) {
-            return null;
-        }
-        
-        String val = comments.get(0);
-        
-        this.hasComment = false;
-        this.comments.clear();
-        
-        return val;
-    }
-    
     public final List<String> readAndResetComments() {
         List<String> comments = this.comments;
         
-        this.hasComment = false;
         this.comments = null;
         
         return comments;
@@ -1153,7 +1136,7 @@ public class Lexer {
     }
 
     public boolean hasComment() {
-        return hasComment;
+        return comments != null;
     }
     
     public void skipToEOF() {

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

@@ -185,6 +185,11 @@ public class SQLExprParser extends SQLParser {
     }
 
     public SQLExpr primary() {
+      List<String> beforeComments = null;
+      if (lexer.isKeepComments() && lexer.hasComment()) {
+          beforeComments = lexer.readAndResetComments();
+      }
+
         SQLExpr sqlExpr = null;
 
         final Token tok = lexer.token();
@@ -555,7 +560,13 @@ public class SQLExprParser extends SQLParser {
                 throw new ParserException("ERROR. token : " + tok + ", pos : " + lexer.pos());
         }
 
-        return primaryRest(sqlExpr);
+        SQLExpr expr = primaryRest(sqlExpr);
+        
+        if (beforeComments != null) {
+            expr.addBeforeComment(beforeComments);
+        }
+        
+        return expr;
     }
 
     protected SQLExpr parseAll() {
@@ -1184,6 +1195,7 @@ public class SQLExprParser extends SQLParser {
                 }
                 
                 lexer.nextToken();
+                
                 SQLExpr rightExp = relational();
 
                 expr = new SQLBinaryOpExpr(expr, SQLBinaryOperator.BooleanAnd, rightExp, getDbType());

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

@@ -162,7 +162,7 @@ public class SQLSelectParser extends SQLParser {
         SQLSelectQueryBlock queryBlock = new SQLSelectQueryBlock();
         
         if (lexer.hasComment() && lexer.isKeepComments()) {
-            queryBlock.addBeforeComment(lexer.readAndResetCommentValue());
+            queryBlock.addBeforeComment(lexer.readAndResetComments());
         }
 
         accept(Token.SELECT);
@@ -238,12 +238,20 @@ public class SQLSelectParser extends SQLParser {
         if (lexer.token() == Token.WHERE) {
             lexer.nextToken();
 
+            List<String> beforeComments = null;
+            if (lexer.hasComment() && lexer.isKeepComments()) {
+                beforeComments = lexer.readAndResetComments();
+            }
             SQLExpr where = expr();
             
+            if (where != null && beforeComments != null) {
+                where.addBeforeComment(beforeComments);
+            }
+            
             if (lexer.hasComment() && lexer.isKeepComments() //
                     && lexer.token() != Token.INSERT // odps multi-insert
                     ) {
-                where.addAfterComment(lexer.readAndResetCommentValue());
+                where.addAfterComment(lexer.readAndResetComments());
             }
             
             queryBlock.setWhere(where);