support select result data update

chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/DlTemplateService.java

@@ -2,6 +2,7 @@ package ai.chat2db.server.domain.api.service;
 
 import ai.chat2db.server.domain.api.param.DlCountParam;
 import ai.chat2db.server.domain.api.param.DlExecuteParam;
+import ai.chat2db.server.domain.api.param.UpdateSelectResultParam;
 import ai.chat2db.spi.model.ExecuteResult;
 import ai.chat2db.server.tools.base.wrapper.result.DataResult;
 
@@ -32,4 +33,12 @@ public interface DlTemplateService {
      */
     DataResult<Long> count(DlCountParam param);
 
+
+    /**
+     * 更新查询结果
+     * @param param
+     * @return
+     */
+    DataResult<String> updateSelectResult(UpdateSelectResultParam param);
+
 }

chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java

@@ -6,6 +6,8 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 
+import ai.chat2db.server.domain.api.param.UpdateSelectResultParam;
+import ai.chat2db.spi.MetaData;
 import com.alibaba.druid.DbType;
 import com.alibaba.druid.sql.PagerUtils;
 import com.alibaba.druid.sql.SQLUtils;
@@ -52,12 +54,12 @@ public class DlTemplateServiceImpl implements DlTemplateService {
             return ListResult.empty();
         }
         // 解析sql
-        SqlAnalyseParam sqlAnalyseParam = new SqlAnalyseParam();
+        RemoveSpecialGO(param);
-        sqlAnalyseParam.setDataSourceId(param.getDataSourceId());
-        sqlAnalyseParam.setSql(param.getSql());
         DbType dbType =
-            JdbcUtils.parse2DruidDbType(Chat2DBContext.getConnectInfo().getDbType());
+                JdbcUtils.parse2DruidDbType(Chat2DBContext.getConnectInfo().getDbType());
         List<String> sqlList = SqlUtils.parse(param.getSql(), dbType);
+
+
         if (CollectionUtils.isEmpty(sqlList)) {
             throw new BusinessException("dataSource.sqlAnalysisError");
         }
@@ -77,6 +79,16 @@ public class DlTemplateServiceImpl implements DlTemplateService {
         return listResult;
     }
 
