From a7d845f224b50ad1fc08ef0a438df4cebbed272b Mon Sep 17 00:00:00 2001
From: LXW <1264174498@qq.com>
Date: Mon, 27 May 2024 20:52:13 +0800
Subject: [PATCH] (improvement)(Headless) Automatically populate model field
 names with comments from table fields (#1036)

Co-authored-by: jolunoluo
---
 .../supersonic/common/pojo/QueryColumn.java   |  1 +
 .../service/impl/DatabaseServiceImpl.java     | 47 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/common/src/main/java/com/tencent/supersonic/common/pojo/QueryColumn.java b/common/src/main/java/com/tencent/supersonic/common/pojo/QueryColumn.java
index d8faae44e..c2c5374cf 100644
--- a/common/src/main/java/com/tencent/supersonic/common/pojo/QueryColumn.java
+++ b/common/src/main/java/com/tencent/supersonic/common/pojo/QueryColumn.java
@@ -16,6 +16,7 @@ public class QueryColumn {
     private Boolean authorized = true;
     private String dataFormatType;
     private DataFormat dataFormat;
+    private String comment;
 
     public QueryColumn(String nameEn, String type) {
         this.type = type;
diff --git a/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/DatabaseServiceImpl.java b/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/DatabaseServiceImpl.java
index 91aa06c88..44df0066e 100644
--- a/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/DatabaseServiceImpl.java
+++ b/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/DatabaseServiceImpl.java
@@ -3,6 +3,10 @@ package com.tencent.supersonic.headless.server.service.impl;
 import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
 import com.google.common.collect.Lists;
 import com.tencent.supersonic.auth.api.authentication.pojo.User;
+import com.tencent.supersonic.common.pojo.Constants;
+import com.tencent.supersonic.common.pojo.Pair;
+import com.tencent.supersonic.common.pojo.QueryColumn;
+import com.tencent.supersonic.common.util.jsqlparser.SqlSelectHelper;
 import com.tencent.supersonic.headless.api.pojo.request.DatabaseReq;
 import com.tencent.supersonic.headless.api.pojo.request.SqlExecuteReq;
 import com.tencent.supersonic.headless.api.pojo.response.DatabaseResp;
@@ -23,11 +27,13 @@ import com.tencent.supersonic.headless.server.service.DatabaseService;
 import com.tencent.supersonic.headless.server.service.ModelService;
 import com.tencent.supersonic.headless.server.utils.DatabaseConverter;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Lazy;
 import org.springframework.stereotype.Service;
 import org.springframework.util.CollectionUtils;
 
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -128,7 +134,9 @@ public class DatabaseServiceImpl extends ServiceImpl<DatabaseDOMapper, DatabaseD
         checkPermission(databaseResp, user);
         String sql = sqlExecuteReq.getSql();
         sql = SqlVariableParseUtils.parse(sql, sqlExecuteReq.getSqlVariables(), Lists.newArrayList());
-        return executeSql(sql, databaseResp);
+        SemanticQueryResp semanticQueryResp = executeSql(sql, databaseResp);
+        fillColumnComment(sql, databaseResp, semanticQueryResp);
+        return semanticQueryResp;
     }
 
     @Override
@@ -136,6 +144,43 @@ public class DatabaseServiceImpl extends ServiceImpl<DatabaseDOMapper, DatabaseD
         return queryWithColumns(sql, DatabaseConverter.convert(databaseResp));
     }
 
+    private void fillColumnComment(String sql, DatabaseResp databaseResp,
+                                   SemanticQueryResp semanticQueryResp) {
+        Pair<String, String> dbTableName = getDbTableName(sql, databaseResp);
+        String db = dbTableName.first;
+        String table = dbTableName.second;
+        if (StringUtils.isBlank(db) || StringUtils.isBlank(table)) {
+            return;
+        }
+        SemanticQueryResp columnsWithComment = getColumns(databaseResp, db, table);
+        Map<String, String> columnCommentMap = getColumnCommentMap(columnsWithComment.getResultList());
+        List<QueryColumn> columns = semanticQueryResp.getColumns();
+        for (QueryColumn column : columns) {
+            column.setComment(columnCommentMap.get(column.getNameEn()));
+        }
+    }
+
+    private Map<String, String> getColumnCommentMap(List<Map<String, Object>> resultList) {
+        Map<String, String> map = new HashMap<>();
+        for (Map<String, Object> result : resultList) {
+            map.put(String.valueOf(result.get("name")), String.valueOf(result.get("comment")));
+        }
+        return map;
+    }
+
+    private Pair<String, String> getDbTableName(String sql, DatabaseResp databaseResp) {
+        String dbTableName = SqlSelectHelper.getDbTableName(sql);
+        if (StringUtils.isBlank(dbTableName)) {
+            return Pair.pair("", "");
+        }
+        if (dbTableName.contains(Constants.DOT)) {
+            String db = dbTableName.split("\\.")[0];
+            String table = dbTableName.split("\\.")[1];
+            return Pair.pair(db, table);
+        }
+        return Pair.pair(databaseResp.getDatabase(), dbTableName);
+    }
+
     @Override
     public Map<String, List<DatabaseParameter>> getDatabaseParameters() {
         return DbParameterFactory.getMap().entrySet().stream().collect(LinkedHashMap::new,