From 5a163936a81a02874b492362ffdc22e108fc4de1 Mon Sep 17 00:00:00 2001
From: lexluo09 <39718951+lexluo09@users.noreply.github.com>
Date: Sun, 29 Sep 2024 11:28:31 +0800
Subject: [PATCH] (improvement)(headless) Determine if there are any 'GROUP BY'
 fields that require compatibility with self-queries in WITH clauses. (#1734)

---
 .../common/jsqlparser/SqlSelectHelper.java      | 17 +++++++++++------
 .../common/jsqlparser/SqlSelectHelperTest.java  |  9 +++++++++
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java
index 2e509671d..cc5305ea7 100644
--- a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java
+++ b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java
@@ -536,12 +536,17 @@ public class SqlSelectHelper {
         if (!(selectStatement instanceof PlainSelect)) {
             return false;
         }
-        PlainSelect plainSelect = (PlainSelect) selectStatement;
-        GroupByElement groupBy = plainSelect.getGroupBy();
-        if (Objects.nonNull(groupBy)) {
-            GroupByVisitor replaceVisitor = new GroupByVisitor();
-            groupBy.accept(replaceVisitor);
-            return replaceVisitor.isHasAggregateFunction();
+
+        List<PlainSelect> withItem = getWithItem(selectStatement);
+        withItem.add((PlainSelect) selectStatement);
+
+        for (PlainSelect plainSelect : withItem) {
+            GroupByElement groupBy = plainSelect.getGroupBy();
+            if (Objects.nonNull(groupBy)) {
+                GroupByVisitor replaceVisitor = new GroupByVisitor();
+                groupBy.accept(replaceVisitor);
+                return replaceVisitor.isHasAggregateFunction();
+            }
         }
         return false;
     }
diff --git a/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelperTest.java b/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelperTest.java
index ebf8fe1a7..44e841d6d 100644
--- a/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelperTest.java
+++ b/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelperTest.java
@@ -333,4 +333,13 @@ class SqlSelectHelperTest {
         selectFields = SqlSelectHelper.gePureSelectFields(sql);
         Assert.assertEquals(selectFields.size(), 2);
     }
+
+    @Test
+    void testHasGroupBy() {
+        String sql =
+                "WITH DepartmentVisits AS (SELECT 部门, SUM(访问次数) AS pv FROM 超音数数据集 WHERE 数据日期 >= '2024-08-29' "
+                        + "AND 数据日期 <= '2024-09-29' GROUP BY 部门) SELECT COUNT(*) FROM DepartmentVisits WHERE pv > 100";
+        Boolean hasGroupBy = SqlSelectHelper.hasGroupBy(sql);
+        Assert.assertEquals(hasGroupBy, true);
+    }
 }