+    private void RemoveSpecialGO(DlExecuteParam param) {
+        String sql = param.getSql();
+        if (StringUtils.isBlank(sql)) {
+            return;
+        }
+        sql = sql.replaceAll("(?i)\\s*go\\s*", ";");
+        param.setSql(sql);
+    }
+
+
     private ExecuteResult executeSQL(String originalSql, DbType dbType, DlExecuteParam param) {
         int pageNo = 1;
         int pageSize = 0;
@@ -103,12 +115,13 @@ public class DlTemplateServiceImpl implements DlTemplateService {
         ExecuteResult executeResult = execute(originalSql, offset, count);
         executeResult.setSqlType(sqlType);
         executeResult.setOriginalSql(originalSql);
+        executeResult.setCanEdit(SqlUtils.canEdit(originalSql));
 
         if (SqlTypeEnum.SELECT.getCode().equals(sqlType)) {
             executeResult.setPageNo(pageNo);
             executeResult.setPageSize(pageSize);
             executeResult.setHasNextPage(
-                CollectionUtils.size(executeResult.getDataList()) >= executeResult.getPageSize());
+                    CollectionUtils.size(executeResult.getDataList()) >= executeResult.getPageSize());
         } else {
             executeResult.setPageNo(pageNo);
             executeResult.setPageSize(CollectionUtils.size(executeResult.getDataList()));
@@ -118,9 +131,9 @@ public class DlTemplateServiceImpl implements DlTemplateService {
         // Splice row numbers
         List<Header> newHeaderList = new ArrayList<>();
         newHeaderList.add(Header.builder()
-            .name(I18nUtils.getMessage("sqlResult.rowNumber"))
+                .name(I18nUtils.getMessage("sqlResult.rowNumber"))
-            .dataType(DataTypeEnum.CHAT2DB_ROW_NUMBER
+                .dataType(DataTypeEnum.CHAT2DB_ROW_NUMBER
-                .getCode()).build());
+                        .getCode()).build());
         if (executeResult.getHeaderList() != null) {
             newHeaderList.addAll(executeResult.getHeaderList());
         }
@@ -158,7 +171,7 @@ public class DlTemplateServiceImpl implements DlTemplateService {
             return DataResult.of(0L);
         }
         DbType dbType =
-            JdbcUtils.parse2DruidDbType(Chat2DBContext.getConnectInfo().getDbType());
+                JdbcUtils.parse2DruidDbType(Chat2DBContext.getConnectInfo().getDbType());
         String sql = param.getSql();
         // 解析sql分页
         SQLStatement sqlStatement = SQLUtils.parseSingleStatement(sql, dbType);
@@ -173,14 +186,66 @@ public class DlTemplateServiceImpl implements DlTemplateService {
             return DataResult.of(0L);
         }
         String count = EasyCollectionUtils.stream(executeResult.getDataList())
-            .findFirst()
+                .findFirst()
-            .orElse(Collections.emptyList())
+                .orElse(Collections.emptyList())
-            .stream()
+                .stream()
-            .findFirst()
+                .findFirst()
-            .orElse("0");
+                .orElse("0");
         return DataResult.of(Long.valueOf(count));
     }
 
+    @Override
+    public DataResult<String> updateSelectResult(UpdateSelectResultParam param) {
+        StringBuilder stringBuilder = new StringBuilder();
+        MetaData metaSchema = Chat2DBContext.getMetaData();
+        for (int i = 0; i < param.getDataList().size(); i++) {
+            List<String> row = param.getDataList().get(i);
+            if (CollectionUtils.isEmpty(row)) {
+                continue;
+            }
+            List<String> odlRow = param.getOldDataList().get(i);
+
+            String sql = getUpdateSql(param, row, odlRow, metaSchema);
+            stringBuilder.append(sql+";\n");
+        }
+        return DataResult.of(stringBuilder.toString());
+    }
+
+    private String getUpdateSql(UpdateSelectResultParam param, List<String> row, List<String> odlRow, MetaData metaSchema) {
+        StringBuilder script = new StringBuilder();
+        script.append("UPDATE ").append(metaSchema.getMetaDataName(param.getDatabaseName(), param.getSchemaName(), param.getTableName()))
+                .append(" set ");
+        for (int i = 0; i < row.size(); i++) {
+            String newValue = row.get(i);
+            String oldValue = odlRow.get(i);
+            if (StringUtils.equals(newValue, oldValue)) {
+                continue;
+            }
+            Header header = param.getHeaderList().get(i);
+            String newSqlValue = SqlUtils.getSqlValue(newValue, header.getDataType());
+            script.append(metaSchema.getMetaDataName(header.getName()))
+                    .append(" = ")
+                    .append(newSqlValue)
+                    .append(",");
+        }
+        script.deleteCharAt(script.length() - 1);
+        script.append(" where ");
+        for (int i = 0; i < odlRow.size(); i++) {
+            String oldValue = odlRow.get(i);
+            if (oldValue == null) {
+                continue;
+            }
+            Header header = param.getHeaderList().get(i);
+            script.append(metaSchema.getMetaDataName(header.getName()))
+                    .append(" = ")
+                    .append(SqlUtils.getSqlValue(oldValue, header.getDataType()))
+                    .append(" and ");
+        }
+
+        script.delete(script.length() - 4, script.length());
+        return script.toString();
+    }
+
     private ExecuteResult execute(String sql, Integer offset, Integer count) {
         ExecuteResult executeResult;
         try {
@@ -188,10 +253,10 @@ public class DlTemplateServiceImpl implements DlTemplateService {
         } catch (SQLException e) {
             log.warn("执行sql:{}异常", sql, e);
             executeResult = ExecuteResult.builder()
-                .sql(sql)
+                    .sql(sql)
-                .success(Boolean.FALSE)
+                    .success(Boolean.FALSE)
-                .message(e.getMessage())
+                    .message(e.getMessage())
-                .build();
+                    .build();
         }
         return executeResult;
     }

chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/enums/DataTypeEnum.java

@@ -96,4 +96,62 @@ public enum DataTypeEnum implements BaseEnum<String> {
     public String getCode() {
         return this.name();
     }
+
+    public static DataTypeEnum getByCode(String code) {
+        for (DataTypeEnum value : DataTypeEnum.values()) {
+            if (value.getCode().equals(code)) {
+                return value;
+            }
+        }
+        return DataTypeEnum.UNKNOWN;
+    }
+
+    public String getSqlValue(String value) {
+        if (this == DataTypeEnum.BOOLEAN) {
+            return value;
+        }
+        if (this == DataTypeEnum.NUMERIC) {
+            return value;
+        }
+        if (this == DataTypeEnum.STRING) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.DATETIME) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.BINARY) {
+            return "''";
+        }
+        if (this == DataTypeEnum.CONTENT) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.STRUCT) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.DOCUMENT) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.ARRAY) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.OBJECT) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.REFERENCE) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.ROWID) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.ANY) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.UNKNOWN) {
+            return "'" + value + "'";
+        }
+        if (this == DataTypeEnum.CHAT2DB_ROW_NUMBER) {
+            return "'" + value + "'";
+        }
+        return "'" + value + "'";
+    }
 }

chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java

@@ -1,8 +1,10 @@
 package ai.chat2db.spi.jdbc;
 
 import java.sql.Connection;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import ai.chat2db.spi.MetaData;
 import ai.chat2db.spi.SqlBuilder;
@@ -105,4 +107,12 @@ public class DefaultMetaService implements MetaData {
     public TableMeta getTableMeta(String databaseName, String schemaName, String tableName) {
         return null;
     }
+
+    @Override
+    public String getMetaDataName(String... names) {
+        return Arrays.stream(names).collect(Collectors.joining("."));
+    }
+
+
+
 }

chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/ColumnType.java

@@ -6,7 +6,6 @@ import lombok.Data;
 @Data
 @AllArgsConstructor
 public class ColumnType {
-
     private String typeName;
     private boolean supportLength;
     private boolean supportScale;
@@ -17,8 +16,7 @@ public class ColumnType {
     private boolean supportComments;
     private boolean supportDefaultValue;
     private boolean supportExtent;
-
     private boolean supportValue;
-
     private boolean supportUnit;
+
 }

chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/ExecuteResult.java

@@ -94,4 +94,10 @@ public class ExecuteResult {
      * 执行持续时间
      */
     private Long duration;
+
+
+    /**
+     * 返回结果是否可以编辑
+     */
+    private boolean canEdit;
 }