improved sql format output

src/main/java/com/alibaba/druid/sql/SQLUtils.java

@@ -226,7 +226,10 @@ public class SQLUtils {
 
     public static String format(String sql, String dbType, List<Object> parameters) {
         try {
-            List<SQLStatement> statementList = toStatementList(sql, dbType);
+            SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
+            parser.setKeepComments(true);
+            
+            List<SQLStatement> statementList = parser.parseStatementList();
 
             return toSQLString(statementList, dbType, parameters);
         } catch (ParserException ex) {
@@ -247,10 +250,27 @@ public class SQLUtils {
         }
 
         for (int i = 0; i < statementList.size(); i++) {
+            SQLStatement stmt = statementList.get(i);
+            
             if (i > 0) {
-                out.append(";\n");
+                visitor.println(";");
+            }
+            
+            List<String> comments = stmt.getBeforeCommentsDirect();
+            if (comments != null){
+                for(String comment : comments) {
+                    visitor.println(comment);
+                }
+            }
+            stmt.accept(visitor);
+            
+            if (i == statementList.size() - 1) {
+                Boolean semi = (Boolean) stmt.getAttribute("format.semi");
+                if (semi != null && semi.booleanValue()) {
+                    visitor.println();
+                    visitor.print(";");
+                }
             }
-            statementList.get(i).accept(visitor);
         }
 
         return out.toString();

src/main/java/com/alibaba/druid/sql/ast/SQLObject.java

@@ -15,6 +15,7 @@
  */
 package com.alibaba.druid.sql.ast;
 
+import java.util.List;
 import java.util.Map;
 
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
@@ -34,6 +35,18 @@ public interface SQLObject {
     void putAttribute(String name, Object value);
 
     Map<String, Object> getAttributesDirect();
+    
+    void addBeforeComment(String comment);
+    
+    List<String> getBeforeCommentsDirect();
+    
+    void addAfterComment(String comment);
+    
+    List<String> getAfterCommentsDirect();
+    
+    boolean hasBeforeComment();
+    
+    boolean hasAfterComment()
 
     void output(StringBuffer buf);
 }

src/main/java/com/alibaba/druid/sql/ast/SQLObjectImpl.java

@@ -15,6 +15,7 @@
  */
 package com.alibaba.druid.sql.ast;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -107,4 +108,70 @@ public abstract class SQLObjectImpl implements SQLObject {
     public Map<String, Object> getAttributesDirect() {
         return attributes;
     }
+    
+    @SuppressWarnings("unchecked")
+    public void addBeforeComment(String comment) {
+        if (attributes == null) {
+            attributes = new HashMap<String, Object>(1);
+        }
+        
+        List<String> comments = (List<String>) attributes.get("format.before_comment");
+        if (comments == null) {
+            comments = new ArrayList<String>(2);
+            attributes.put("format.before_comment", comments);
+        }
+        
+        comments.add(comment);
+    }
+    
+    @SuppressWarnings("unchecked")
+    public List<String> getBeforeCommentsDirect() {
+        if (attributes == null) {
+            return null;
+        }
+        
+        return (List<String>) attributes.get("format.before_comment");
+    }
+    
+    @SuppressWarnings("unchecked")
+    public void addAfterComment(String comment) {
+        if (attributes == null) {
+            attributes = new HashMap<String, Object>(1);
+        }
+        
+        List<String> comments = (List<String>) attributes.get("format.after_comment");
+        if (comments == null) {
+            comments = new ArrayList<String>(2);
+            attributes.put("format.after_comment", comments);
+        }
+        
+        comments.add(comment);
+    }
+    
+    @SuppressWarnings("unchecked")
+    public List<String> getAfterCommentsDirect() {
+        if (attributes == null) {
+            return null;
+        }
+        
+        return (List<String>) attributes.get("format.after_comment");
+    }
+    
+    public boolean hasBeforeComment() {
+        List<String> comments = getBeforeCommentsDirect();
+        if (comments == null) {
+            return false;
+        }
+        
+        return !comments.isEmpty();
+    }
+    
+    public boolean hasAfterComment() {
+        List<String> comments = getAfterCommentsDirect();
+        if (comments == null) {
+            return false;
+        }
+        
+        return !comments.isEmpty();
+    }
 }

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

@@ -109,6 +109,9 @@ public class MySqlLexer extends Lexer {
         stringVal = subString(mark, bufPos);
         token = Token.LINE_COMMENT;
         hasComment = true;
+        if (keepComments) {
+            commentVal = stringVal;
+        }
 
         if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
             return;
@@ -448,13 +451,16 @@ public class MySqlLexer extends Lexer {
                 stringVal = subString(mark, bufPos);
                 token = Token.MULTI_LINE_COMMENT;
                 hasComment = true;
-            }
-
-            if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
-                return;
+                if (keepComments) {
+                    commentVal = stringVal;
+                }
             }
 
             endOfComment = isEOF();
+            
+            if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
+                return;
+            }
 
             if (!isHint && !isAllowComment() && !isSafeComment(stringVal)) {
                 throw new NotAllowCommentException();
@@ -492,6 +498,9 @@ public class MySqlLexer extends Lexer {
             stringVal = subString(mark, bufPos + 1);
             token = Token.LINE_COMMENT;
             hasComment = true;
+            if (keepComments) {
+                commentVal = stringVal;
+            }
 
             if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
                 return;

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

@@ -48,6 +48,11 @@ public class OdpsExprParser extends SQLExprParser {
         this.lexer.nextToken();
     }
     
+    public OdpsExprParser(String sql, Lexer.CommentHandler commentHandler){
+        this(new OdpsLexer(sql, commentHandler));
+        this.lexer.nextToken();
+    }
+    
     protected SQLExpr parseAliasExpr(String alias) {
         return new SQLCharExpr(alias);
     }
@@ -97,7 +102,13 @@ public class OdpsExprParser extends SQLExprParser {
         }
 
         final String alias = as();
+        
+        SQLSelectItem item = new SQLSelectItem(expr, alias);
+        
+        if (lexer.hasComment() && lexer.isKeepComments()) {
+            item.addAfterComment(lexer.commentVal());
+        }
 
-        return new SQLSelectItem(expr, alias);
+        return item;
     }
 }

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

@@ -52,10 +52,17 @@ public class OdpsLexer extends Lexer {
         super.keywods = DEFAULT_ODPS_KEYWORDS;
     }
     
+    public OdpsLexer(String input, CommentHandler commentHandler){
+        super(input, commentHandler);
+        super.keywods = DEFAULT_ODPS_KEYWORDS;
+    }
+    
     public void scanComment() {
         if (ch != '/' && ch != '-') {
             throw new IllegalStateException();
         }
+        
+        Token lastToken = this.token;
 
         mark = pos;
         bufPos = 0;
@@ -95,9 +102,16 @@ public class OdpsLexer extends Lexer {
                 stringVal = subString(mark + startHintSp, (bufPos - startHintSp) - 1);
                 token = Token.HINT;
             } else {
-                stringVal = subString(mark, bufPos);
+                stringVal = subString(mark, bufPos + 1);
                 token = Token.MULTI_LINE_COMMENT;
                 hasComment = true;
+                if (keepComments) {
+                    commentVal = stringVal;
+                }
+            }
+            
+            if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
+                return;
             }
 
             if (token != Token.HINT && !isAllowComment()) {
@@ -138,10 +152,18 @@ public class OdpsLexer extends Lexer {
                 bufPos++;
             }
 
-            stringVal = subString(mark + 1, bufPos);
+            stringVal = subString(mark, bufPos);
             token = Token.LINE_COMMENT;
             hasComment = true;
+            if (keepComments) {
+                commentVal = stringVal;
+            }
             endOfComment = isEOF();
+            
+            if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
+                return;
+            }
+            
             return;
         }
     }

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

