fix add comment & print comment

alibaba · May 13, 2023 · e8a37b5 · e8a37b5
1 parent 634ad1a
commit e8a37b5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectStatement.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectStatement.java
@@ -120,4 +120,16 @@ public List<String> computeSelecteListAlias() {
     public boolean addWhere(SQLExpr where) {
         return select.addWhere(where);
     }
+
+    public void addBeforeComment(String comment) {
+        if (headHints == null && attributes == null && select.hints == null && select.getBeforeCommentsDirect() == null) {
+            SQLSelectQueryBlock queryBlock = select.getQueryBlock();
+            List<String> queryBlockBeforeComments = queryBlock.getBeforeCommentsDirect();
+            if (queryBlockBeforeComments != null) {
+                queryBlockBeforeComments.add(comment);
+                return;
+            }
+        }
+        super.addBeforeComment(comment);
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
@@ -3022,6 +3022,10 @@ public boolean visit(SQLSelectStatement stmt) {
             }
         }
 
+        if (stmt.hasBeforeComment()) {
+            printlnComments(stmt.getBeforeCommentsDirect());
+        }
+
         SQLSelect select = stmt.getSelect();
         if (select != null) {
             this.visit(select);