[fix][headless] Fix order by and group by not enclosed in backticks. (#…

…2041)
tencentmusic · Feb 10, 2025 · 3ca46be · 3ca46be
1 parent a8157ee
commit 3ca46be
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 3 deletions.
diff --git a/.../java/com/tencent/supersonic/common/jsqlparser/FieldAliasReplaceWithBackticksVisitor.java b/.../java/com/tencent/supersonic/common/jsqlparser/FieldAliasReplaceWithBackticksVisitor.java
@@ -5,8 +5,13 @@
 import net.sf.jsqlparser.statement.select.SelectItemVisitorAdapter;
 import org.apache.commons.lang3.StringUtils;
 
+import java.util.HashMap;
+import java.util.Map;
+
 public class FieldAliasReplaceWithBackticksVisitor extends SelectItemVisitorAdapter {
 
+    private Map<String, String> fieldAliasReplacedMap = new HashMap<>();
+
     @Override
     public void visit(SelectItem selectExpressionItem) {
         Alias alias = selectExpressionItem.getAlias();
@@ -19,6 +24,7 @@ public void visit(SelectItem selectExpressionItem) {
             return;
         }
         alias.setName(replaceValue);
+        fieldAliasReplacedMap.put(aliasName, replaceValue);
     }
 
     private String addBackticks(String aliasName) {
@@ -30,4 +36,8 @@ private String addBackticks(String aliasName) {
         }
         return "`" + aliasName + "`";
     }
+
+    public Map<String, String> getFieldAliasReplacedMap() {
+        return fieldAliasReplacedMap;
+    }
 }
diff --git a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java
@@ -487,7 +487,7 @@ public static String replaceAliasFieldName(String sql, Map<String, String> field
     }
 
     public static String replaceAliasWithBackticks(String sql) {
-          Select selectStatement = SqlSelectHelper.getSelect(sql);
+        Select selectStatement = SqlSelectHelper.getSelect(sql);
         if (!(selectStatement instanceof PlainSelect)) {
             return sql;
         }
@@ -496,6 +496,26 @@ public static String replaceAliasWithBackticks(String sql) {
         for (SelectItem selectItem : plainSelect.getSelectItems()) {
             selectItem.accept(visitor);
         }
+        // Replace `order by` and `group by`
+        // Get the map of field aliases that have been replaced
+        Map<String, String> aliasReplacedMap = visitor.getFieldAliasReplacedMap();
+
+        // If no aliases have been replaced, return the original SQL statement as a string
+        if (aliasReplacedMap.isEmpty()) {
+            return selectStatement.toString();
+        }
+        // Order by elements
+        List<OrderByElement> orderByElements = selectStatement.getOrderByElements();
+        if (!CollectionUtils.isEmpty(orderByElements)) {
+            for (OrderByElement orderByElement : orderByElements) {
+                orderByElement.accept(new OrderByReplaceVisitor(aliasReplacedMap, true));
+            }
+        }
+        // Group by elements
+        GroupByElement groupByElement = plainSelect.getGroupBy();
+        if (Objects.nonNull(groupByElement)) {
+            groupByElement.accept(new GroupByReplaceVisitor(aliasReplacedMap, true));
+        }
         return selectStatement.toString();
     }
 

diff --git a/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelperTest.java b/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelperTest.java
@@ -309,7 +309,7 @@ void testReplaceAliasWithBackticks() {
         String replaceSql = SqlReplaceHelper.replaceAliasWithBackticks(sql);
         System.out.println(replaceSql);
         Assert.assertEquals("SELECT 部门, SUM(访问次数) AS `总访问次数` FROM 超音数 WHERE "
-                + "datediff('day', 数据日期, '2023-09-05') <= 3 GROUP BY 部门 ORDER BY 总访问次数 DESC LIMIT 10",
+                + "datediff('day', 数据日期, '2023-09-05') <= 3 GROUP BY 部门 ORDER BY `总访问次数` DESC LIMIT 10",
                 replaceSql);
 
         sql = "select 部门, sum(访问次数) as 访问次数 from 超音数 where "
@@ -318,7 +318,16 @@ void testReplaceAliasWithBackticks() {
         replaceSql = SqlReplaceHelper.replaceAliasWithBackticks(sql);
         System.out.println(replaceSql);
         Assert.assertEquals("SELECT 部门, sum(访问次数) AS `访问次数` FROM 超音数 WHERE (datediff('day', 数据日期, "
-                + "'2023-09-05') <= 3) AND 数据日期 = '2023-10-10' GROUP BY 部门 ORDER BY 访问次数 DESC LIMIT 10",
+                + "'2023-09-05') <= 3) AND 数据日期 = '2023-10-10' GROUP BY 部门 ORDER BY `访问次数` DESC LIMIT 10",
+                replaceSql);
+
+        sql = "select 部门, sum(访问次数) as 访问次数, count(部门) as 部门数, count(部门) as 部门数2 from 超音数 where "
+                + "(datediff('day', 数据日期, '2023-09-05') <= 3) and 数据日期 = '2023-10-10' "
+                + "group by 部门, 部门数, 部门数2  order by 访问次数 desc limit 10";
+        replaceSql = SqlReplaceHelper.replaceAliasWithBackticks(sql);
+        System.out.println(replaceSql);
+        Assert.assertEquals("SELECT 部门, sum(访问次数) AS `访问次数`, count(部门) AS `部门数`, count(部门) AS `部门数2` FROM 超音数 WHERE (datediff('day', 数据日期, "
+                + "'2023-09-05') <= 3) AND 数据日期 = '2023-10-10' GROUP BY 部门, `部门数`, `部门数2` ORDER BY `访问次数` DESC LIMIT 10",
                 replaceSql);
     }