@@ -40,14 +40,18 @@ public class OdpsSelectParser extends SQLSelectParser {
             return queryRest(select);
         }
 
+        OdpsSelectQueryBlock queryBlock = new OdpsSelectQueryBlock();
+        
+        if (lexer.hasComment() && lexer.isKeepComments()) {
+            queryBlock.addBeforeComment(lexer.commentVal());
+        }
+        
         accept(Token.SELECT);
 
         if (lexer.token() == Token.COMMENT) {
             lexer.nextToken();
         }
 
-        OdpsSelectQueryBlock queryBlock = new OdpsSelectQueryBlock();
-
         if (lexer.token() == Token.DISTINCT) {
             queryBlock.setDistionOption(SQLSetQuantifier.DISTINCT);
             lexer.nextToken();

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

@@ -42,9 +42,11 @@ import com.alibaba.druid.sql.parser.Token;
 import com.alibaba.druid.util.JdbcConstants;
 
 public class OdpsStatementParser extends SQLStatementParser {
-
     public OdpsStatementParser(String sql){
-        super(new OdpsExprParser(sql));
+        super (new OdpsLexer(sql), JdbcConstants.ODPS);
+        this.exprParser = new OdpsExprParser(this.lexer);
+        this.lexer.setCommentHandler(commentCallBack);
+        this.lexer.nextToken();
     }
 
     public OdpsStatementParser(SQLExprParser exprParser){
@@ -108,6 +110,11 @@ public class OdpsStatementParser extends SQLStatementParser {
 
     public OdpsInsert parseOdpsInsert() {
         OdpsInsert insert = new OdpsInsert();
+        
+        if (lexer.hasComment()) {
+            insert.addBeforeComment(lexer.commentVal());
+        }
+        
         SQLSelectParser selectParser = createSQLSelectParser();
 
         accept(Token.INSERT);

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

@@ -153,6 +153,10 @@ public class OdpsOutputVisitor extends SQLASTOutputVisitor implements OdpsASTVis
 
     @Override
     public boolean visit(OdpsInsert x) {
+        if (x.hasBeforeComment()) {
+            printComment(x.getBeforeCommentsDirect(), "\n");
+            println();
+        }
         if (x.isOverwrite()) {
             print("INSERT OVERWRITE TABLE ");
         } else {
@@ -231,12 +235,22 @@ public class OdpsOutputVisitor extends SQLASTOutputVisitor implements OdpsASTVis
     protected void printSelectList(List<SQLSelectItem> selectList) {
         incrementIndent();
         for (int i = 0, size = selectList.size(); i < size; ++i) {
+            SQLSelectItem selectItem = selectList.get(i);
+            
             if (i != 0) {
                 print(", ");
+                
+                printComment(selectList.get(i - 1).getAfterCommentsDirect(), "\n");
+                
                 println();
             }
-
-            selectList.get(i).accept(this);
+            
+            selectItem.accept(this);
+            
+            if (i == selectList.size() - 1 && selectItem.hasAfterComment()) {
+                print(' ');
+                printComment(selectItem.getAfterCommentsDirect(), "\n");
+            }
         }
         decrementIndent();
     }
@@ -374,6 +388,11 @@ public class OdpsOutputVisitor extends SQLASTOutputVisitor implements OdpsASTVis
 
     @Override
     public boolean visit(OdpsSelectQueryBlock x) {
+        if (isPrettyFormat() && x.hasBeforeComment()) {
+            printComment(x.getBeforeCommentsDirect(), "\n");
+            println();
+        }
+        
         print("SELECT ");
 
         if (SQLSetQuantifier.ALL == x.getDistionOption()) {
@@ -397,6 +416,10 @@ public class OdpsOutputVisitor extends SQLASTOutputVisitor implements OdpsASTVis
             print("WHERE ");
             x.getWhere().setParent(x);
             x.getWhere().accept(this);
+            if (x.getWhere().hasAfterComment() && isPrettyFormat()) {
+                print(' ');
+                printComment(x.getWhere().getAfterCommentsDirect(), "\n");
+            }
         }
 
         if (x.getGroupBy() != null) {

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

@@ -231,6 +231,9 @@ public class OracleLexer extends Lexer {
                 stringVal = subString(mark, bufPos);
                 token = Token.MULTI_LINE_COMMENT;
                 hasComment = true;
+                if (keepComments) {
+                    commentVal = stringVal;
+                }
             }
 
             if (token != Token.HINT && !isAllowComment()) {
@@ -274,6 +277,9 @@ public class OracleLexer extends Lexer {
             stringVal = subString(mark + 1, bufPos);
             token = Token.LINE_COMMENT;
             hasComment = true;
+            if (keepComments) {
+                commentVal = stringVal;
+            }
             endOfComment = isEOF();
             return;
         }

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

@@ -106,6 +106,9 @@ public class SQLServerLexer extends Lexer {
                 stringVal = subString(mark, bufPos);
                 token = Token.MULTI_LINE_COMMENT;
                 hasComment = true;
+                if (keepComments) {
+                    commentVal = stringVal;
+                }
             }
 
             if (token != Token.HINT && !isAllowComment() && !isSafeComment(stringVal)) {
@@ -145,6 +148,9 @@ public class SQLServerLexer extends Lexer {
             stringVal = subString(mark + 1, bufPos);
             token = Token.LINE_COMMENT;
             hasComment = true;
+            if (keepComments) {
+                commentVal = stringVal;
+            }
             endOfComment = isEOF();
             
             if (!isAllowComment() && (isEOF() || !isSafeComment(stringVal))) {

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

@@ -56,6 +56,8 @@ public class Lexer {
     protected Keywords     keywods      = Keywords.DEFAULT_KEYWORDS;
 
     protected String       stringVal;
+    
+    protected String       commentVal;
 
     protected boolean      skipComment  = true;
 
@@ -73,9 +75,24 @@ public class Lexer {
     protected boolean        hasComment   = false;
 
     protected boolean        endOfComment = false;
+    
+    protected boolean        keepComments = false;
 
     public Lexer(String input){
+        this(input, null);
+    }
+    
+    public Lexer(String input, CommentHandler commentHandler){
         this(input, true);
+        this.commentHandler = commentHandler;
+    }
+    
+    public boolean isKeepComments() {
+        return keepComments;
+    }
+    
+    public void setKeepComments(boolean keepComments) {
+        this.keepComments = keepComments;
     }
 
     public CommentHandler getCommentHandler() {
@@ -270,6 +287,7 @@ public class Lexer {
 
     public final void nextToken() {
         bufPos = 0;
+        hasComment = false;
 
         for (;;) {
             if (isWhitespace(ch)) {
@@ -759,6 +777,8 @@ public class Lexer {
     }
 
     private void scanMultiLineComment() {
+        Token lastToken = this.token;
+        
         scanChar();
         scanChar();
         mark = pos;
@@ -782,12 +802,25 @@ public class Lexer {
         stringVal = subString(mark, bufPos);
         token = Token.MULTI_LINE_COMMENT;
         hasComment = true;
+        if (keepComments) {
+            commentVal = stringVal;
+        }
+        
+        if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
+            return;
+        }
+        
+        if (!isAllowComment() && !isSafeComment(stringVal)) {
+            throw new NotAllowCommentException();
+        }
     }
 
     private void scanSingleLineComment() {
+        Token lastToken = this.token;
+        
         scanChar();
         scanChar();
-        mark = pos;
+        mark = pos - 1;
         bufPos = 0;
 
         for (;;) {
@@ -817,6 +850,17 @@ public class Lexer {
         stringVal = subString(mark, bufPos);
         token = Token.LINE_COMMENT;
         hasComment = true;
+        if (keepComments) {
+            commentVal = stringVal;
+        }
+        
+        if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
+            return;
+        }
+        
+        if (!isAllowComment() && !isSafeComment(stringVal)) {
+            throw new NotAllowCommentException();
+        }
     }
 
     public void scanIdentifier() {
@@ -972,6 +1016,10 @@ public class Lexer {
     public final String stringVal() {
         return stringVal;
     }
+    
+    public final String commentVal() {
+        return commentVal;
+    }
 
     private boolean isOperator(char ch) {
         switch (ch) {
@@ -1080,11 +1128,10 @@ public class Lexer {
     }
 
     public static interface CommentHandler {
-
         boolean handle(Token lastToken, String comment);
     }
 
-    public boolean isHasComment() {
+    public boolean hasComment() {
         return hasComment;
     }
     

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

@@ -1179,6 +1179,10 @@ public class SQLExprParser extends SQLParser {
     public SQLExpr andRest(SQLExpr expr) {
         for (;;) {
             if (lexer.token() == Token.AND || lexer.token() == Token.AMPAMP) {
+                if (lexer.isKeepComments() && lexer.hasComment()) {
+                    expr.addAfterComment(lexer.commentVal());
+                }
+                
                 lexer.nextToken();
                 SQLExpr rightExp = relational();
 
@@ -1682,7 +1686,12 @@ public class SQLExprParser extends SQLParser {
             item.setValue(new SQLIdentifierExpr(lexer.stringVal()));
             lexer.nextToken();
         } else {
-            item.setValue(expr());
+            if (lexer.token() == Token.ALL) {
+                item.setValue(new SQLIdentifierExpr(lexer.stringVal()));
+                lexer.nextToken();
+            } else {
+                item.setValue(expr());
+            }
         }
         
         return item;

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

@@ -158,6 +158,12 @@ public class SQLSelectParser extends SQLParser {
 
             return queryRest(select);
         }
+        
+        SQLSelectQueryBlock queryBlock = new SQLSelectQueryBlock();
+        
+        if (lexer.hasComment() && lexer.isKeepComments()) {
+            queryBlock.addBeforeComment(lexer.commentVal());
+        }
 
         accept(Token.SELECT);
 
@@ -165,8 +171,6 @@ public class SQLSelectParser extends SQLParser {
             lexer.nextToken();
         }
 
-        SQLSelectQueryBlock queryBlock = new SQLSelectQueryBlock();
-
         if (lexer.token() == Token.DISTINCT) {
             queryBlock.setDistionOption(SQLSetQuantifier.DISTINCT);
             lexer.nextToken();
@@ -234,7 +238,15 @@ public class SQLSelectParser extends SQLParser {
         if (lexer.token() == Token.WHERE) {
             lexer.nextToken();
 
-            queryBlock.setWhere(expr());
+            SQLExpr where = expr();
+            
+            if (lexer.hasComment() && lexer.isKeepComments() //
+                    && lexer.token() != Token.INSERT // odps multi-insert
+                    ) {
+                where.addAfterComment(lexer.commentVal());
+            }
+            
+            queryBlock.setWhere(where);
         }
     }
 
@@ -326,7 +338,13 @@ public class SQLSelectParser extends SQLParser {
 
         parseTableSourceQueryTableExpr(tableReference);
 
-        return parseTableSourceRest(tableReference);
+        SQLTableSource tableSrc = parseTableSourceRest(tableReference);
+        
+        if (lexer.hasComment() && lexer.isKeepComments()) {
+            tableSrc.addAfterComment(lexer.commentVal());
+        }
+        
+        return tableSrc;
     }
 
     private void parseTableSourceQueryTableExpr(SQLExprTableSource tableReference) {

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

@@ -88,12 +88,16 @@ import com.alibaba.druid.sql.dialect.odps.parser.OdpsSelectParser;
 
 public class SQLStatementParser extends SQLParser {
 
-    protected SQLExprParser exprParser;
+    protected SQLExprParser   exprParser;
 
-    protected boolean       parseCompleteValues = true;
-
-    protected int           parseValuesSize     = 3;
+    protected boolean         parseCompleteValues = true;
 
+    protected int             parseValuesSize     = 3;
+    
+    protected List<String>    comments;
+    
+    protected CommentCallback commentCallBack     = new CommentCallback();
+    
     public SQLStatementParser(String sql){
         this(sql, null);
     }
@@ -101,11 +105,23 @@ public class SQLStatementParser extends SQLParser {
     public SQLStatementParser(String sql, String dbType){
         this(new SQLExprParser(sql, dbType));
     }
-
+    
     public SQLStatementParser(SQLExprParser exprParser){
         super(exprParser.getLexer(), exprParser.getDbType());
         this.exprParser = exprParser;
     }
+    
+    protected SQLStatementParser(Lexer lexer, String dbType){
+        super(lexer, dbType);
+    }
+    
+    public boolean isKeepComments() {
+        return lexer.isKeepComments();
+    }
+    
+    public void setKeepComments(boolean keepComments) {
+        this.lexer.setKeepComments(keepComments);
+    }
 
     public SQLExprParser getExprParser() {
         return exprParser;
@@ -122,13 +138,30 @@ public class SQLStatementParser extends SQLParser {
     }
 
     public void parseStatementList(List<SQLStatement> statementList, int max) {
+        List<String> comments = null;
+        if (this.comments != null) {
+            comments = this.comments;
+            this.comments = null;
+        }
+        
+        int size_original = statementList.size();
+        
         for (;;) {
+            boolean isFirstStmt = statementList.size() == size_original + 1;
+            if (isFirstStmt && comments != null && lexer.isKeepComments()) {
+                SQLStatement firstStmt = statementList.get(size_original);
+                for (String comment : comments) {
+                    firstStmt.addBeforeComment(comment);
+                }
+                comments = null;
+            }
+            
             if (max != -1) {
                 if (statementList.size() >= max) {
                     return;
                 }
             }
-
+            
             if (lexer.token() == Token.EOF) {
                 return;
             }
@@ -136,8 +169,15 @@ public class SQLStatementParser extends SQLParser {
                 return;
             }
 
-            if (lexer.token() == (Token.SEMI)) {
+            if (lexer.token() == Token.SEMI) {
+                size_original = statementList.size();
+                this.comments = null;
                 lexer.nextToken();
+                comments = this.comments;
+                if (lexer.isKeepComments()) {
+                    SQLStatement stmt = statementList.get(statementList.size() - 1);
+                    stmt.getAttributes().put("format.semi", Boolean.TRUE);
+                }
                 continue;
             }
 
@@ -1781,4 +1821,22 @@ public class SQLStatementParser extends SQLParser {
     public void setParseValuesSize(int parseValuesSize) {
         this.parseValuesSize = parseValuesSize;
     }
+    
+    class CommentCallback implements Lexer.CommentHandler {
+
+        @Override
+        public boolean handle(Token lastToken, String comment) {
+            addComment(comment);
+            
+            return true;
+        }
+        
+    }
+    
+    void addComment(String comment) {
+        if (comments == null) {
+            comments = new ArrayList<String>();
+        }
+        comments.add(comment);
+    }
 }

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

@@ -66,7 +66,7 @@ public class ParameterizedOutputVisitorUtils {
             stmt.accept(visitor);
         }
 
-        if (visitor.getReplaceCount() == 0 && !parser.getLexer().isHasComment()) {
+        if (visitor.getReplaceCount() == 0 && !parser.getLexer().hasComment()) {
             return sql;
         }
 

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

@@ -350,6 +350,10 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
             visitBinaryLeft(item, x.getOperator());
 
             if (relational) {
+                if (isPrettyFormat() && item.hasAfterComment()) {
+                    print(' ');
+                    printComment(item.getAfterCommentsDirect(), "\n");
+                }
                 println();
             } else {
                 print(" ");
@@ -684,6 +688,10 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
     }
 
     public boolean visit(SQLSelectQueryBlock x) {
+        if (isPrettyFormat() && x.hasBeforeComment()) {
+            printComment(x.getBeforeCommentsDirect(), "\n");
+        }
+        
         print("SELECT ");
 
         if (SQLSetQuantifier.ALL == x.getDistionOption()) {
@@ -763,11 +771,16 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
 
     public boolean visit(SQLExprTableSource x) {
         x.getExpr().accept(this);
-
+        
         if (x.getAlias() != null) {
             print(' ');
             print(x.getAlias());
         }
+        
+        if (isPrettyFormat() && x.hasAfterComment()) {
+            print(' ');
+            printComment(x.getAfterCommentsDirect(), "\n");
+        }
 
         return false;
     }
@@ -2182,4 +2195,16 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab
         }
         return false;
     }
+    
+    protected void printComment(List<String> comments, String seperator) {
+        if (comments != null) {
+            for (int i = 0; i < comments.size(); ++i) {
+                if (i != 0) {
+                    print(seperator);
+                }
+                String comment = comments.get(i);
+                print(comment);
+            }
+        }
+    }
 }