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

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

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

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

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<String, String> functionMap = new HashMap<>();
-        functionMap.put("MONTH".toLowerCase(), "toMonth");
+        functionMap.put("MONTH".toLowerCase(), "TO_CHAR");
-        functionMap.put("DAY".toLowerCase(), "toDayOfMonth");
+        functionMap.put("DAY".toLowerCase(), "TO_CHAR");
-        functionMap.put("YEAR".toLowerCase(), "toYear");
+        functionMap.put("YEAR".toLowerCase(), "TO_CHAR");
-        return SqlReplaceHelper.replaceFunction(sql, functionMap);
+        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