improved odps sql formatter

src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlLexer.java

@@ -108,7 +108,6 @@ public class MySqlLexer extends Lexer {
 
         stringVal = subString(mark - 1, bufPos + 1);
         token = Token.LINE_COMMENT;
-        hasComment = true;
         if (keepComments) {
             addComment(stringVal);
         }
@@ -450,7 +449,6 @@ public class MySqlLexer extends Lexer {
             } else {
                 stringVal = subString(mark, bufPos);
                 token = Token.MULTI_LINE_COMMENT;
-                hasComment = true;
                 if (keepComments) {
                     addComment(stringVal);
                 }
@@ -497,7 +495,6 @@ public class MySqlLexer extends Lexer {
 
             stringVal = subString(mark, bufPos + 1);
             token = Token.LINE_COMMENT;
-            hasComment = true;
             if (keepComments) {
                 addComment(stringVal);
             }

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

@@ -112,7 +112,6 @@ public class OdpsLexer extends Lexer {
             } else {
                 stringVal = subString(mark, bufPos + 1);
                 token = Token.MULTI_LINE_COMMENT;
-                hasComment = true;
                 if (keepComments) {
                     addComment(stringVal);
                 }
@@ -160,9 +159,8 @@ public class OdpsLexer extends Lexer {
                 bufPos++;
             }
 
-            stringVal = subString(mark, bufPos);
+            stringVal = subString(mark, ch != EOI ? bufPos : bufPos + 1);
             token = Token.LINE_COMMENT;
-            hasComment = true;
             if (keepComments) {
                 addComment(stringVal);
             }

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

@@ -450,6 +450,10 @@ public class OdpsOutputVisitor extends SQLASTOutputVisitor implements OdpsASTVis
         if (x.getWhere() != null) {
             println();
             print("WHERE ");
+            if (x.getWhere().hasBeforeComment() && isPrettyFormat()) {
+                printlnComments(x.getWhere().getBeforeCommentsDirect());
+            }
+            
             x.getWhere().setParent(x);
             x.getWhere().accept(this);
             if (x.getWhere().hasAfterComment() && isPrettyFormat()) {

src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleLexer.java

@@ -230,7 +230,6 @@ public class OracleLexer extends Lexer {
             } else {
                 stringVal = subString(mark, bufPos);
                 token = Token.MULTI_LINE_COMMENT;
-                hasComment = true;
                 if (keepComments) {
                     addComment(stringVal);
                 }
@@ -276,7 +275,6 @@ public class OracleLexer extends Lexer {
 
             stringVal = subString(mark + 1, bufPos);
             token = Token.LINE_COMMENT;
-            hasComment = true;
             if (keepComments) {
                 addComment(stringVal);
             }

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/parser/SQLServerLexer.java

@@ -105,7 +105,6 @@ public class SQLServerLexer extends Lexer {
             } else {
                 stringVal = subString(mark, bufPos);
                 token = Token.MULTI_LINE_COMMENT;
-                hasComment = true;
                 if (keepComments) {
                     addComment(stringVal);
                 }
@@ -147,7 +146,6 @@ public class SQLServerLexer extends Lexer {
 
             stringVal = subString(mark + 1, bufPos);
             token = Token.LINE_COMMENT;
-            hasComment = true;
             if (keepComments) {
                 addComment(stringVal);
             }

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);

src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java

@@ -356,13 +356,30 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
 
         for (int i = groupList.size() - 1; i >= 0; --i) {
             SQLExpr item = groupList.get(i);
+            
+            if (relational) {
+                if (isPrettyFormat() && item.hasBeforeComment()) {
+                    printlnComments(item.getBeforeCommentsDirect());
+                }
+            }
+            
+            if (isPrettyFormat() && item.hasBeforeComment()) {
+                printlnComments(item.getBeforeCommentsDirect());
+            }
+            
             visitBinaryLeft(item, x.getOperator());
 
+            if (isPrettyFormat() && item.hasAfterComment()) {
+                print(' ');
+                printComment(item.getAfterCommentsDirect(), "\n");
+            }
+            
+            if (i != groupList.size() - 1 && isPrettyFormat() && item.getParent().hasAfterComment()) {
+                print(' ');
+                printComment(item.getParent().getAfterCommentsDirect(), "\n");
+            }
+            
             if (relational) {
-                if (isPrettyFormat() && item.hasAfterComment()) {
-                    print(' ');
-                    printComment(item.getAfterCommentsDirect(), "\n");
-                }
                 println();
             } else {
                 print(" ");
@@ -381,6 +398,10 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
     }
 
     private void visitorBinaryRight(SQLBinaryOpExpr x) {
+        if (isPrettyFormat() && x.getRight().hasBeforeComment()) {
+            printlnComments(x.getRight().getBeforeCommentsDirect());
+        }
+        
         if (x.getRight() instanceof SQLBinaryOpExpr) {
             SQLBinaryOpExpr right = (SQLBinaryOpExpr) x.getRight();
             boolean rightRational = right.getOperator() == SQLBinaryOperator.BooleanAnd
@@ -404,6 +425,11 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
         } else {
             x.getRight().accept(this);
         }
+        
+        if (x.getRight().hasAfterComment() && isPrettyFormat()) {
+            print(' ');
+            printlnComments(x.getRight().getAfterCommentsDirect());
+        }
     }
 
     private void visitBinaryLeft(SQLExpr left, SQLBinaryOperator op) {
@@ -2140,9 +2166,10 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
     public boolean visit(SQLUnionQueryTableSource x) {
         print("(");
         incrementIndent();
-        x.getUnion().accept(this);
         println();
+        x.getUnion().accept(this);
         decrementIndent();
+        println();
         print(")");
 
         if (x.getAlias() != null) {
@@ -2217,6 +2244,15 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
         }
     }
     
+    protected void printlnComments(List<String> comments) {
+        if (comments != null) {
+            for (int i = 0; i < comments.size(); ++i) {
+                String comment = comments.get(i);
+                println(comment);
+            }
+        }
+    }
+    
     @Override
     public boolean visit(SQLAlterViewRenameStatement x) {
         print("ALTER VIEW ");

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

@@ -0,0 +1,20 @@
+package com.alibaba.druid.bvt.sql.odps;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.SQLUtils;
+
+import junit.framework.TestCase;
+
+public class OdpsFormatCommentTest11 extends TestCase {
+    public void test_column_comment() throws Exception {
+        String sql = "select * from t where -- comment_xx"
+                + "\n f1 > 0";
+        Assert.assertEquals("SELECT *" //
+                + "\nFROM t" //
+                + "\nWHERE -- comment_xx"
+                + "\nf1 > 0", SQLUtils.formatOdps(sql));
+    }
+
+   
+}

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

@@ -0,0 +1,17 @@
+package com.alibaba.druid.bvt.sql.odps;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.SQLUtils;
+
+import junit.framework.TestCase;
+
+public class OdpsFormatCommentTest12 extends TestCase {
+    public void test_column_comment() throws Exception {
+        String sql = "select * from t --abc";
+        Assert.assertEquals("SELECT *" //
+                + "\nFROM t --abc", SQLUtils.formatOdps(sql));
+    }
+
+   
+}

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

@@ -0,0 +1,26 @@
+package com.alibaba.druid.bvt.sql.odps;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.SQLUtils;
+
+import junit.framework.TestCase;
+
+public class OdpsFormatCommentTest13 extends TestCase {
+    public void test_column_comment() throws Exception {
+        String sql = "select * from t where f0 > 0 -- comment_0"
+                + "\n and -- comment_1"
+                + "\n f1 > 1 -- comment_2"
+                + "\n and -- comment_3"
+                + "\n f2 > 2 -- comment_4";
+        Assert.assertEquals("SELECT *"
+                + "\nFROM t"
+                + "\nWHERE f0 > 0 -- comment_0"
+                + "\n\tAND -- comment_1"
+                + "\n\tf1 > 1 -- comment_2"
+                + "\n\tAND -- comment_3"
+                + "\n\tf2 > 2 -- comment_4", SQLUtils.formatOdps(sql));
+    }
+
+   
+}

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

@@ -0,0 +1,38 @@
+package com.alibaba.druid.bvt.sql.odps;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.SQLUtils;
+
+import junit.framework.TestCase;
+
+public class OdpsFormatCommentTest14 extends TestCase {
+    public void test_column_comment() throws Exception {
+        String sql = "select * from (" //
+                + "select 1 from t1"
+                + "\n --c_0"
+                + "\n union all "
+                + "\n --c_1" //
+                + "\nselect 2 from t2" //
+                + "\n --c_2"
+                + "\n union all "
+                + "\n --c_3" //
+                + "\nselect 3 from t3" //
+                + ") xx";
+        Assert.assertEquals("SELECT *"
+                + "\nFROM ("
+                + "\n\tSELECT 1"
+                + "\n\tFROM t1 --c_0"
+                + "\n\tUNION ALL"
+                + "\n\t--c_1"
+                + "\n\tSELECT 2"
+                + "\n\tFROM t2 --c_2"
+                + "\n\tUNION ALL"
+                + "\n\t--c_3"
+                + "\n\tSELECT 3"
+                + "\n\tFROM t3"
+                + "\n) xx", SQLUtils.formatOdps(sql));
+    }
+
+   
+}

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

@@ -20,7 +20,7 @@ public class OdpsFormatCommentTest8 extends TestCase {
                 + "\nFROM table2;"
                 + "\nSELECT *"
                 + "\nFROM table3"
-                + "\n;--comment2 goes her", SQLUtils.formatOdps(sql));
+                + "\n;--comment2 goes here", SQLUtils.formatOdps(sql));
     }
 
    

src/test/resources/bvt/parser/mysql-6.txt

@@ -38,7 +38,8 @@ from (
 
 ---------------------------
 SELECT COUNT(*)
-FROM (SELECT a.activity_offer_id AS adResourceId, a.title AS adResourceName, a.b2boffer_url AS adResourceURL, a.custid, b.memberid
+FROM (
+	SELECT a.activity_offer_id AS adResourceId, a.title AS adResourceName, a.b2boffer_url AS adResourceURL, a.custid, b.memberid
 		, c.position_name AS positionName, c.charge_setting AS feeRegulation, a.priority_weight AS priorityWeight, a.audit_state AS adResourceState, 0 AS adResourceType
 		, a.gmt_create, a.refuse_remark
 	FROM q_activity_offer_enroll a, q_custinfo b, q_activity_position c
@@ -53,4 +54,4 @@ FROM (SELECT a.activity_offer_id AS adResourceId, a.title AS adResourceName, a.b
 	WHERE a.custid = b.custid
 		AND a.activity_position_id = c.activity_position_id
 		AND c.activity_id = ?
-	) a
+) a

src/test/resources/bvt/parser/mysql-8.txt

@@ -7,13 +7,14 @@ WHERE a.examine_id = b.examine_id
 
 ---------------------------
 SELECT a.examine_id AS examineId, a.examine_name AS examineName, a.level_limit AS levelLimit
-FROM c_examinepaper a, (SELECT examine_id
+FROM c_examinepaper a, (
+		SELECT examine_id
 		FROM c_examinepaper
 		WHERE level_limit = ?
 			AND status = 1
 		MINUS
 		SELECT DISTINCT examine_id
 		FROM course
-		) b
+	) b
 WHERE a.examine_id = b.examine_id