Merge aaf2d46a56 into 94e853f57e

[improvement][headless]Support dataSetNames that contain dash.

[improvement][headless]Support dataSetNames that contain dash.

chat/server/src/main/java/com/tencent/supersonic/chat/server/processor/parse/ErrorMsgRewriteProcessor.java

@@ -27,7 +27,7 @@ public class ErrorMsgRewriteProcessor implements ParseResultProcessor {
 
     private static final Logger keyPipelineLog = LoggerFactory.getLogger("keyPipeline");
 
-    public static final String APP_KEY_ERROR_MESSAGE = "REWRITE_ERROR_MESSAGE";
+    public static final String APP_KEY = "REWRITE_ERROR_MESSAGE";
     private static final String REWRITE_ERROR_MESSAGE_INSTRUCTION = ""
             + "#Role: You are a data business partner who closely interacts with business people.\n"
             + "#Task: Your will be provided with user input, system output and some examples, "
@@ -38,7 +38,7 @@ public class ErrorMsgRewriteProcessor implements ParseResultProcessor {
             + "#Examples: {{examples}}\n" + "#Response: ";
 
     public ErrorMsgRewriteProcessor() {
-        ChatAppManager.register(APP_KEY_ERROR_MESSAGE,
+        ChatAppManager.register(APP_KEY,
                 ChatApp.builder().prompt(REWRITE_ERROR_MESSAGE_INSTRUCTION).name("异常提示改写")
                         .appModule(AppModule.CHAT).description("通过大模型将异常信息改写为更友好和引导性的提示用语")
                         .enable(true).build());
@@ -46,7 +46,7 @@ public class ErrorMsgRewriteProcessor implements ParseResultProcessor {
 
     @Override
     public boolean accept(ParseContext parseContext) {
-        ChatApp chatApp = parseContext.getAgent().getChatAppConfig().get(APP_KEY_ERROR_MESSAGE);
+        ChatApp chatApp = parseContext.getAgent().getChatAppConfig().get(APP_KEY);
         return StringUtils.isNotBlank(parseContext.getResponse().getErrorMsg())
                 && Objects.nonNull(chatApp) && chatApp.isEnable();
     }
@@ -54,16 +54,20 @@ public class ErrorMsgRewriteProcessor implements ParseResultProcessor {
     @Override
     public void process(ParseContext parseContext) {
         String errMsg = parseContext.getResponse().getErrorMsg();
-        ChatApp chatApp = parseContext.getAgent().getChatAppConfig().get(APP_KEY_ERROR_MESSAGE);
+        ChatApp chatApp = parseContext.getAgent().getChatAppConfig().get(APP_KEY);
         Map<String, Object> variables = new HashMap<>();
         variables.put("user_question", parseContext.getRequest().getQueryText());
         variables.put("system_message", errMsg);
 
         StringBuilder exampleStr = new StringBuilder();
-        parseContext.getResponse().getUsedExemplars().forEach(e -> exampleStr.append(
+        if (parseContext.getResponse().getUsedExemplars() != null) {
-                String.format("<Question:{%s},Schema:{%s}> ", e.getQuestion(), e.getDbSchema())));
+            parseContext.getResponse().getUsedExemplars().forEach(e -> exampleStr.append(String
-        parseContext.getAgent().getExamples()
+                    .format("<Question:{%s},Schema:{%s}> ", e.getQuestion(), e.getDbSchema())));
-                .forEach(e -> exampleStr.append(String.format("<Question:{%s}> ", e)));
+        }
+        if (parseContext.getAgent().getExamples() != null) {
+            parseContext.getAgent().getExamples()
+                    .forEach(e -> exampleStr.append(String.format("<Question:{%s}> ", e)));
+        }
         variables.put("examples", exampleStr);
 
         Prompt prompt = PromptTemplate.from(chatApp.getPrompt()).apply(variables);

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

@@ -38,6 +38,10 @@ public class SqlReplaceHelper {
 
     private final static double replaceColumnThreshold = 0.4;
 
+    public static String escapeTableName(String table) {
+        return String.format("`%s`", table);
+    }
+
     public static String replaceAggFields(String sql,
             Map<String, Pair<String, String>> fieldNameToAggMap) {
         Select selectStatement = SqlSelectHelper.getSelect(sql);

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

@@ -228,7 +228,7 @@ public class SqlSelectHelper {
             statement = CCJSqlParserUtil.parse(sql);
         } catch (JSQLParserException e) {
             log.error("parse error, sql:{}", sql, e);
-            return null;
+            throw new RuntimeException(e);
         }
 
         if (statement instanceof ParenthesedSelect) {

headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/request/QueryStructReq.java

@@ -2,6 +2,7 @@ package com.tencent.supersonic.headless.api.pojo.request;
 
 import com.google.common.collect.Lists;
 import com.tencent.supersonic.common.jsqlparser.SqlAddHelper;
+import com.tencent.supersonic.common.jsqlparser.SqlReplaceHelper;
 import com.tencent.supersonic.common.pojo.Aggregator;
 import com.tencent.supersonic.common.pojo.Constants;
 import com.tencent.supersonic.common.pojo.DateConf;
@@ -281,7 +282,7 @@ public class QueryStructReq extends SemanticQueryReq {
 
     public String getTableName() {
         if (StringUtils.isNotBlank(dataSetName)) {
-            return dataSetName;
+            return SqlReplaceHelper.escapeTableName(dataSetName);
         }
         if (dataSetId != null) {
             return Constants.TABLE_PREFIX + dataSetId;

headless/chat/src/main/java/com/tencent/supersonic/headless/chat/utils/QueryReqBuilder.java

@@ -15,12 +15,19 @@ import com.tencent.supersonic.headless.api.pojo.request.QueryFilter;
 import com.tencent.supersonic.headless.api.pojo.request.QueryMultiStructReq;
 import com.tencent.supersonic.headless.api.pojo.request.QuerySqlReq;
 import com.tencent.supersonic.headless.api.pojo.request.QueryStructReq;
+import com.tencent.supersonic.headless.api.pojo.response.DataSetResp;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.BeanUtils;
 import org.springframework.util.CollectionUtils;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 @Slf4j
@@ -97,15 +104,16 @@ public class QueryReqBuilder {
      * convert to QueryS2SQLReq
      *
      * @param querySql
-     * @param dataSetId
+     * @param dataSet
      * @return
      */
-    public static QuerySqlReq buildS2SQLReq(String querySql, Long dataSetId) {
+    public static QuerySqlReq buildS2SQLReq(String querySql, DataSetResp dataSet) {
         QuerySqlReq querySQLReq = new QuerySqlReq();
         if (Objects.nonNull(querySql)) {
             querySQLReq.setSql(querySql);
         }
-        querySQLReq.setDataSetId(dataSetId);
+        querySQLReq.setDataSetId(dataSet.getId());
+        querySQLReq.setDataSetName(dataSet.getName());
         return querySQLReq;
     }
 

headless/server/src/main/java/com/tencent/supersonic/headless/server/aspect/S2DataPermissionAspect.java

@@ -7,6 +7,7 @@ import com.tencent.supersonic.auth.api.authorization.request.QueryAuthResReq;
 import com.tencent.supersonic.auth.api.authorization.response.AuthorizedResourceResp;
 import com.tencent.supersonic.auth.api.authorization.service.AuthService;
 import com.tencent.supersonic.common.jsqlparser.SqlAddHelper;
+import com.tencent.supersonic.common.jsqlparser.SqlReplaceHelper;
 import com.tencent.supersonic.common.pojo.Filter;
 import com.tencent.supersonic.common.pojo.QueryAuthorization;
 import com.tencent.supersonic.common.pojo.User;
@@ -73,6 +74,15 @@ public class S2DataPermissionAspect {
         SemanticQueryReq queryReq = null;
         if (objects[0] instanceof SemanticQueryReq) {
             queryReq = (SemanticQueryReq) objects[0];
+            if (queryReq instanceof QuerySqlReq) {
+                QuerySqlReq sqlReq = (QuerySqlReq) queryReq;
+                if (sqlReq.getDataSetName() != null) {
+                    String escapedTable = SqlReplaceHelper.escapeTableName(sqlReq.getDataSetName());
+                    sqlReq.setSql(sqlReq.getSql().replaceAll(
+                            String.format(" %s ", sqlReq.getDataSetName()),
+                            String.format(" %s ", escapedTable)));
+                }
+            }
         }
         if (queryReq == null) {
             throw new InvalidArgumentException("queryReq is not Invalid");

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

@@ -140,7 +140,8 @@ public class QueryUtils {
                 || type.equalsIgnoreCase("float") || type.equalsIgnoreCase("double")
                 || type.equalsIgnoreCase("real") || type.equalsIgnoreCase("numeric")
                 || type.toLowerCase().startsWith("decimal") || type.toLowerCase().startsWith("uint")
-                || type.toLowerCase().startsWith("int") || type.toLowerCase().equalsIgnoreCase("decfloat");
+                || type.toLowerCase().startsWith("int")
+                || type.toLowerCase().equalsIgnoreCase("decfloat");
     }
 
     private String getName(String nameEn) {

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

@@ -20,6 +20,7 @@ import com.tencent.supersonic.headless.api.pojo.response.SemanticQueryResp;
 import com.tencent.supersonic.headless.server.facade.service.SemanticLayerService;
 import com.tencent.supersonic.headless.server.persistence.dataobject.DomainDO;
 import com.tencent.supersonic.headless.server.persistence.repository.DomainRepository;
+import com.tencent.supersonic.headless.server.service.DataSetService;
 import com.tencent.supersonic.headless.server.service.DatabaseService;
 import com.tencent.supersonic.headless.server.service.SchemaService;
 import com.tencent.supersonic.util.DataUtils;
@@ -46,6 +47,8 @@ public class BaseTest extends BaseApplication {
     private AgentService agentService;
     @Autowired
     protected DatabaseService databaseService;
+    @Autowired
+    protected DataSetService dataSetService;
 
     protected Agent agent;
     protected SemanticSchema schema;

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

@@ -2,6 +2,7 @@ package com.tencent.supersonic.headless;
 
 import com.tencent.supersonic.common.pojo.User;
 import com.tencent.supersonic.demo.S2VisitsDemo;
+import com.tencent.supersonic.headless.api.pojo.response.DataSetResp;
 import com.tencent.supersonic.headless.api.pojo.response.SemanticTranslateResp;
 import com.tencent.supersonic.headless.chat.utils.QueryReqBuilder;
 import org.junit.jupiter.api.BeforeEach;
@@ -18,14 +19,15 @@ import static org.junit.Assert.assertTrue;
 
 public class TranslatorTest extends BaseTest {
 
-    private Long dataSetId;
+    private DataSetResp dataSet;
 
     @BeforeEach
     public void init() {
         agent = getAgentByName(S2VisitsDemo.AGENT_NAME);
         schema = schemaService.getSemanticSchema(agent.getDataSetIds());
         if (Objects.nonNull(agent)) {
-            dataSetId = agent.getDataSetIds().stream().findFirst().get();
+            long dataSetId = agent.getDataSetIds().stream().findFirst().get();
+            dataSet = dataSetService.getDataSet(dataSetId);
         }
     }
 
@@ -34,7 +36,7 @@ public class TranslatorTest extends BaseTest {
         String sql =
                 "SELECT SUM(访问次数) AS _总访问次数_ FROM 超音数数据集 WHERE 数据日期 >= '2024-11-15' AND 数据日期 <= '2024-12-15'";
         SemanticTranslateResp explain = semanticLayerService
-                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSetId), User.getDefaultUser());
+                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSet), User.getDefaultUser());
         assertNotNull(explain);
         assertNotNull(explain.getQuerySQL());
         assertTrue(explain.getQuerySQL().contains("count(1)"));
@@ -45,7 +47,7 @@ public class TranslatorTest extends BaseTest {
     public void testSql_1() throws Exception {
         String sql = "SELECT 部门, SUM(访问次数) AS 总访问次数 FROM 超音数PVUV统计 GROUP BY 部门 ";
         SemanticTranslateResp explain = semanticLayerService
-                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSetId), User.getDefaultUser());
+                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSet), User.getDefaultUser());
         assertNotNull(explain);
         assertNotNull(explain.getQuerySQL());
         assertTrue(explain.getQuerySQL().contains("department"));
@@ -59,7 +61,7 @@ public class TranslatorTest extends BaseTest {
         String sql =
                 "WITH _department_visits_ AS (SELECT 部门, SUM(访问次数) AS _total_visits_ FROM 超音数数据集 WHERE 数据日期 >= '2024-11-15' AND 数据日期 <= '2024-12-15' GROUP BY 部门) SELECT 部门 FROM _department_visits_ ORDER BY _total_visits_ DESC LIMIT 2";
         SemanticTranslateResp explain = semanticLayerService
-                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSetId), User.getDefaultUser());
+                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSet), User.getDefaultUser());
         assertNotNull(explain);
         assertNotNull(explain.getQuerySQL());
         assertTrue(explain.getQuerySQL().toLowerCase().contains("department"));
@@ -73,7 +75,7 @@ public class TranslatorTest extends BaseTest {
         String sql =
                 "WITH recent_data AS (SELECT 用户名, 访问次数 FROM 超音数数据集 WHERE 部门 = 'marketing' AND 数据日期 >= '2024-12-01' AND 数据日期 <= '2024-12-15') SELECT 用户名 FROM recent_data ORDER BY 访问次数 DESC LIMIT 1";
         SemanticTranslateResp explain = semanticLayerService
-                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSetId), User.getDefaultUser());
+                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSet), User.getDefaultUser());
         assertNotNull(explain);
         assertNotNull(explain.getQuerySQL());
         assertTrue(explain.getQuerySQL().toLowerCase().contains("department"));
@@ -89,7 +91,7 @@ public class TranslatorTest extends BaseTest {
                         Paths.get(ClassLoader.getSystemResource("sql/testUnion.sql").toURI())),
                 StandardCharsets.UTF_8);
         SemanticTranslateResp explain = semanticLayerService
-                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSetId), User.getDefaultUser());
+                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSet), User.getDefaultUser());
         assertNotNull(explain);
         assertNotNull(explain.getQuerySQL());
         assertTrue(explain.getQuerySQL().contains("user_name"));
@@ -105,7 +107,7 @@ public class TranslatorTest extends BaseTest {
                         Paths.get(ClassLoader.getSystemResource("sql/testWith.sql").toURI())),
                 StandardCharsets.UTF_8);
         SemanticTranslateResp explain = semanticLayerService
-                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSetId), User.getDefaultUser());
+                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSet), User.getDefaultUser());
         assertNotNull(explain);
         assertNotNull(explain.getQuerySQL());
         executeSql(explain.getQuerySQL());
@@ -119,7 +121,7 @@ public class TranslatorTest extends BaseTest {
                         Paths.get(ClassLoader.getSystemResource("sql/testSubquery.sql").toURI())),
                 StandardCharsets.UTF_8);
         SemanticTranslateResp explain = semanticLayerService
-                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSetId), User.getDefaultUser());
+                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSet), User.getDefaultUser());
         assertNotNull(explain);
         assertNotNull(explain.getQuerySQL());
         executeSql(explain.getQuerySQL());