From 58fa713694e026f17f0d2af2c97e04d65299e879 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=AB=98=E9=93=81?= <shaojin.wensj@alibaba-inc.com>
Date: Sat, 26 Dec 2020 06:08:25 +0800
Subject: [PATCH] support clickhouse limit by, fix #4110

---
 .../com/alibaba/druid/sql/ast/SQLLimit.java   | 53 +++++++++++++++----
 .../druid/sql/parser/SQLExprParser.java       | 15 ++++++
 .../sql/visitor/SQLASTOutputVisitor.java      | 11 ++++
 3 files changed, 69 insertions(+), 10 deletions(-)

diff --git a/src/main/java/com/alibaba/druid/sql/ast/SQLLimit.java b/src/main/java/com/alibaba/druid/sql/ast/SQLLimit.java
index 4ed6031e92..2bc82adeae 100644
--- a/src/main/java/com/alibaba/druid/sql/ast/SQLLimit.java
+++ b/src/main/java/com/alibaba/druid/sql/ast/SQLLimit.java
@@ -15,11 +15,9 @@
  */
 package com.alibaba.druid.sql.ast;
 
-import com.alibaba.druid.DbType;
-import com.alibaba.druid.sql.SQLUtils;
+import java.util.ArrayList;
+import java.util.List;
 import com.alibaba.druid.sql.ast.expr.SQLIntegerExpr;
-import com.alibaba.druid.sql.dialect.mysql.ast.MySqlObject;
-import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 import java.util.HashMap;
@@ -30,6 +28,7 @@ import java.util.HashMap;
 public final class SQLLimit extends SQLObjectImpl implements SQLReplaceable {
     private SQLExpr rowCount;
     private SQLExpr offset;
+    private List<SQLExpr> by;
 
     public SQLLimit() {
 
@@ -79,20 +78,26 @@ public final class SQLLimit extends SQLObjectImpl implements SQLReplaceable {
         this.offset = offset;
     }
 
-    public void merge(SQLLimit limit) {
-        if (limit == null) {
+    public void merge(SQLLimit other) {
+        if (other == null) {
             return;
         }
 
-        if (limit.offset != null) {
+        if (other.offset != null) {
             if (this.offset == null) {
-                this.offset = limit.offset.clone();
+                this.offset = other.offset.clone();
             }
         }
 
-        if (limit.rowCount != null) {
+        if (other.rowCount != null) {
             if (this.rowCount == null) {
-                this.rowCount = limit.rowCount.clone();
+                this.rowCount = other.rowCount.clone();
+            }
+        }
+
+        if (other.by != null) {
+            for (SQLExpr item : other.by) {
+                addBy(item.clone());
             }
         }
     }
@@ -107,6 +112,12 @@ public final class SQLLimit extends SQLObjectImpl implements SQLReplaceable {
             if (rowCount != null) {
                 rowCount.accept(visitor);
             }
+
+            if (by != null) {
+                for (SQLExpr item : by) {
+                    item.accept(visitor);
+                }
+            }
         }
         visitor.endVisit(this);
     }
@@ -122,6 +133,12 @@ public final class SQLLimit extends SQLObjectImpl implements SQLReplaceable {
             x.setRowCount(rowCount.clone());
         }
 
+        if (by != null) {
+            for (SQLExpr item : by) {
+                x.addBy(item);
+            }
+        }
+
         if (attributes != null) {
             x.attributes = (HashMap) ((HashMap) attributes).clone();
         }
@@ -129,6 +146,22 @@ public final class SQLLimit extends SQLObjectImpl implements SQLReplaceable {
         return x;
     }
 
+    public void addBy(SQLExpr item) {
+        if (item == null) {
+            return;
+        }
+
+        if (by == null) {
+            by = new ArrayList<>(1);
+        }
+        by.add(item);
+        item.setParent(this);
+    }
+
+    public List<SQLExpr> getBy() {
+        return by;
+    }
+
     @Override
     public boolean replace(SQLExpr expr, SQLExpr target) {
         if (rowCount == expr) {
diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java
index 7e39babb44..ffd61d48da 100644
--- a/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java
+++ b/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java
@@ -5711,6 +5711,21 @@ public class SQLExprParser extends SQLParser {
         } else {
             limit.setRowCount(temp);
         }
+
+        if (lexer.token == Token.BY && dbType == DbType.clickhouse) {
+            lexer.nextToken();
+
+            for (;;) {
+                SQLExpr item = this.expr();
+                limit.addBy(item);
+                if (lexer.token == Token.COMMA) {
+                    lexer.nextToken();
+                    continue;
+                }
+                break;
+            }
+        }
+
         return limit;
     }
 
diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
index d6c62747e3..20ee6a1960 100644
--- a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
+++ b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
@@ -7950,6 +7950,17 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Paramet
             printExpr(rowCount, parameterized);
         }
 
+        final List<SQLExpr> by = x.getBy();
+        if (by != null) {
+            print0(ucase ? " BY " : " by ");
+            for (int i = 0; i < by.size(); i++) {
+                if (i != 0) {
+                    print0(", ");
+                }
+                by.get(i).accept(this);
+            }
+        }
+
         return false;
     }