From 5c70607851cbe8b222f1fd63ad5e3dac40854fff Mon Sep 17 00:00:00 2001 From: lexluo09 <39718951+lexluo09@users.noreply.github.com> Date: Wed, 30 Oct 2024 13:17:28 +0800 Subject: [PATCH] [improvement][common] Field replacement is performed using a recursive approach, and it supports field replacement with complex expressions (#1859) --- .../jsqlparser/FieldReplaceVisitor.java | 16 +++ .../common/jsqlparser/SqlReplaceHelper.java | 63 ++++------- .../common/jsqlparser/SqlSelectHelper.java | 104 +++++++++++++++++- .../jsqlparser/SqlReplaceFieldsTest.java | 17 +++ .../jsqlparser/SqlReplaceHelperTest.java | 1 + 5 files changed, 161 insertions(+), 40 deletions(-) diff --git a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/FieldReplaceVisitor.java b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/FieldReplaceVisitor.java index 87f50cb7e..ccd0a64da 100644 --- a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/FieldReplaceVisitor.java +++ b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/FieldReplaceVisitor.java @@ -2,9 +2,13 @@ package com.tencent.supersonic.common.jsqlparser; import com.tencent.supersonic.common.util.ContextUtils; import lombok.extern.slf4j.Slf4j; +import net.sf.jsqlparser.expression.AnalyticExpression; import net.sf.jsqlparser.expression.ExpressionVisitorAdapter; import net.sf.jsqlparser.expression.Function; +import net.sf.jsqlparser.expression.WindowDefinition; import net.sf.jsqlparser.schema.Column; +import net.sf.jsqlparser.statement.select.OrderByElement; +import org.springframework.util.CollectionUtils; import java.util.Map; @@ -34,4 +38,16 @@ public class FieldReplaceVisitor extends ExpressionVisitorAdapter { exactReplace.set(originalExactReplace); } } + + @Override + public void visit(AnalyticExpression expr) { + super.visit(expr); + WindowDefinition windowDefinition = expr.getWindowDefinition(); + if (windowDefinition != null + && !CollectionUtils.isEmpty(windowDefinition.getOrderByElements())) { + for (OrderByElement element : windowDefinition.getOrderByElements()) { + element.getExpression().accept(this); + } + } + } } diff --git a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java index a150a73bc..e9fc3dbcd 100644 --- a/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java +++ b/common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java @@ -154,37 +154,20 @@ public class SqlReplaceHelper { public static String replaceFields(String sql, Map fieldNameMap, boolean exactReplace) { Select selectStatement = SqlSelectHelper.getSelect(sql); - List plainSelectList = SqlSelectHelper.getWithItem(selectStatement); - if (selectStatement instanceof PlainSelect) { - PlainSelect plainSelect = (PlainSelect) selectStatement; - plainSelectList.add(plainSelect); - getFromSelect(plainSelect.getFromItem(), plainSelectList); - } 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); - getFromSelect(subPlainSelect.getFromItem(), plainSelectList); - }); + Set getAllSelect(Select selectStatement) { + Set selects) { + if (select == null) { + return; + } + if (select instanceof PlainSelect) { + PlainSelect plainSelect = (PlainSelect) select; + selects.add(plainSelect); + collectFromItemPlainSelects(plainSelect.getFromItem(), selects); + collectWithItemPlainSelects(plainSelect.getWithItemsList(), selects); + collectJoinsPlainSelects(plainSelect.getJoins(), selects); + collectNestedPlainSelects(plainSelect, selects); + } else if (select instanceof SetOperationList) { + SetOperationList setOperationList = (SetOperationList) select; + selects.add(setOperationList); + if (!CollectionUtils.isEmpty(setOperationList.getSelects())) { + for (Select subSelectBody : setOperationList.getSelects()) { + collectSelects(subSelectBody, selects); + } + } + } else if (select instanceof WithItem) { + WithItem withItem = (WithItem) select; + collectSelects(withItem.getSelect(), selects); + } else if (select instanceof ParenthesedSelect) { + ParenthesedSelect parenthesedSelect = (ParenthesedSelect) select; + collectSelects(parenthesedSelect.getPlainSelect(), selects); + } + } + + private static void collectJoinsPlainSelects(List joins, Set selects) { + if (fromItem instanceof ParenthesedSelect) { + ParenthesedSelect parenthesedSelect = (ParenthesedSelect) fromItem; + collectSelects(parenthesedSelect.getSelect(), selects); + } + } + + public static void collectWithItemPlainSelects(List withItemList, + Set selects) { + ExpressionVisitorAdapter expressionVisitor = new ExpressionVisitorAdapter() { + @Override + public void visit(Select subSelect) { + if (subSelect instanceof ParenthesedSelect) { + ParenthesedSelect parenthesedSelect = (ParenthesedSelect) subSelect; + if (parenthesedSelect.getSelect() instanceof PlainSelect) { + selects.add(parenthesedSelect.getPlainSelect()); + } + } + } + }; + + plainSelect.accept(new SelectVisitorAdapter() { + @Override + public void visit(PlainSelect plainSelect) { + Expression whereExpression = plainSelect.getWhere(); + if (whereExpression != null) { + whereExpression.accept(expressionVisitor); + } + Expression having = plainSelect.getHaving(); + if (Objects.nonNull(having)) { + having.accept(expressionVisitor); + } + List> selectItems = plainSelect.getSelectItems(); + if (!CollectionUtils.isEmpty(selectItems)) { + for (SelectItem selectItem : selectItems) { + selectItem.accept(expressionVisitor); + } + } + } + }); + } } diff --git a/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceFieldsTest.java b/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceFieldsTest.java index 4585ede1e..95281a21d 100644 --- a/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceFieldsTest.java +++ b/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceFieldsTest.java @@ -263,4 +263,21 @@ class SqlReplaceFieldsTest extends SqlReplaceHelperTest { + "SELECT * FROM daily_visits", replaceSql); } + @Test + void testReplaceFields18() { + + String replaceSql = "WITH\n" + " latest_data AS (\n" + " SELECT\n" + " 粉丝数,\n" + + " ROW_NUMBER() OVER (\n" + " ORDER BY\n" + " 数据日期 DESC\n" + + " ) AS __row_num__\n" + " FROM\n" + " 问答艺人数据集\n" + " WHERE\n" + + " (TME歌手ID = '1')\n" + " AND (\n" + " 数据日期 >= '2024-10-22'\n" + + " AND 数据日期 <= '2024-10-29'\n" + " )\n" + " )\n" + "SELECT\n" + + " AVG(__粉丝数__)\n" + "FROM\n" + " latest_data\n" + "WHERE\n" + + " __row_num__ = 1"; + replaceSql = SqlReplaceHelper.replaceFields(replaceSql, fieldToBizName); + + Assert.assertEquals("WITH latest_data AS (SELECT fans_cnt, ROW_NUMBER() OVER " + + "(ORDER BY sys_imp_date DESC) AS __row_num__ FROM 问答艺人数据集 WHERE (TME歌手ID = '1') " + + "AND (sys_imp_date >= '2024-10-22' AND sys_imp_date <= '2024-10-29')) SELECT AVG(__粉丝数__) " + + "FROM latest_data WHERE __row_num__ = 1", replaceSql); + } } diff --git a/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelperTest.java b/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelperTest.java index 66d51a6d5..a641b2539 100644 --- a/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelperTest.java +++ b/common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelperTest.java @@ -348,6 +348,7 @@ class SqlReplaceHelperTest { fieldToBizName.put("歌曲发布时间", "song_publis_date"); fieldToBizName.put("歌曲发布年份", "song_publis_year"); fieldToBizName.put("访问次数", "pv"); + fieldToBizName.put("粉丝数", "fans_cnt"); return fieldToBizName; }