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

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());
+            }
+        }
+    }
+}

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);
+        }
+    }
+}

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

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

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);
     }

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

common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlAsHelperTest.java

@@ -0,0 +1,42 @@
+package com.tencent.supersonic.common.jsqlparser;
+
+import org.junit.Assert;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+class SqlAsHelperTest {
+
+    @Test
+    void testReplaceAggField() {
+        String sql =
+                "WITH SalesData AS (\n"
+                        + "    SELECT \n"
+                        + "        SalesID,\n"
+                        + "        ProductID,\n"
+                        + "        Quantity,\n"
+                        + "        Price,\n"
+                        + "        (Quantity * Price) AS TotalSales\n"
+                        + "    FROM \n"
+                        + "        Sales\n"
+                        + ")\n"
+                        + "SELECT \n"
+                        + "    ProductID,\n"
+                        + "    SUM(TotalSales) AS TotalRevenue,\n"
+                        + "    COUNT(SalesID) AS NumberOfSales\n"
+                        + "FROM \n"
+                        + "    SalesData\n"
+                        + "WHERE \n"
+                        + "    Quantity > 10\n"
+                        + "GROUP BY \n"
+                        + "    ProductID\n"
+                        + "HAVING \n"
+                        + "    SUM(TotalSales) > 1000\n"
+                        + "ORDER BY \n"
+                        + "    TotalRevenue DESC";
+        List<String> asFields = SqlAsHelper.getAsFields(sql);
+        Assert.assertTrue(asFields.contains("NumberOfSales"));
+        Assert.assertTrue(asFields.contains("TotalRevenue"));
+        Assert.assertTrue(asFields.contains("TotalSales"));
+    }
+}

common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelperTest.java

