(improvement)(headless) fix SqlReplaceHelper replaceFields，improve subselect and supports join (#1210)

.gitignore

@@ -3,7 +3,6 @@ target/
 .vscode/
 logs/
 log/
-venv/
 *.DS_Store
 *.iml
 *.bin

common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java

@@ -2,6 +2,7 @@ package com.tencent.supersonic.common.jsqlparser;
 
 import com.tencent.supersonic.common.pojo.enums.AggOperatorEnum;
 import com.tencent.supersonic.common.util.StringUtil;
+
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -9,6 +10,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.UnaryOperator;
+
 import lombok.extern.slf4j.Slf4j;
 import net.sf.jsqlparser.JSQLParserException;
 import net.sf.jsqlparser.expression.Alias;
@@ -38,6 +40,7 @@ import net.sf.jsqlparser.statement.select.Select;
 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.FromItem;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.springframework.util.CollectionUtils;
@@ -124,9 +127,9 @@ public class SqlReplaceHelper {
         if (!(selectStatement instanceof PlainSelect)) {
             return sql;
         }
-        List<PlainSelect> plainSelectList = new ArrayList<>();
+        //List<PlainSelect> plainSelectList = new ArrayList<>();
-        plainSelectList.add((PlainSelect) selectStatement);
+        //plainSelectList.add((PlainSelect) selectStatement);
-        List<PlainSelect> plainSelects = SqlSelectHelper.getPlainSelects(plainSelectList);
+        List<PlainSelect> plainSelects = SqlSelectHelper.getPlainSelect(selectStatement);
         for (PlainSelect plainSelect : plainSelects) {
             Expression where = plainSelect.getWhere();
             FieldlValueReplaceVisitor visitor = new FieldlValueReplaceVisitor(exactReplace, filedNameToValueMap);
@@ -155,6 +158,28 @@ public class SqlReplaceHelper {
         return selectStatement.toString();
     }
 
+    public static void getFromSelect(FromItem fromItem, List<PlainSelect> plainSelectList) {
+        if (!(fromItem instanceof ParenthesedSelect)) {
+            return;
+        }
+        ParenthesedSelect parenthesedSelect = (ParenthesedSelect) fromItem;
+        Select select = parenthesedSelect.getSelect();
+        if (select instanceof PlainSelect) {
+            PlainSelect plainSelect = (PlainSelect) select;
+            plainSelectList.add(plainSelect);
+            getFromSelect(plainSelect.getFromItem(), plainSelectList);
+        } else if (select instanceof SetOperationList) {
+            SetOperationList setOperationList = (SetOperationList) select;
+            if (!CollectionUtils.isEmpty(setOperationList.getSelects())) {
+                setOperationList.getSelects().forEach(subSelectBody -> {
+                    PlainSelect subPlainSelect = (PlainSelect) subSelectBody;
+                    plainSelectList.add(subPlainSelect);
+                    getFromSelect(subPlainSelect.getFromItem(), plainSelectList);
+                });
+            }
+        }
+    }
+
     public static String replaceFields(String sql, Map<String, String> fieldNameMap) {
         return replaceFields(sql, fieldNameMap, false);
     }
@@ -164,13 +189,19 @@ public class SqlReplaceHelper {
         List<PlainSelect> plainSelectList = SqlSelectHelper.getWithItem(selectStatement);
         //plainSelectList.add(selectStatement.getPlainSelect());
         if (selectStatement instanceof PlainSelect) {
-            plainSelectList.add((PlainSelect) selectStatement);
+            PlainSelect plainSelect = (PlainSelect) selectStatement;
+            plainSelectList.add(plainSelect);
+            getFromSelect(plainSelect.getFromItem(), plainSelectList);
+            //plainSelectList.add((PlainSelect) selectStatement);
         } else if (selectStatement instanceof SetOperationList) {
             SetOperationList setOperationList = (SetOperationList) selectStatement;
             if (!CollectionUtils.isEmpty(setOperationList.getSelects())) {
                 setOperationList.getSelects().forEach(subSelectBody -> {
+                    //PlainSelect subPlainSelect = (PlainSelect) subSelectBody;
+                    //plainSelectList.add(subPlainSelect);
                     PlainSelect subPlainSelect = (PlainSelect) subSelectBody;
                     plainSelectList.add(subPlainSelect);
+                    getFromSelect(subPlainSelect.getFromItem(), plainSelectList);
                 });
             }
             List<OrderByElement> orderByElements = setOperationList.getOrderByElements();
@@ -206,8 +237,19 @@ public class SqlReplaceHelper {
 
         if (plainSelect.getFromItem() instanceof ParenthesedSelect) {
             ParenthesedSelect parenthesedSelect = (ParenthesedSelect) plainSelect.getFromItem();
-            PlainSelect subPlainSelect = parenthesedSelect.getPlainSelect();
+            Select select = parenthesedSelect.getSelect();
+            if (select instanceof PlainSelect) {
+                PlainSelect subPlainSelect = (PlainSelect) select;
                 replaceFieldsInPlainOneSelect(fieldNameMap, exactReplace, subPlainSelect);
+            } else if (select instanceof SetOperationList) {
+                SetOperationList setOperationList = (SetOperationList) select;
+                if (!CollectionUtils.isEmpty(setOperationList.getSelects())) {
+                    setOperationList.getSelects().forEach(subSelectBody -> {
+                        PlainSelect subPlainSelect = (PlainSelect) subSelectBody;
+                        replaceFieldsInPlainOneSelect(fieldNameMap, exactReplace, subPlainSelect);
+                    });
+                }
+            }
         }
 
         //3. replace oder by fields
@@ -238,8 +280,9 @@ public class SqlReplaceHelper {
                 if (!(join.getRightItem() instanceof ParenthesedSelect)) {
                     continue;
                 }
+                ParenthesedSelect parenthesedSelect = (ParenthesedSelect) join.getRightItem();
                 List<PlainSelect> plainSelectList = new ArrayList<>();
-                plainSelectList.add((PlainSelect) join.getRightItem());
+                plainSelectList.add(parenthesedSelect.getPlainSelect());
                 List<PlainSelect> subPlainSelects = SqlSelectHelper.getPlainSelects(plainSelectList);
                 for (PlainSelect subPlainSelect : subPlainSelects) {
                     replaceFieldsInPlainOneSelect(fieldNameMap, exactReplace, subPlainSelect);
@@ -435,6 +478,16 @@ public class SqlReplaceHelper {
             PlainSelect subPlainSelect = parenthesedSelect.getPlainSelect();
             replaceSingleTable(subPlainSelect, tableName);
         }
+        List<Join> joinList = plainSelect.getJoins();
+        if (CollectionUtils.isEmpty(joinList)) {
+            return;
+        }
+        for (Join join : joinList) {
+            if (join.getFromItem() instanceof ParenthesedSelect) {
+                ParenthesedSelect parenthesedSelect = (ParenthesedSelect) join.getFromItem();
+                replaceSingleTable(parenthesedSelect.getPlainSelect(), tableName);
+            }
+        }
     }
 
     public static void replaceSingleTable(PlainSelect plainSelect, String tableName) {
@@ -648,6 +701,10 @@ public class SqlReplaceHelper {
     public static String dealAliasToOrderBy(String querySql) {
         Select selectStatement = SqlSelectHelper.getSelect(querySql);
         List<PlainSelect> plainSelectList = new ArrayList<>();
+        //List<PlainSelect> withPlainSelectList = SqlSelectHelper.getWithItem(selectStatement);
+        //if (!CollectionUtils.isEmpty(withPlainSelectList)) {
+        //    plainSelectList.addAll(withPlainSelectList);
+        //}
         if (selectStatement instanceof PlainSelect) {
             plainSelectList.add((PlainSelect) selectStatement);
         } else if (selectStatement instanceof SetOperationList) {

common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java

@@ -116,8 +116,7 @@ public class SqlSelectHelper {
         return result;
     }
 
-    public static List<PlainSelect> getPlainSelect(String sql) {
+    public static List<PlainSelect> getPlainSelect(Select selectStatement) {
-        Select selectStatement = getSelect(sql);
         if (selectStatement == null) {
             return null;
         }
@@ -139,6 +138,11 @@ public class SqlSelectHelper {
         return plainSelectList;
     }
 
+    public static List<PlainSelect> getPlainSelect(String sql) {
+        Select selectStatement = getSelect(sql);
+        return getPlainSelect(selectStatement);
+    }
+
     public static Boolean hasSubSelect(String sql) {
         Select selectStatement = getSelect(sql);
         if (selectStatement == null) {
@@ -162,6 +166,16 @@ public class SqlSelectHelper {
                 Select subSelect = parenthesedSelect.getSelect();
                 getSubPlainSelect(subSelect, plainSelectList);
             }
+            List<Join> joinList = plainSelect.getJoins();
+            if (CollectionUtils.isEmpty(joinList)) {
+                return;
+            }
+            for (Join join : joinList) {
+                if (join.getRightItem() instanceof ParenthesedSelect) {
+                    ParenthesedSelect parenthesedSelect = (ParenthesedSelect) join.getRightItem();
+                    plainSelectList.add(parenthesedSelect.getPlainSelect());
+                }
+            }
         }
         if (select instanceof SetOperationList) {
             SetOperationList setOperationList = (SetOperationList) select;
@@ -313,11 +327,29 @@ public class SqlSelectHelper {
             }
             if (plainSelect.getFromItem() instanceof ParenthesedSelect) {
                 ParenthesedSelect parenthesedSelect = (ParenthesedSelect) plainSelect.getFromItem();
+                Select subSelect = parenthesedSelect.getSelect();
+                if (subSelect instanceof PlainSelect) {
                     PlainSelect subPlainSelect = parenthesedSelect.getPlainSelect();
                     Expression subWhere = subPlainSelect.getWhere();
                     if (Objects.nonNull(subWhere)) {
                         subWhere.accept(new FieldAndValueAcquireVisitor(result));
                     }
+                } else if (subSelect instanceof ParenthesedSelect) {
+                    ParenthesedSelect subParenthesedSelect = (ParenthesedSelect) subSelect;
+                    Expression subWhere = subParenthesedSelect.getPlainSelect().getWhere();
+                    if (Objects.nonNull(subWhere)) {
+                        subWhere.accept(new FieldAndValueAcquireVisitor(result));
+                    }
+                } else if (subSelect instanceof SetOperationList) {
+                    SetOperationList setOperationList = (SetOperationList) subSelect;
+                    List<Select> selectList = setOperationList.getSelects();
+                    for (Select select : selectList) {
+                        Expression subWhere = select.getPlainSelect().getWhere();
+                        if (Objects.nonNull(subWhere)) {
+                            subWhere.accept(new FieldAndValueAcquireVisitor(result));
+                        }
+                    }
+                }
             }
         }
         return new ArrayList<>(result);
@@ -767,7 +799,9 @@ public class SqlSelectHelper {
         if (plainSelect.getFromItem() instanceof ParenthesedSelect) {
             ParenthesedSelect parenthesedSelect = (ParenthesedSelect) plainSelect.getFromItem();
             getFieldsWithSubQuery(parenthesedSelect.getPlainSelect(), fields);
-            if (!CollectionUtils.isEmpty(plainSelect.getJoins())) {
+            if (CollectionUtils.isEmpty(plainSelect.getJoins())) {
+                return;
+            }
             for (Join join : plainSelect.getJoins()) {
                 if (join.getRightItem() instanceof ParenthesedSelect) {
                     getFieldsWithSubQuery(((ParenthesedSelect) join.getRightItem()).getPlainSelect(), fields);
@@ -778,6 +812,5 @@ public class SqlSelectHelper {
             }
         }
     }
-    }
 }