From 96e94b30960f1e506d6a3656d9124c02d34e6a0e Mon Sep 17 00:00:00 2001
From: jipeli <54889677+jipeli@users.noreply.github.com>
Date: Thu, 6 Jun 2024 17:00:34 +0800
Subject: [PATCH] (improvement)(Headless) fixed get group error (#1107)

---
 .../util/jsqlparser/SqlSelectHelper.java      | 23 +++++++++++++------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlSelectHelper.java b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlSelectHelper.java
index 14c8b2f4e..dcb538611 100644
--- a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlSelectHelper.java
+++ b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlSelectHelper.java
@@ -38,6 +38,7 @@ import net.sf.jsqlparser.statement.select.Distinct;
 import net.sf.jsqlparser.statement.select.GroupByElement;
 import net.sf.jsqlparser.statement.select.Join;
 import net.sf.jsqlparser.statement.select.LateralView;
+import net.sf.jsqlparser.statement.select.Limit;
 import net.sf.jsqlparser.statement.select.OrderByElement;
 import net.sf.jsqlparser.statement.select.ParenthesedSelect;
 import net.sf.jsqlparser.statement.select.PlainSelect;
@@ -46,7 +47,6 @@ import net.sf.jsqlparser.statement.select.SelectItem;
 import net.sf.jsqlparser.statement.select.SelectVisitorAdapter;
 import net.sf.jsqlparser.statement.select.SetOperationList;
 import net.sf.jsqlparser.statement.select.WithItem;
-import net.sf.jsqlparser.statement.select.Limit;
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.util.CollectionUtils;
 
@@ -139,12 +139,21 @@ public class SqlSelectHelper {
         return plainSelectList;
     }
 
-    public static void getSubPlainSelect(PlainSelect plainSelect, List<PlainSelect> plainSelectList) {
-        plainSelectList.add(plainSelect);
-        if (plainSelect.getFromItem() instanceof ParenthesedSelect) {
-            ParenthesedSelect parenthesedSelect = (ParenthesedSelect) plainSelect.getFromItem();
-            PlainSelect subPlainSelect = parenthesedSelect.getPlainSelect();
-            getSubPlainSelect(subPlainSelect, plainSelectList);
+    public static void getSubPlainSelect(Select select, List<PlainSelect> plainSelectList) {
+        if (select instanceof PlainSelect) {
+            PlainSelect plainSelect = (PlainSelect) select;
+            plainSelectList.add(plainSelect);
+            if (plainSelect.getFromItem() instanceof ParenthesedSelect) {
+                ParenthesedSelect parenthesedSelect = (ParenthesedSelect) plainSelect.getFromItem();
+                Select subSelect = parenthesedSelect.getSelect();
+                getSubPlainSelect(subSelect, plainSelectList);
+            }
+        }
+        if (select instanceof SetOperationList) {
+            SetOperationList setOperationList = (SetOperationList) select;
+            for (Select subSelect : setOperationList.getSelects()) {
+                getSubPlainSelect(subSelect, plainSelectList);
+            }
         }
     }