@@ -14,33 +14,6 @@ import java.util.Set;
 /** SqlParserReplaceHelperTest */
 class SqlReplaceHelperTest {
 
-    @Test
-    void testReplaceSelectField() {
-
-        String sql =
-                "SELECT 维度1,sum(播放量) FROM 数据库 "
-                        + "WHERE (歌手名 = '张三') AND 数据日期 = '2023-11-17' GROUP BY 维度1";
-        Map<String, String> fieldMap = new HashMap<>();
-        fieldMap.put("播放量", "播放量1");
-        sql = SqlReplaceHelper.replaceSelectFields(sql, fieldMap);
-        System.out.println(sql);
-        Assert.assertEquals(
-                "SELECT 维度1, sum(播放量1) FROM 数据库 "
-                        + "WHERE (歌手名 = '张三') AND 数据日期 = '2023-11-17' GROUP BY 维度1",
-                sql);
-
-        sql =
-                "SELECT 维度1,播放量 FROM 数据库 "
-                        + "WHERE (歌手名 = '张三') AND 数据日期 = '2023-11-17' GROUP BY 维度1";
-        fieldMap = new HashMap<>();
-        fieldMap.put("播放量", "播放量1");
-        sql = SqlReplaceHelper.replaceSelectFields(sql, fieldMap);
-        System.out.println(sql);
-        Assert.assertEquals(
-                "SELECT 维度1, 播放量1 FROM 数据库 WHERE (歌手名 = '张三') AND 数据日期 = '2023-11-17' GROUP BY 维度1",
-                sql);
-    }
-
     @Test
     void testReplaceAggField() {
         String sql =

headless/chat/src/main/java/com/tencent/supersonic/headless/chat/corrector/SchemaCorrector.java

@@ -2,6 +2,7 @@ package com.tencent.supersonic.headless.chat.corrector;
 
 import com.tencent.supersonic.common.jsqlparser.AggregateEnum;
 import com.tencent.supersonic.common.jsqlparser.FieldExpression;
+import com.tencent.supersonic.common.jsqlparser.SqlAsHelper;
 import com.tencent.supersonic.common.jsqlparser.SqlRemoveHelper;
 import com.tencent.supersonic.common.jsqlparser.SqlReplaceHelper;
 import com.tencent.supersonic.common.jsqlparser.SqlSelectHelper;
@@ -38,8 +39,6 @@ public class SchemaCorrector extends BaseSemanticCorrector {
 
         correctAggFunction(semanticParseInfo);
 
-        replaceAlias(semanticParseInfo);
-
         updateFieldNameByLinkingValue(semanticParseInfo);
 
         updateFieldValueByLinkingValue(semanticParseInfo);
@@ -62,17 +61,16 @@ public class SchemaCorrector extends BaseSemanticCorrector {
         sqlInfo.setCorrectedS2SQL(sql);
     }
 
-    private void replaceAlias(SemanticParseInfo semanticParseInfo) {
-        SqlInfo sqlInfo = semanticParseInfo.getSqlInfo();
-        String replaceAlias = SqlReplaceHelper.replaceAlias(sqlInfo.getCorrectedS2SQL());
-        sqlInfo.setCorrectedS2SQL(replaceAlias);
-    }
-
     private void correctFieldName(
             ChatQueryContext chatQueryContext, SemanticParseInfo semanticParseInfo) {
         Map<String, String> fieldNameMap =
                 getFieldNameMap(chatQueryContext, semanticParseInfo.getDataSetId());
+        // add as fieldName
         SqlInfo sqlInfo = semanticParseInfo.getSqlInfo();
+        List<String> asFields = SqlAsHelper.getAsFields(sqlInfo.getCorrectedS2SQL());
+        for (String asField : asFields) {
+            fieldNameMap.put(asField, asField);
+        }
         String sql = SqlReplaceHelper.replaceFields(sqlInfo.getCorrectedS2SQL(), fieldNameMap);
         sqlInfo.setCorrectedS2SQL(sql);
     }

headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/TagQueryServiceImpl.java

@@ -43,8 +43,8 @@ public class TagQueryServiceImpl implements TagQueryService {
     @Value("${s2.item.value.date.format:yyyy-MM-dd}")
     private String itemValueDateFormat;
 
-    //因有些数据库(如Postgresql)不支持列名大小写，所以列名统一使用小写，以兼容更多数据库
+    // 因有些数据库(如Postgresql)不支持列名大小写，所以列名统一使用小写，以兼容更多数据库
-    //private final String tagValueAlias = "internalTagCount";
+    // private final String tagValueAlias = "internalTagCount";
     private final String tagValueAlias = "internal_tag_count";
     private final String maxDateAlias = "internal_max_date";
     private final TagMetaService tagMetaService;

headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/QueryUtils.java

@@ -116,6 +116,10 @@ public class QueryUtils {
             column.setDataFormatType(metricRespMap.get(nameEn).getDataFormatType());
             column.setDataFormat(metricRespMap.get(nameEn).getDataFormat());
         }
+        // set name by NameEn
+        if (StringUtils.isBlank(column.getName()) && StringUtils.isNotBlank(column.getNameEn())) {
+            column.setName(column.getNameEn());
+        }
     }
 
     private boolean isNumberType(String type) {

launchers/standalone/src/test/java/com/tencent/supersonic/headless/QueryBySqlTest.java

@@ -31,23 +31,23 @@ public class QueryBySqlTest extends BaseTest {
     @Test
     public void testSumQuery() throws Exception {
         SemanticQueryResp semanticQueryResp =
-                queryBySql("SELECT SUM(访问次数) AS 访问次数 FROM 超音数PVUV统计 ");
+                queryBySql("SELECT SUM(访问次数) AS 总访问次数 FROM 超音数PVUV统计 ");
 
         assertEquals(1, semanticQueryResp.getColumns().size());
         QueryColumn queryColumn = semanticQueryResp.getColumns().get(0);
-        assertEquals("访问次数", queryColumn.getName());
+        assertEquals("总访问次数", queryColumn.getName());
         assertEquals(1, semanticQueryResp.getResultList().size());
     }
 
     @Test
     public void testGroupByQuery() throws Exception {
         SemanticQueryResp result =
-                queryBySql("SELECT 部门, SUM(访问次数) AS 访问次数 FROM 超音数PVUV统计  GROUP BY 部门 ");
+                queryBySql("SELECT 部门, SUM(访问次数) AS 总访问次数 FROM 超音数PVUV统计  GROUP BY 部门 ");
         assertEquals(2, result.getColumns().size());
         QueryColumn firstColumn = result.getColumns().get(0);
         QueryColumn secondColumn = result.getColumns().get(1);
         assertEquals("部门", firstColumn.getName());
-        assertEquals("访问次数", secondColumn.getName());
+        assertEquals("总访问次数", secondColumn.getName());
         assertEquals(4, result.getResultList().size());
     }
 
@@ -55,12 +55,12 @@ public class QueryBySqlTest extends BaseTest {
     public void testFilterQuery() throws Exception {
         SemanticQueryResp result =
                 queryBySql(
-                        "SELECT 部门, SUM(访问次数) AS 访问次数 FROM 超音数PVUV统计 WHERE 部门 ='HR' GROUP BY 部门 ");
+                        "SELECT 部门, SUM(访问次数) AS 总访问次数 FROM 超音数PVUV统计 WHERE 部门 ='HR' GROUP BY 部门 ");
         assertEquals(2, result.getColumns().size());
         QueryColumn firstColumn = result.getColumns().get(0);
         QueryColumn secondColumn = result.getColumns().get(1);
         assertEquals("部门", firstColumn.getName());
-        assertEquals("访问次数", secondColumn.getName());
+        assertEquals("总访问次数", secondColumn.getName());
         assertEquals(1, result.getResultList().size());
         assertEquals("HR", result.getResultList().get(0).get("department").toString());
     }
@@ -69,11 +69,12 @@ public class QueryBySqlTest extends BaseTest {
     public void testDateSumQuery() throws Exception {
         String startDate = now().plusDays(-365).toString();
         String endDate = now().plusDays(0).toString();
-        String sql = "SELECT SUM(访问次数) AS 访问次数 FROM 超音数PVUV统计 WHERE 数据日期 >= '%s' AND 数据日期 <= '%s' ";
+        String sql =
+                "SELECT SUM(访问次数) AS 总访问次数 FROM 超音数PVUV统计 WHERE 数据日期 >= '%s' AND 数据日期 <= '%s' ";
         SemanticQueryResp semanticQueryResp = queryBySql(String.format(sql, startDate, endDate));
         assertEquals(1, semanticQueryResp.getColumns().size());
         QueryColumn queryColumn = semanticQueryResp.getColumns().get(0);
-        assertEquals("访问次数", queryColumn.getName());
+        assertEquals("总访问次数", queryColumn.getName());
         assertEquals(1, semanticQueryResp.getResultList().size());
     }