From 11c2e0505b4a807a26f64fb600c49a11b6694166 Mon Sep 17 00:00:00 2001 From: jipeli <54889677+jipeli@users.noreply.github.com> Date: Fri, 26 Apr 2024 15:34:36 +0800 Subject: [PATCH] (improvement)(Headless) fix pg date function name corrector (#950) --- .../FunctionNameReplaceVisitor.java | 19 +++++++++- .../GroupByFunctionReplaceVisitor.java | 18 ++++++++- .../util/jsqlparser/SqlReplaceHelper.java | 15 ++++++-- .../core/adaptor/db/PostgresqlAdaptor.java | 37 +++++++++++++++++-- 4 files changed, 79 insertions(+), 10 deletions(-) diff --git a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/FunctionNameReplaceVisitor.java b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/FunctionNameReplaceVisitor.java index 76fc5bd25..d0cd3476e 100644 --- a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/FunctionNameReplaceVisitor.java +++ b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/FunctionNameReplaceVisitor.java @@ -1,24 +1,41 @@ package com.tencent.supersonic.common.util.jsqlparser; import java.util.Map; +import java.util.Objects; +import java.util.function.UnaryOperator; import lombok.extern.slf4j.Slf4j; import net.sf.jsqlparser.expression.ExpressionVisitorAdapter; import net.sf.jsqlparser.expression.Function; +import net.sf.jsqlparser.expression.operators.relational.ExpressionList; import org.apache.commons.lang3.StringUtils; @Slf4j public class FunctionNameReplaceVisitor extends ExpressionVisitorAdapter { private Map functionMap; + private Map functionCallMap; public FunctionNameReplaceVisitor(Map functionMap) { this.functionMap = functionMap; } + public FunctionNameReplaceVisitor(Map functionMap, Map functionCallMap) { + this.functionMap = functionMap; + this.functionCallMap = functionCallMap; + } + public void visit(Function function) { - String replaceFunctionName = functionMap.get(function.getName().toLowerCase()); + String functionName = function.getName().toLowerCase(); + String replaceFunctionName = functionMap.get(functionName); if (StringUtils.isNotBlank(replaceFunctionName)) { function.setName(replaceFunctionName); + if (Objects.nonNull(functionCallMap) && functionCallMap.containsKey(functionName)) { + Object ret = functionCallMap.get(functionName).apply(function.getParameters()); + if (Objects.nonNull(ret) && ret instanceof ExpressionList) { + ExpressionList expressionList = (ExpressionList) ret; + function.setParameters(expressionList); + } + } } } diff --git a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/GroupByFunctionReplaceVisitor.java b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/GroupByFunctionReplaceVisitor.java index 769fcb18d..423bf1c46 100644 --- a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/GroupByFunctionReplaceVisitor.java +++ b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/GroupByFunctionReplaceVisitor.java @@ -2,6 +2,8 @@ package com.tencent.supersonic.common.util.jsqlparser; import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.function.UnaryOperator; import lombok.extern.slf4j.Slf4j; import net.sf.jsqlparser.expression.Expression; import net.sf.jsqlparser.expression.Function; @@ -14,11 +16,17 @@ import org.apache.commons.lang3.StringUtils; public class GroupByFunctionReplaceVisitor implements GroupByVisitor { private Map functionMap; + private Map functionCallMap; public GroupByFunctionReplaceVisitor(Map functionMap) { this.functionMap = functionMap; } + public GroupByFunctionReplaceVisitor(Map functionMap, Map functionCallMap) { + this.functionMap = functionMap; + this.functionCallMap = functionCallMap; + } + public void visit(GroupByElement groupByElement) { groupByElement.getGroupByExpressionList(); ExpressionList groupByExpressionList = groupByElement.getGroupByExpressionList(); @@ -28,9 +36,17 @@ public class GroupByFunctionReplaceVisitor implements GroupByVisitor { Expression expression = groupByExpressions.get(i); if (expression instanceof Function) { Function function = (Function) expression; - String replaceName = functionMap.get(function.getName().toLowerCase()); + String functionName = function.getName().toLowerCase(); + String replaceName = functionMap.get(functionName); if (StringUtils.isNotBlank(replaceName)) { function.setName(replaceName); + if (Objects.nonNull(functionCallMap) && functionCallMap.containsKey(functionName)) { + Object ret = functionCallMap.get(functionName).apply(function.getParameters()); + if (Objects.nonNull(ret) && ret instanceof ExpressionList) { + ExpressionList expressionList = (ExpressionList) ret; + function.setParameters(expressionList); + } + } } } } diff --git a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlReplaceHelper.java b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlReplaceHelper.java index e05dd5a2b..d3e903f48 100644 --- a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlReplaceHelper.java +++ b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlReplaceHelper.java @@ -7,6 +7,7 @@ import java.util.List; 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; @@ -255,6 +256,11 @@ public class SqlReplaceHelper { } public static String replaceFunction(String sql, Map functionMap) { + return replaceFunction(sql, functionMap, null); + } + + public static String replaceFunction(String sql, Map functionMap, + Map functionCall) { Select selectStatement = SqlSelectHelper.getSelect(sql); //SelectBody selectBody = selectStatement.getSelectBody(); if (!(selectStatement instanceof PlainSelect)) { @@ -264,23 +270,24 @@ public class SqlReplaceHelper { plainSelectList.add((PlainSelect) selectStatement); List plainSelects = SqlSelectHelper.getPlainSelects(plainSelectList); for (PlainSelect plainSelect : plainSelects) { - replaceFunction(functionMap, plainSelect); + replaceFunction(functionMap, functionCall, plainSelect); } return selectStatement.toString(); } - private static void replaceFunction(Map functionMap, PlainSelect selectBody) { + private static void replaceFunction(Map functionMap, Map functionCall, + PlainSelect selectBody) { PlainSelect plainSelect = selectBody; //1. replace where dataDiff function Expression where = plainSelect.getWhere(); - FunctionNameReplaceVisitor visitor = new FunctionNameReplaceVisitor(functionMap); + FunctionNameReplaceVisitor visitor = new FunctionNameReplaceVisitor(functionMap, functionCall); if (Objects.nonNull(where)) { where.accept(visitor); } GroupByElement groupBy = plainSelect.getGroupBy(); if (Objects.nonNull(groupBy)) { - GroupByFunctionReplaceVisitor replaceVisitor = new GroupByFunctionReplaceVisitor(functionMap); + GroupByFunctionReplaceVisitor replaceVisitor = new GroupByFunctionReplaceVisitor(functionMap, functionCall); groupBy.accept(replaceVisitor); } diff --git a/headless/core/src/main/java/com/tencent/supersonic/headless/core/adaptor/db/PostgresqlAdaptor.java b/headless/core/src/main/java/com/tencent/supersonic/headless/core/adaptor/db/PostgresqlAdaptor.java index 0524c3dc1..e432e54f3 100644 --- a/headless/core/src/main/java/com/tencent/supersonic/headless/core/adaptor/db/PostgresqlAdaptor.java +++ b/headless/core/src/main/java/com/tencent/supersonic/headless/core/adaptor/db/PostgresqlAdaptor.java @@ -5,6 +5,10 @@ import com.tencent.supersonic.common.pojo.enums.TimeDimensionEnum; import com.tencent.supersonic.common.util.jsqlparser.SqlReplaceHelper; import java.util.HashMap; import java.util.Map; +import java.util.Objects; +import java.util.function.UnaryOperator; +import net.sf.jsqlparser.expression.StringValue; +import net.sf.jsqlparser.expression.operators.relational.ExpressionList; public class PostgresqlAdaptor extends DbAdaptor { @@ -44,10 +48,35 @@ public class PostgresqlAdaptor extends DbAdaptor { @Override public String functionNameCorrector(String sql) { Map functionMap = new HashMap<>(); - functionMap.put("MONTH".toLowerCase(), "toMonth"); - functionMap.put("DAY".toLowerCase(), "toDayOfMonth"); - functionMap.put("YEAR".toLowerCase(), "toYear"); - return SqlReplaceHelper.replaceFunction(sql, functionMap); + functionMap.put("MONTH".toLowerCase(), "TO_CHAR"); + functionMap.put("DAY".toLowerCase(), "TO_CHAR"); + functionMap.put("YEAR".toLowerCase(), "TO_CHAR"); + Map functionCall = new HashMap<>(); + functionCall.put("MONTH".toLowerCase(), o -> { + if (Objects.nonNull(o) && o instanceof ExpressionList) { + ExpressionList expressionList = (ExpressionList) o; + expressionList.add(new StringValue("MM")); + return expressionList; + } + return o; + }); + functionCall.put("DAY".toLowerCase(), o -> { + if (Objects.nonNull(o) && o instanceof ExpressionList) { + ExpressionList expressionList = (ExpressionList) o; + expressionList.add(new StringValue("dd")); + return expressionList; + } + return o; + }); + functionCall.put("YEAR".toLowerCase(), o -> { + if (Objects.nonNull(o) && o instanceof ExpressionList) { + ExpressionList expressionList = (ExpressionList) o; + expressionList.add(new StringValue("YYYY")); + return expressionList; + } + return o; + }); + return SqlReplaceHelper.replaceFunction(sql, functionMap, functionCall); } @Override