(improvement)(headless) Optimize jsqlparser code and add support for complex queries like WITH in the corrector (#1730)

2025-12-13 04:57:28 +00:00 · 2024-09-28 23:26:39 +08:00
parent e351a1c541
commit 3a11ccb6e9
11 changed files with 219 additions and 177 deletions
--- a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/FunctionAliasVisitor.java
+++ b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/FunctionAliasVisitor.java
@@ -0,0 +1,26 @@
+package com.tencent.supersonic.common.jsqlparser;
+
+import net.sf.jsqlparser.expression.Function;
+import net.sf.jsqlparser.statement.select.SelectItem;
+import net.sf.jsqlparser.statement.select.SelectItemVisitorAdapter;
+
+import java.util.Objects;
+import java.util.Set;
+
+public class FunctionAliasVisitor extends SelectItemVisitorAdapter {
+
+    private Set<String> fieldNames;
+
+    public FunctionAliasVisitor(Set<String> fieldNames) {
+        this.fieldNames = fieldNames;
+    }
+
+    @Override
+    public void visit(SelectItem selectExpressionItem) {
+        if (selectExpressionItem.getExpression() instanceof Function) {
+            if (Objects.nonNull(selectExpressionItem.getAlias())) {
+                fieldNames.add(selectExpressionItem.getAlias().getName());
+            }
+        }
+    }
+}
--- a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlAsHelper.java
+++ b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlAsHelper.java
@@ -0,0 +1,61 @@
+package com.tencent.supersonic.common.jsqlparser;
+
+import lombok.extern.slf4j.Slf4j;
+import net.sf.jsqlparser.expression.Alias;
+import net.sf.jsqlparser.statement.select.PlainSelect;
+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.WithItem;
+import org.springframework.util.CollectionUtils;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+/** Sql Parser Select Helper */
+@Slf4j
+public class SqlAsHelper {
+
+    public static List<String> getAsFields(String sql) {
+        List<PlainSelect> plainSelectList = SqlSelectHelper.getPlainSelect(sql);
+        if (CollectionUtils.isEmpty(plainSelectList)) {
+            return new ArrayList<>();
+        }
+        Set<String> aliases = new HashSet<>();
+        for (PlainSelect plainSelect : plainSelectList) {
+            if (plainSelect instanceof Select) {
+                Select select = plainSelect;
+                Select selectBody = select.getSelectBody();
+                selectBody.accept(
+                        new SelectVisitorAdapter() {
+                            @Override
+                            public void visit(PlainSelect plainSelect) {
+                                extractAliasesFromSelect(plainSelect, aliases);
+                            }
+
+                            @Override
+                            public void visit(WithItem withItem) {
+                                withItem.getSelectBody().accept(this);
+                            }
+                        });
+            }
+        }
+        return new ArrayList<>(aliases);
+    }
+
+    private static void extractAliasesFromSelect(PlainSelect plainSelect, Set<String> aliases) {
+        // Extract aliases from SELECT items
+        for (SelectItem selectItem : plainSelect.getSelectItems()) {
+            Alias alias = selectItem.getAlias();
+            if (alias != null) {
+                aliases.add(alias.getName());
+            }
+        }
+        FunctionAliasVisitor visitor = new FunctionAliasVisitor(aliases);
+        for (SelectItem selectItem : plainSelect.getSelectItems()) {
+            selectItem.accept(visitor);
+        }
+    }
+}
--- a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlRemoveHelper.java
+++ b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlRemoveHelper.java
@@ -154,68 +154,79 @@ public class SqlRemoveHelper {

    private static void removeExpressionWithConstant(
            Expression expression, Set<String> removeFieldNames) {
-        if (expression instanceof EqualsTo
-                || expression instanceof GreaterThanEquals
-                || expression instanceof GreaterThan
-                || expression instanceof MinorThanEquals
-                || expression instanceof MinorThan) {
-            ComparisonOperator comparisonOperator = (ComparisonOperator) expression;
-            String columnName =
-                    SqlSelectHelper.getColumnName(
-                            comparisonOperator.getLeftExpression(),
-                            comparisonOperator.getRightExpression());
-            if (!removeFieldNames.contains(columnName)) {
-                return;
-            }
-            String constant = getConstant(expression);
-            try {
-                ComparisonOperator constantExpression =
-                        (ComparisonOperator) CCJSqlParserUtil.parseCondExpression(constant);
-                comparisonOperator.setLeftExpression(constantExpression.getLeftExpression());
-                comparisonOperator.setRightExpression(constantExpression.getRightExpression());
-                comparisonOperator.setASTNode(constantExpression.getASTNode());
-            } catch (JSQLParserException e) {
-                log.error("JSQLParserException", e);
+        try {
+            if (expression instanceof ComparisonOperator) {
+                handleComparisonOperator((ComparisonOperator) expression, removeFieldNames);
+            } else if (expression instanceof InExpression) {
+                handleInExpression((InExpression) expression, removeFieldNames);
+            } else if (expression instanceof LikeExpression) {
+                handleLikeExpression((LikeExpression) expression, removeFieldNames);
            }
+        } catch (JSQLParserException e) {
+            log.error("JSQLParserException", e);
        }
-        if (expression instanceof InExpression) {
-            InExpression inExpression = (InExpression) expression;
-            String columnName =
-                    SqlSelectHelper.getColumnName(
-                            inExpression.getLeftExpression(), inExpression.getRightExpression());
-            if (!removeFieldNames.contains(columnName)) {
-                return;
-            }
-            try {
-                InExpression constantExpression =
-                        (InExpression)
-                                CCJSqlParserUtil.parseCondExpression(JsqlConstants.IN_CONSTANT);
-                inExpression.setLeftExpression(constantExpression.getLeftExpression());
-                inExpression.setRightExpression(constantExpression.getRightExpression());
-                inExpression.setASTNode(constantExpression.getASTNode());
-            } catch (JSQLParserException e) {
-                log.error("JSQLParserException", e);
-            }
+    }
+
+    private static void handleComparisonOperator(
+            ComparisonOperator comparisonOperator, Set<String> removeFieldNames)
+            throws JSQLParserException {
+        String columnName =
+                SqlSelectHelper.getColumnName(
+                        comparisonOperator.getLeftExpression(),
+                        comparisonOperator.getRightExpression());
+        if (!removeFieldNames.contains(columnName)) {
+            return;
        }
-        if (expression instanceof LikeExpression) {
-            LikeExpression likeExpression = (LikeExpression) expression;
-            String columnName =
-                    SqlSelectHelper.getColumnName(
-                            likeExpression.getLeftExpression(),
-                            likeExpression.getRightExpression());
-            if (!removeFieldNames.contains(columnName)) {
-                return;
-            }
-            try {
-                LikeExpression constantExpression =
-                        (LikeExpression)
-                                CCJSqlParserUtil.parseCondExpression(JsqlConstants.LIKE_CONSTANT);
-                likeExpression.setLeftExpression(constantExpression.getLeftExpression());
-                likeExpression.setRightExpression(constantExpression.getRightExpression());
-            } catch (JSQLParserException e) {
-                log.error("JSQLParserException", e);
-            }
+        String constant = getConstant(comparisonOperator);
+        ComparisonOperator constantExpression =
+                (ComparisonOperator) CCJSqlParserUtil.parseCondExpression(constant);
+        updateComparisonOperator(comparisonOperator, constantExpression);
+    }
+
+    private static void handleInExpression(InExpression inExpression, Set<String> removeFieldNames)
+            throws JSQLParserException {
+        String columnName =
+                SqlSelectHelper.getColumnName(
+                        inExpression.getLeftExpression(), inExpression.getRightExpression());
+        if (!removeFieldNames.contains(columnName)) {
+            return;
        }
+        InExpression constantExpression =
+                (InExpression) CCJSqlParserUtil.parseCondExpression(JsqlConstants.IN_CONSTANT);
+        updateInExpression(inExpression, constantExpression);
+    }
+
+    private static void handleLikeExpression(
+            LikeExpression likeExpression, Set<String> removeFieldNames)
+            throws JSQLParserException {
+        String columnName =
+                SqlSelectHelper.getColumnName(
+                        likeExpression.getLeftExpression(), likeExpression.getRightExpression());
+        if (!removeFieldNames.contains(columnName)) {
+            return;
+        }
+        LikeExpression constantExpression =
+                (LikeExpression) CCJSqlParserUtil.parseCondExpression(JsqlConstants.LIKE_CONSTANT);
+        updateLikeExpression(likeExpression, constantExpression);
+    }
+
+    private static void updateComparisonOperator(
+            ComparisonOperator original, ComparisonOperator constantExpression) {
+        original.setLeftExpression(constantExpression.getLeftExpression());
+        original.setRightExpression(constantExpression.getRightExpression());
+        original.setASTNode(constantExpression.getASTNode());
+    }
+
+    private static void updateInExpression(InExpression original, InExpression constantExpression) {
+        original.setLeftExpression(constantExpression.getLeftExpression());
+        original.setRightExpression(constantExpression.getRightExpression());
+        original.setASTNode(constantExpression.getASTNode());
+    }
+
+    private static void updateLikeExpression(
+            LikeExpression original, LikeExpression constantExpression) {
+        original.setLeftExpression(constantExpression.getLeftExpression());
+        original.setRightExpression(constantExpression.getRightExpression());
    }

    public static String removeHavingCondition(String sql, Set<String> removeFieldNames) {
--- a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java
+++ b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java
@@ -4,7 +4,6 @@ import com.tencent.supersonic.common.pojo.enums.AggOperatorEnum;
 import com.tencent.supersonic.common.util.StringUtil;
 import lombok.extern.slf4j.Slf4j;
 import net.sf.jsqlparser.JSQLParserException;
-import net.sf.jsqlparser.expression.Alias;
 import net.sf.jsqlparser.expression.Expression;
 import net.sf.jsqlparser.expression.Function;
 import net.sf.jsqlparser.expression.LongValue;
@@ -47,53 +46,6 @@ import java.util.function.UnaryOperator;
 /** Sql Parser replace Helper */
@Slf4j
 public class SqlReplaceHelper {
-
-    public static String replaceSelectFields(String sql, Map<String, String> fieldNameMap) {
-        Select selectStatement = SqlSelectHelper.getSelect(sql);
-        if (!(selectStatement instanceof PlainSelect)) {
-            return sql;
-        }
-        ((PlainSelect) selectStatement)
-                .getSelectItems().stream()
-                        .forEach(
-                                o -> {
-                                    SelectItem selectExpressionItem = (SelectItem) o;
-                                    String alias = "";
-                                    if (selectExpressionItem.getExpression() instanceof Function) {
-                                        Function function =
-                                                (Function) selectExpressionItem.getExpression();
-                                        Column column =
-                                                (Column)
-                                                        function.getParameters()
-                                                                .getExpressions()
-                                                                .get(0);
-                                        if (fieldNameMap.containsKey(column.getColumnName())) {
-                                            String value = fieldNameMap.get(column.getColumnName());
-                                            alias = value;
-                                            function.withParameters(new Column(value));
-                                        }
-                                    }
-                                    if (selectExpressionItem.getExpression() instanceof Column) {
-                                        Column column =
-                                                (Column) selectExpressionItem.getExpression();
-                                        String columnName = column.getColumnName();
-                                        if (fieldNameMap.containsKey(columnName)) {
-                                            String value = fieldNameMap.get(columnName);
-                                            alias = value;
-                                            if (StringUtils.isNotBlank(value)) {
-                                                selectExpressionItem.setExpression(
-                                                        new Column(value));
-                                            }
-                                        }
-                                    }
-                                    if (Objects.nonNull(selectExpressionItem.getAlias())
-                                            && StringUtils.isNotBlank(alias)) {
-                                        selectExpressionItem.getAlias().setName(alias);
-                                    }
-                                });
-        return selectStatement.toString();
-    }
-
    public static String replaceAggFields(
            String sql, Map<String, Pair<String, String>> fieldNameToAggMap) {
        Select selectStatement = SqlSelectHelper.getSelect(sql);
@@ -257,7 +209,6 @@ public class SqlReplaceHelper {
        // 2. replace select fields
        for (SelectItem selectItem : plainSelect.getSelectItems()) {
            selectItem.accept(visitor);
-            replaceAsName(fieldNameMap, selectItem);
        }

        if (plainSelect.getFromItem() instanceof ParenthesedSelect) {
@@ -323,19 +274,6 @@ public class SqlReplaceHelper {
        }
    }

-    private static void replaceAsName(Map<String, String> fieldNameMap, SelectItem selectItem) {
-
-        Alias alias = selectItem.getAlias();
-        if (Objects.isNull(alias)) {
-            return;
-        }
-        String aliasName = alias.getName();
-        String replaceFieldName = fieldNameMap.get(aliasName);
-        if (StringUtils.isNotBlank(replaceFieldName)) {
-            alias.setName(replaceFieldName);
-        }
-    }
-
    public static String replaceFunction(String sql, Map<String, String> functionMap) {
        return replaceFunction(sql, functionMap, null);
    }
--- a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java
+++ b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java
@@ -26,7 +26,6 @@ import net.sf.jsqlparser.parser.CCJSqlParserUtil;
 import net.sf.jsqlparser.schema.Column;
 import net.sf.jsqlparser.schema.Table;
 import net.sf.jsqlparser.statement.Statement;
-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;
@@ -547,17 +546,6 @@ public class SqlSelectHelper {
        return false;
    }

-    public static boolean hasDistinct(String sql) {
-        Select selectStatement = getSelect(sql);
-
-        if (!(selectStatement instanceof PlainSelect)) {
-            return false;
-        }
-        PlainSelect plainSelect = (PlainSelect) selectStatement;
-        Distinct distinct = plainSelect.getDistinct();
-        return Objects.nonNull(distinct);
-    }
-
    public static boolean isLogicExpression(Expression whereExpression) {
        return whereExpression instanceof AndExpression
                || (whereExpression instanceof OrExpression