(improvement)(Headless) fix pg date function name corrector (#950)

This commit is contained in:
jipeli
2024-04-26 15:34:36 +08:00
committed by GitHub
parent a551f3cc16
commit 11c2e0505b
4 changed files with 79 additions and 10 deletions

View File

@@ -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<String, String> functionMap;
private Map<String, UnaryOperator> functionCallMap;
public FunctionNameReplaceVisitor(Map<String, String> functionMap) {
this.functionMap = functionMap;
}
public FunctionNameReplaceVisitor(Map<String, String> functionMap, Map<String, UnaryOperator> 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);
}
}
}
}

View File

@@ -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<String, String> functionMap;
private Map<String, UnaryOperator> functionCallMap;
public GroupByFunctionReplaceVisitor(Map<String, String> functionMap) {
this.functionMap = functionMap;
}
public GroupByFunctionReplaceVisitor(Map<String, String> functionMap, Map<String, UnaryOperator> 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);
}
}
}
}
}

View File

@@ -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<String, String> functionMap) {
return replaceFunction(sql, functionMap, null);
}
public static String replaceFunction(String sql, Map<String, String> functionMap,
Map<String, UnaryOperator> 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<PlainSelect> plainSelects = SqlSelectHelper.getPlainSelects(plainSelectList);
for (PlainSelect plainSelect : plainSelects) {
replaceFunction(functionMap, plainSelect);
replaceFunction(functionMap, functionCall, plainSelect);
}
return selectStatement.toString();
}
private static void replaceFunction(Map<String, String> functionMap, PlainSelect selectBody) {
private static void replaceFunction(Map<String, String> functionMap, Map<String, UnaryOperator> 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);
}

View File

@@ -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<String, String> 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<String, UnaryOperator> 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