From 6f968086c9c68a7e304b43f8e263b8fc10d6b365 Mon Sep 17 00:00:00 2001 From: SwallowGG <1558143046@qq.com> Date: Wed, 31 Jan 2024 19:16:15 +0800 Subject: [PATCH] feat:support download task --- .../domain/api/enums/DeletedTypeEnum.java | 6 + .../domain/api/enums/TaskStatusEnum.java | 6 + .../server/domain/api/enums/TaskTypeEnum.java | 24 ++ .../chat2db/server/domain/api/model/Task.java | 84 +++++ .../domain/api/param/TaskCreateParam.java | 54 +++ .../domain/api/param/TaskPageParam.java | 19 ++ .../domain/api/param/TaskUpdateParam.java | 47 +++ .../domain/api/service/TaskService.java | 45 +++ .../domain/core/converter/TaskConverter.java | 21 ++ .../domain/core/impl/TaskServiceImpl.java | 69 ++++ .../server/domain/repository/MapperUtils.java | 10 + .../domain/repository/entity/TaskDO.java | 101 ++++++ .../domain/repository/mapper/TaskMapper.java | 19 ++ .../resources/db/migration/V2_1_9__task.sql | 20 ++ .../src/main/resources/mapper/TaskMapper.xml | 36 ++ .../ai/chat2db/server/start/Application.java | 1 - .../server/start/config/StdinReader.java | 30 -- .../start/config/oauth/SaLogForSlf4j.java | 44 --- .../start/config/oauth/SaTokenConfigure.java | 21 -- .../test/mybatis/MybatisGeneratorTest.java | 22 +- .../api/controller/rdb/RdbDmlController.java | 34 -- .../api/controller/task/ExportController.java | 42 +++ .../api/controller/task/TaskController.java | 77 +++++ .../controller/task/biz/TaskBizService.java | 312 ++++++++++++++++++ .../web/api/http/GatewayClientService.java | 11 +- .../chat2db/server/web/api/ws/WsService.java | 33 -- .../java/ai/chat2db/spi/sql/ConnectInfo.java | 26 ++ 27 files changed, 1040 insertions(+), 174 deletions(-) create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/DeletedTypeEnum.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/TaskStatusEnum.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/TaskTypeEnum.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/model/Task.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskCreateParam.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskPageParam.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskUpdateParam.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/TaskService.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/converter/TaskConverter.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/TaskServiceImpl.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/MapperUtils.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/entity/TaskDO.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/mapper/TaskMapper.java create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/db/migration/V2_1_9__task.sql create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/mapper/TaskMapper.xml delete mode 100644 chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/StdinReader.java delete mode 100644 chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/oauth/SaLogForSlf4j.java delete mode 100644 chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/oauth/SaTokenConfigure.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/ExportController.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/TaskController.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/biz/TaskBizService.java diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/DeletedTypeEnum.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/DeletedTypeEnum.java new file mode 100644 index 00000000..aac6872b --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/DeletedTypeEnum.java @@ -0,0 +1,6 @@ +package ai.chat2db.server.domain.api.enums; + +public enum DeletedTypeEnum { + + Y,N +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/TaskStatusEnum.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/TaskStatusEnum.java new file mode 100644 index 00000000..821ac74b --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/TaskStatusEnum.java @@ -0,0 +1,6 @@ +package ai.chat2db.server.domain.api.enums; + +public enum TaskStatusEnum { + + INIT, PROCESSING, FINISH, ERROR +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/TaskTypeEnum.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/TaskTypeEnum.java new file mode 100644 index 00000000..7d3b9341 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/TaskTypeEnum.java @@ -0,0 +1,24 @@ +package ai.chat2db.server.domain.api.enums; + +public enum TaskTypeEnum { + + /** + * download table data + */ + DOWNLOAD_TABLE_DATA, + + /** + * upload table data + */ + UPLOAD_TABLE_DATA, + + /** + * download table structure + */ + DOWNLOAD_TABLE_STRUCTURE, + + /** + * upload table structure + */ + UPLOAD_TABLE_STRUCTURE, +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/model/Task.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/model/Task.java new file mode 100644 index 00000000..52c2a1b3 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/model/Task.java @@ -0,0 +1,84 @@ +package ai.chat2db.server.domain.api.model; + +import lombok.Data; + +import java.io.Serializable; +import java.util.Date; + +@Data +public class Task implements Serializable { + /** + * 主键 + */ + private Long id; + + /** + * 创建时间 + */ + private Date gmtCreate; + + /** + * 修改时间 + */ + private Date gmtModified; + + /** + * 数据源连接ID + */ + private Long dataSourceId; + + /** + * db名称 + */ + private String databaseName; + + /** + * schema名称 + */ + private String schemaName; + + /** + * table_name + */ + private String tableName; + + /** + * 是否被删除,y表示删除,n表示未删除 + */ + private String deleted; + + /** + * 用户id + */ + private Long userId; + + /** + * task type, such as: DOWNLOAD_DATA, UPLOAD_TABLE_DATA, DOWNLOAD_TABLE_STRUCTURE, UPLOAD_TABLE_STRUCTURE, + */ + private String taskType; + + /** + * task status + */ + private String taskStatus; + + /** + * task progress + */ + private String taskProgress; + + /** + * task name + */ + private String taskName; + + /** + * download url + */ + private String downloadUrl; + + /** + * task content + */ + private byte[] content; +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskCreateParam.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskCreateParam.java new file mode 100644 index 00000000..553beb34 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskCreateParam.java @@ -0,0 +1,54 @@ +package ai.chat2db.server.domain.api.param; + +import lombok.Data; + +import java.io.Serializable; + +@Data +public class TaskCreateParam implements Serializable { + + private static final long serialVersionUID = 1L; + + /** + * 数据源连接ID + */ + private Long dataSourceId; + + /** + * db名称 + */ + private String databaseName; + + /** + * schema名称 + */ + private String schemaName; + + /** + * table_name + */ + private String tableName; + + /** + * 用户id + */ + private Long userId; + + + /** + * task progress + */ + private String taskProgress; + + /** + * task name + */ + private String taskName; + + /** + * task type, such as: DOWNLOAD_DATA, UPLOAD_TABLE_DATA, DOWNLOAD_TABLE_STRUCTURE, UPLOAD_TABLE_STRUCTURE, + */ + private String taskType; + + +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskPageParam.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskPageParam.java new file mode 100644 index 00000000..ecf20755 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskPageParam.java @@ -0,0 +1,19 @@ +package ai.chat2db.server.domain.api.param; + +import ai.chat2db.server.tools.base.wrapper.param.PageQueryParam; +import lombok.Data; + +import java.io.Serializable; +import java.util.List; + +@Data +public class TaskPageParam extends PageQueryParam implements Serializable { + + + private Long userId; + + private List taskType; + + private String taskStatus; + +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskUpdateParam.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskUpdateParam.java new file mode 100644 index 00000000..9465bc19 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/TaskUpdateParam.java @@ -0,0 +1,47 @@ +package ai.chat2db.server.domain.api.param; + +import lombok.Data; + +import java.io.Serializable; + +@Data +public class TaskUpdateParam implements Serializable { + + + private static final long serialVersionUID = 1L; + + /** + * task id + */ + private Long id; + + /** + * 用户id + */ + private Long userId; + + /** + * task type, such as: DOWNLOAD_DATA, UPLOAD_TABLE_DATA, DOWNLOAD_TABLE_STRUCTURE, UPLOAD_TABLE_STRUCTURE, + */ + private String taskStatus; + + /** + * task progress + */ + private String taskProgress; + + /** + * task name + */ + private String taskName; + + /** + * task description + */ + private String downloadUrl; + + /** + * task content + */ + private byte[] content; +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/TaskService.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/TaskService.java new file mode 100644 index 00000000..e8d44bd8 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/TaskService.java @@ -0,0 +1,45 @@ +package ai.chat2db.server.domain.api.service; + +import ai.chat2db.server.domain.api.model.Task; +import ai.chat2db.server.domain.api.param.TaskCreateParam; +import ai.chat2db.server.domain.api.param.TaskPageParam; +import ai.chat2db.server.domain.api.param.TaskUpdateParam; +import ai.chat2db.server.tools.base.wrapper.result.ActionResult; +import ai.chat2db.server.tools.base.wrapper.result.DataResult; +import ai.chat2db.server.tools.base.wrapper.result.PageResult; + +public interface TaskService { + + /** + * create task + * + * @param param task param + * @return task id + */ + DataResult create(TaskCreateParam param); + + /** + * update task status + * + * @param param task param + * @return action result + */ + ActionResult updateStatus(TaskUpdateParam param); + + + /** + * get task list + * + * @param param task id + * @return task + */ + PageResult page(TaskPageParam param); + + /** + * get task + * + * @param id task id + * @return task + */ + DataResult get(Long id); +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/converter/TaskConverter.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/converter/TaskConverter.java new file mode 100644 index 00000000..e1f66aa8 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/converter/TaskConverter.java @@ -0,0 +1,21 @@ +package ai.chat2db.server.domain.core.converter; + +import ai.chat2db.server.domain.api.model.Task; +import ai.chat2db.server.domain.api.param.TaskCreateParam; +import ai.chat2db.server.domain.repository.entity.TaskDO; +import lombok.extern.slf4j.Slf4j; +import org.mapstruct.Mapper; + +import java.util.List; +@Slf4j +@Mapper(componentModel = "spring") +public abstract class TaskConverter { + + public abstract TaskDO todo(TaskCreateParam param); + + + public abstract Task toModel(TaskDO param); + + + public abstract List toModel(List param); +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/TaskServiceImpl.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/TaskServiceImpl.java new file mode 100644 index 00000000..25ba9a9f --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/TaskServiceImpl.java @@ -0,0 +1,69 @@ +package ai.chat2db.server.domain.core.impl; + +import ai.chat2db.server.domain.api.enums.DeletedTypeEnum; +import ai.chat2db.server.domain.api.enums.TaskStatusEnum; +import ai.chat2db.server.domain.api.model.Task; +import ai.chat2db.server.domain.api.param.TaskCreateParam; +import ai.chat2db.server.domain.api.param.TaskPageParam; +import ai.chat2db.server.domain.api.param.TaskUpdateParam; +import ai.chat2db.server.domain.api.service.TaskService; +import ai.chat2db.server.domain.core.converter.TaskConverter; +import ai.chat2db.server.domain.repository.MapperUtils; +import ai.chat2db.server.domain.repository.entity.TaskDO; +import ai.chat2db.server.tools.base.wrapper.result.ActionResult; +import ai.chat2db.server.tools.base.wrapper.result.DataResult; +import ai.chat2db.server.tools.base.wrapper.result.PageResult; +import com.baomidou.mybatisplus.core.metadata.IPage; +import com.baomidou.mybatisplus.core.metadata.OrderItem; +import com.baomidou.mybatisplus.extension.plugins.pagination.Page; +import com.google.common.collect.Lists; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Service; + +@Service +public class TaskServiceImpl implements TaskService { + + /** + * task converter + */ + @Autowired + private TaskConverter taskConverter; + + @Override + public DataResult create(TaskCreateParam param) { + TaskDO taskDO = taskConverter.todo(param); + taskDO.setDeleted(DeletedTypeEnum.N.name()); + taskDO.setTaskStatus(TaskStatusEnum.INIT.name()); + MapperUtils.getTaskMapper().insert(taskDO); + return DataResult.of(taskDO.getId()); + } + + @Override + public ActionResult updateStatus(TaskUpdateParam param) { + TaskDO taskDO = new TaskDO(); + taskDO.setId(param.getId()); + taskDO.setTaskStatus(param.getTaskStatus()); + taskDO.setContent(param.getContent()); + MapperUtils.getTaskMapper().updateById(taskDO); + return ActionResult.isSuccess(); + } + + @Override + public PageResult page(TaskPageParam param) { + Page page = new Page<>(); + page.setCurrent(param.getPageNo()); + page.setSize(param.getPageSize()); + page.setOrders(Lists.newArrayList(OrderItem.desc("gmt_create"))); + IPage iPage = MapperUtils.getTaskMapper().pageQuery(page, param.getUserId(), DeletedTypeEnum.N.name()); + if (iPage != null) { + return PageResult.of(taskConverter.toModel(iPage.getRecords()), param); + } + return PageResult.empty(param.getPageNo(), param.getPageSize()); + } + + @Override + public DataResult get(Long id) { + TaskDO task = MapperUtils.getTaskMapper().selectById(id); + return DataResult.of(taskConverter.toModel(task)); + } +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/MapperUtils.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/MapperUtils.java new file mode 100644 index 00000000..5f19d25f --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/MapperUtils.java @@ -0,0 +1,10 @@ +package ai.chat2db.server.domain.repository; + +import ai.chat2db.server.domain.repository.mapper.TaskMapper; + +public class MapperUtils { + + public static TaskMapper getTaskMapper() { + return Dbutils.getMapper(TaskMapper.class); + } +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/entity/TaskDO.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/entity/TaskDO.java new file mode 100644 index 00000000..9e4661a6 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/entity/TaskDO.java @@ -0,0 +1,101 @@ +package ai.chat2db.server.domain.repository.entity; + +import com.baomidou.mybatisplus.annotation.IdType; +import com.baomidou.mybatisplus.annotation.TableId; +import com.baomidou.mybatisplus.annotation.TableName; +import java.io.Serializable; +import java.util.Date; +import lombok.Getter; +import lombok.Setter; + +/** + *

+ * TASK TABLE + *

+ * + * @author chat2db + * @since 2024-01-25 + */ +@Getter +@Setter +@TableName("TASK") +public class TaskDO implements Serializable { + + private static final long serialVersionUID = 1L; + + /** + * 主键 + */ + @TableId(value = "ID", type = IdType.AUTO) + private Long id; + + /** + * 创建时间 + */ + private Date gmtCreate; + + /** + * 修改时间 + */ + private Date gmtModified; + + /** + * 数据源连接ID + */ + private Long dataSourceId; + + /** + * db名称 + */ + private String databaseName; + + /** + * schema名称 + */ + private String schemaName; + + /** + * table_name + */ + private String tableName; + + /** + * 是否被删除,y表示删除,n表示未删除 + */ + private String deleted; + + /** + * 用户id + */ + private Long userId; + + /** + * task type, such as: DOWNLOAD_DATA, UPLOAD_TABLE_DATA, DOWNLOAD_TABLE_STRUCTURE, UPLOAD_TABLE_STRUCTURE, + */ + private String taskType; + + /** + * task status + */ + private String taskStatus; + + /** + * task progress + */ + private String taskProgress; + + /** + * task name + */ + private String taskName; + + /** + * download url + */ + private String downloadUrl; + + /** + * task content + */ + private byte[] content; +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/mapper/TaskMapper.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/mapper/TaskMapper.java new file mode 100644 index 00000000..1a641ab8 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/mapper/TaskMapper.java @@ -0,0 +1,19 @@ +package ai.chat2db.server.domain.repository.mapper; + +import ai.chat2db.server.domain.repository.entity.TaskDO; +import com.baomidou.mybatisplus.core.mapper.BaseMapper; +import com.baomidou.mybatisplus.core.metadata.IPage; +import org.apache.ibatis.annotations.Param; + +/** + *

+ * TASK TABLE Mapper 接口 + *

+ * + * @author chat2db + * @since 2024-01-25 + */ +public interface TaskMapper extends BaseMapper { + + IPage pageQuery(IPage page, @Param("userId") Long userId,@Param("deleted") String deleted); +} diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/db/migration/V2_1_9__task.sql b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/db/migration/V2_1_9__task.sql new file mode 100644 index 00000000..326972a2 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/db/migration/V2_1_9__task.sql @@ -0,0 +1,20 @@ +CREATE TABLE IF NOT EXISTS `task` ( + `id` bigint(20) unsigned NOT NULL AUTO_INCREMENT COMMENT '主键', + `gmt_create` datetime NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间', + `gmt_modified` datetime NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '修改时间', + `data_source_id` bigint(20) unsigned NULL COMMENT '数据源连接ID', + `database_name` varchar(128) DEFAULT NULL COMMENT 'db名称', + `schema_name` varchar(128) DEFAULT NULL COMMENT 'schema名称', + `table_name` varchar(128) DEFAULT NULL COMMENT 'table_name', + `deleted` varchar(10) DEFAULT NULL COMMENT '是否被删除,y表示删除,n表示未删除', + `user_id` bigint(20) unsigned NOT NULL DEFAULT 0 COMMENT '用户id', + `task_type` varchar(128) DEFAULT NULL COMMENT 'task type, such as: DOWNLOAD_DATA, UPLOAD_TABLE_DATA, DOWNLOAD_TABLE_STRUCTURE, UPLOAD_TABLE_STRUCTURE,', + `task_status` varchar(128) DEFAULT NULL COMMENT 'task status', + `task_progress` varchar(128) DEFAULT NULL COMMENT 'task progress', + `task_name` varchar(128) DEFAULT NULL COMMENT 'task name', + `content` blob DEFAULT NULL COMMENT 'task content', + `download_url` varchar(512) DEFAULT NULL COMMENT 'down load url', + PRIMARY KEY (`id`) + ) ENGINE=InnoDB AUTO_INCREMENT=1 DEFAULT CHARSET=utf8mb4 COMMENT='TASK TABLE' +; +create INDEX idx_task_user_id on task(user_id) ; \ No newline at end of file diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/mapper/TaskMapper.xml b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/mapper/TaskMapper.xml new file mode 100644 index 00000000..498d1b6d --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/mapper/TaskMapper.xml @@ -0,0 +1,36 @@ + + + + + + diff --git a/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/Application.java b/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/Application.java index 573a6105..43ab662c 100644 --- a/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/Application.java +++ b/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/Application.java @@ -11,7 +11,6 @@ import org.springframework.scheduling.annotation.EnableAsync; import org.springframework.scheduling.annotation.EnableScheduling; import org.springframework.stereotype.Indexed; -import java.util.concurrent.CompletableFuture; /** * 启动类 diff --git a/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/StdinReader.java b/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/StdinReader.java deleted file mode 100644 index 3505849c..00000000 --- a/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/StdinReader.java +++ /dev/null @@ -1,30 +0,0 @@ -//package ai.chat2db.server.start.config; -// -//import jakarta.annotation.PostConstruct; -//import org.springframework.boot.context.event.ApplicationReadyEvent; -//import org.springframework.context.ApplicationListener; -//import org.springframework.stereotype.Component; -// -//import java.util.Scanner; -// -//@Component -//public class StdinReader implements ApplicationListener { -// -// -// @Override -// public void onApplicationEvent(ApplicationReadyEvent event) { -// // 启动读取 stdin 的线程 -// new Thread(() -> readStdin()).start(); -// } -// -// private void readStdin() { -// Scanner scanner = new Scanner(System.in); -// while (scanner.hasNextLine()) { -// String line = scanner.nextLine(); -// // 处理接收到的数据 -// System.out.println("接收到数据: " + line); -// // 在这里调用其他服务或逻辑 -// } -// } -// -//} diff --git a/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/oauth/SaLogForSlf4j.java b/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/oauth/SaLogForSlf4j.java deleted file mode 100644 index f45df2fe..00000000 --- a/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/oauth/SaLogForSlf4j.java +++ /dev/null @@ -1,44 +0,0 @@ -//package ai.chat2db.server.start.config.oauth; -// -//import cn.dev33.satoken.log.SaLog; -//import lombok.extern.slf4j.Slf4j; -//import org.springframework.stereotype.Component; -// -///** -// * satoken 日志打印 -// * -// * @author 是仪 -// */ -//@Slf4j -//@Component -//public class SaLogForSlf4j implements SaLog { -// @Override -// public void trace(String str, Object... args) { -// log.trace(str, args); -// } -// -// @Override -// public void debug(String str, Object... args) { -// log.debug(str, args); -// } -// -// @Override -// public void info(String str, Object... args) { -// log.info(str, args); -// } -// -// @Override -// public void warn(String str, Object... args) { -// log.trace(str, args); -// } -// -// @Override -// public void error(String str, Object... args) { -// log.error(str, args); -// } -// -// @Override -// public void fatal(String str, Object... args) { -// log.error(str, args); -// } -//} diff --git a/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/oauth/SaTokenConfigure.java b/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/oauth/SaTokenConfigure.java deleted file mode 100644 index 66663a8f..00000000 --- a/chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/config/oauth/SaTokenConfigure.java +++ /dev/null @@ -1,21 +0,0 @@ -//package ai.chat2db.server.start.config.oauth; -// -//import cn.dev33.satoken.jwt.StpLogicJwtForStateless; -//import cn.dev33.satoken.stp.StpLogic; -//import org.springframework.context.annotation.Bean; -//import org.springframework.context.annotation.Configuration; -// -///** -// * satoken配置 -// * -// * @author 是仪 -// */ -//@Configuration -//public class SaTokenConfigure { -// @Bean -// public StpLogic ttpLogic() { -// // 登录展示用 无状态的 这样不用依赖于与redis之类的 -// // 后续可以改成ehcahe 存储磁盘? -// return new StpLogicJwtForStateless(); -// } -//} \ No newline at end of file diff --git a/chat2db-server/chat2db-server-start/src/test/java/ai/chat2db/server/start/test/mybatis/MybatisGeneratorTest.java b/chat2db-server/chat2db-server-start/src/test/java/ai/chat2db/server/start/test/mybatis/MybatisGeneratorTest.java index e07c3c10..09615e5e 100644 --- a/chat2db-server/chat2db-server-start/src/test/java/ai/chat2db/server/start/test/mybatis/MybatisGeneratorTest.java +++ b/chat2db-server/chat2db-server-start/src/test/java/ai/chat2db/server/start/test/mybatis/MybatisGeneratorTest.java @@ -14,8 +14,10 @@ import com.baomidou.mybatisplus.generator.config.converts.MySqlTypeConvert; import com.baomidou.mybatisplus.generator.config.rules.DateType; import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine; import com.google.common.collect.Lists; +import com.zaxxer.hikari.HikariDataSource; import jakarta.annotation.Resource; import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; import org.junit.jupiter.api.Test; /** @@ -30,12 +32,30 @@ public class MybatisGeneratorTest extends BaseTest { @Test public void coreGenerator() { + + HikariDataSource dataSource = new HikariDataSource(); + String environment = StringUtils.defaultString(System.getProperty("spring.profiles.active"), "dev"); + if ("dev".equalsIgnoreCase(environment)) { + dataSource.setJdbcUrl("jdbc:h2:file:~/.chat2db/db/chat2db_dev;MODE=MYSQL"); + }else if ("test".equalsIgnoreCase(environment)) { + dataSource.setJdbcUrl("jdbc:h2:file:~/.chat2db/db/chat2db_test;MODE=MYSQL"); + }else { + dataSource.setJdbcUrl("jdbc:h2:~/.chat2db/db/chat2db;MODE=MYSQL;FILE_LOCK=NO"); + } + dataSource.setDriverClassName("org.h2.Driver"); + dataSource.setIdleTimeout(60000); + dataSource.setAutoCommit(true); + dataSource.setMaximumPoolSize(500); + dataSource.setMinimumIdle(1); + dataSource.setMaxLifetime(60000 * 10); + dataSource.setConnectionTestQuery("SELECT 1"); + this.dataSource = dataSource; //doGenerator(Lists.newArrayList("data_source")); //doGenerator(Lists.newArrayList("operation_log")); //doGenerator(Lists.newArrayList("operation_saved")); //doGenerator(Lists.newArrayList("environment","data_source","team","team_dbhub_user","data_source_access", // "dbhub_user")); - doGenerator(Lists.newArrayList("operation_log")); + doGenerator(Lists.newArrayList("TASK")); } private void doGenerator(List tableList) { diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlController.java index 2197d84e..5b3962f3 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlController.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlController.java @@ -70,33 +70,9 @@ public class RdbDmlController { DlExecuteParam param = rdbWebConverter.request2param(request); ListResult resultDTOListResult = dlTemplateService.execute(param); List resultVOS = rdbWebConverter.dto2vo(resultDTOListResult.getData()); - String type = Chat2DBContext.getConnectInfo().getDbType(); - String clientId = getClientId(); - String sqlContent = request.getSql(); - executorService.submit(() -> { - try { - addOperationLog(clientId, type, sqlContent, resultDTOListResult.getErrorMessage(), resultDTOListResult.getSuccess(), resultVOS); - } catch (Exception e) { - // do nothing - } - }); return ListResult.of(resultVOS); } - private void addOperationLog(String clientId, String sqlType, String sqlContent, String errorMessage, Boolean isSuccess, List executeResultVOS) { - SqlExecuteHistoryCreateRequest createRequest = new SqlExecuteHistoryCreateRequest(); - createRequest.setClientId(clientId); - createRequest.setErrorMessage(errorMessage); - createRequest.setDatabaseType(sqlType); - createRequest.setSqlContent(sqlContent); - createRequest.setExecuteStatus(isSuccess ? "success" : "fail"); - executeResultVOS.forEach(executeResultVO -> { - createRequest.setSqlType(executeResultVO.getSqlType()); - createRequest.setDuration(executeResultVO.getDuration()); - createRequest.setTableName(executeResultVO.getTableName()); - gatewayClientService.addOperationLog(createRequest); - }); - } /** * query chat2db apikey @@ -148,16 +124,6 @@ public class RdbDmlController { return DataResult.error(result.getErrorCode(), result.getErrorMessage()); } ExecuteResultVO executeResultVO = rdbWebConverter.dto2vo(result.getData()); - String type = Chat2DBContext.getConnectInfo().getDbType(); - String sqlContent = request.getSql(); - String clientId = getClientId(); - executorService.submit(() -> { - try { - addOperationLog(clientId, type, sqlContent, result.getErrorMessage(), result.getSuccess(), Lists.newArrayList(executeResultVO)); - } catch (Exception e) { - // do nothing - } - }); return DataResult.of(executeResultVO); } diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/ExportController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/ExportController.java new file mode 100644 index 00000000..2c0d01e7 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/ExportController.java @@ -0,0 +1,42 @@ +package ai.chat2db.server.web.api.controller.task; + +import ai.chat2db.server.tools.base.wrapper.result.DataResult; +import ai.chat2db.server.web.api.aspect.ConnectionInfoAspect; +import ai.chat2db.server.web.api.controller.rdb.request.DataExportRequest; +import ai.chat2db.server.web.api.controller.task.biz.TaskBizService; +import jakarta.validation.Valid; +import lombok.extern.slf4j.Slf4j; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Controller; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RequestMapping; + +@ConnectionInfoAspect +@RequestMapping("/api/export") +@Controller +@Slf4j +public class ExportController { + + @Autowired + private TaskBizService taskBizService; + + + /** + * export data + * + * @param request + * @return + */ + @PostMapping("/export_data") + public DataResult export(@Valid @RequestBody DataExportRequest request) { + return taskBizService.exportResultData(request); + } + + @PostMapping("/export_doc") + public DataResult exportDoc(@Valid @RequestBody DataExportRequest request) { + return taskBizService.exportSchemaDoc(request); + } + + +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/TaskController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/TaskController.java new file mode 100644 index 00000000..b6c78411 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/TaskController.java @@ -0,0 +1,77 @@ +package ai.chat2db.server.web.api.controller.task; + +import ai.chat2db.server.domain.api.model.Task; +import ai.chat2db.server.domain.api.param.TaskPageParam; +import ai.chat2db.server.domain.api.service.TaskService; +import ai.chat2db.server.tools.base.wrapper.result.DataResult; +import ai.chat2db.server.tools.base.wrapper.result.PageResult; +import ai.chat2db.server.tools.base.wrapper.result.web.WebPageResult; +import ai.chat2db.server.tools.common.util.ContextUtils; +import ai.chat2db.server.web.api.aspect.ConnectionInfoAspect; +import lombok.extern.slf4j.Slf4j; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.core.io.Resource; +import org.springframework.core.io.UrlResource; +import org.springframework.http.HttpHeaders; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.stereotype.Controller; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.PathVariable; +import org.springframework.web.bind.annotation.RequestMapping; + +import java.net.MalformedURLException; + +@ConnectionInfoAspect +@RequestMapping("/api/task") +@Controller +@Slf4j +public class TaskController { + + @Autowired + private TaskService taskService; + + + @GetMapping("/list") + public WebPageResult list() { + TaskPageParam taskPageParam = new TaskPageParam(); + taskPageParam.setPageNo(1); + taskPageParam.setPageSize(10); + taskPageParam.setUserId(ContextUtils.getUserId()); + PageResult task = taskService.page(taskPageParam); + return WebPageResult.of(task.getData(), 100L, 1, 10); + } + + @GetMapping("/download/{id}") + public ResponseEntity download(@PathVariable Long id) { + DataResult task = taskService.get(id); + if(task.getData() == null){ + log.error("task is null"); + throw new RuntimeException("task is null"); + } + if(ContextUtils.getUserId() != task.getData().getUserId()){ + log.error("task is not belong to user"); + throw new RuntimeException("task is not belong to user"); + } + + Resource resource = null; + try { + resource = new UrlResource("file://"+task.getData().getDownloadUrl()); + } catch (MalformedURLException e) { + throw new RuntimeException(e); + } + + if (resource.exists() || resource.isReadable()) { + return ResponseEntity.ok() + .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + resource.getFilename() + "\"") + .contentType(MediaType.APPLICATION_OCTET_STREAM) + .body(resource); + } else { + throw new RuntimeException("Could not read the file!"); + } + + } + + + +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/biz/TaskBizService.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/biz/TaskBizService.java new file mode 100644 index 00000000..03da5b96 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/biz/TaskBizService.java @@ -0,0 +1,312 @@ +package ai.chat2db.server.web.api.controller.task.biz; + +import ai.chat2db.server.domain.api.enums.*; +import ai.chat2db.server.domain.api.param.*; +import ai.chat2db.server.domain.api.service.TableService; +import ai.chat2db.server.domain.api.service.TaskService; +import ai.chat2db.server.domain.repository.Dbutils; +import ai.chat2db.server.tools.base.excption.BusinessException; +import ai.chat2db.server.tools.base.wrapper.result.DataResult; +import ai.chat2db.server.tools.base.wrapper.result.PageResult; +import ai.chat2db.server.tools.common.model.Context; +import ai.chat2db.server.tools.common.model.LoginUser; +import ai.chat2db.server.tools.common.util.ContextUtils; +import ai.chat2db.server.tools.common.util.EasyCollectionUtils; +import ai.chat2db.server.web.api.controller.rdb.RdbDmlExportController; +import ai.chat2db.server.web.api.controller.rdb.converter.RdbWebConverter; +import ai.chat2db.server.web.api.controller.rdb.doc.DatabaseExportService; +import ai.chat2db.server.web.api.controller.rdb.doc.conf.ExportOptions; +import ai.chat2db.server.web.api.controller.rdb.factory.ExportServiceFactory; +import ai.chat2db.server.web.api.controller.rdb.request.DataExportRequest; +import ai.chat2db.server.web.api.controller.rdb.vo.TableVO; +import ai.chat2db.spi.jdbc.DefaultValueHandler; +import ai.chat2db.spi.model.Table; +import ai.chat2db.spi.sql.Chat2DBContext; +import ai.chat2db.spi.sql.ConnectInfo; +import ai.chat2db.spi.sql.SQLExecutor; +import ai.chat2db.spi.util.JdbcUtils; +import ai.chat2db.spi.util.SqlUtils; +import cn.hutool.core.date.DatePattern; +import cn.hutool.core.io.FileUtil; +import cn.hutool.core.lang.Assert; +import com.alibaba.druid.DbType; +import com.alibaba.druid.sql.SQLUtils; +import com.alibaba.druid.sql.ast.SQLStatement; +import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; +import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; +import com.alibaba.druid.sql.ast.statement.SQLInsertStatement; +import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; +import com.alibaba.druid.sql.visitor.VisitorFeature; +import com.alibaba.excel.EasyExcel; +import com.alibaba.excel.ExcelWriter; +import com.alibaba.excel.support.ExcelTypeEnum; +import com.alibaba.excel.write.builder.ExcelWriterBuilder; +import com.alibaba.excel.write.metadata.WriteSheet; +import com.google.common.collect.Lists; +import lombok.AllArgsConstructor; +import lombok.Data; +import lombok.NoArgsConstructor; +import lombok.experimental.SuperBuilder; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Component; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.io.PrintWriter; +import java.lang.reflect.Constructor; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.time.LocalDateTime; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +@Slf4j +@Component +public class TaskBizService { + + /** + * Format insert statement + */ + private static final SQLUtils.FormatOption INSERT_FORMAT_OPTION = new SQLUtils.FormatOption(true, false); + + static { + INSERT_FORMAT_OPTION.config(VisitorFeature.OutputNameQuote, true); + } + + + @Autowired + private TaskService taskService; + + + @Autowired + private TableService tableService; + + @Autowired + private RdbWebConverter rdbWebConverter; + + public DataResult exportResultData(DataExportRequest request) { + String sql = ExportSizeEnum.CURRENT_PAGE.getCode().equals(request.getExportSize()) ? request.getSql() : request.getOriginalSql(); + Assert.notBlank(sql, "dataSource.sqlEmpty"); + DbType dbType = JdbcUtils.parse2DruidDbType(Chat2DBContext.getConnectInfo().getDbType()); + String tableName = getTableName(request, sql, dbType); + File file = createTempFile(tableName, request.getExportType()); + DataResult dataResult = createTask(tableName, request.getDatabaseName(), request.getSchemaName(), request.getDataSourceId(), tableName); + + LoginUser loginUser = ContextUtils.getLoginUser(); + ConnectInfo connectInfo = Chat2DBContext.getConnectInfo().copy(); + CompletableFuture.runAsync(() -> { + buildContext(loginUser, connectInfo); + doExport(sql, file, dbType, tableName, request.getExportType()); + }).whenComplete((aVoid, throwable) -> { + updateStatus(dataResult.getData(), file, throwable); + removeContext(); + }); + return dataResult; + } + + public DataResult exportSchemaDoc(DataExportRequest request) { + File file = createTempFile(request.getDatabaseName(), request.getExportType()); + DataResult dataResult = createTask(null, request.getDatabaseName(), request.getSchemaName(), request.getDataSourceId(), "schema_doc"); + LoginUser loginUser = ContextUtils.getLoginUser(); + ConnectInfo connectInfo = Chat2DBContext.getConnectInfo().copy(); + CompletableFuture.runAsync(() -> { + buildContext(loginUser, connectInfo); + doExportDoc(request, file); + }).whenComplete((aVoid, throwable) -> { + updateStatus(dataResult.getData(), file, throwable); + removeContext(); + }); + return dataResult; + } + + private void doExportDoc(DataExportRequest request, File file) { + try { + TablePageQueryParam queryParam = rdbWebConverter.tablePageRequest2param(request); + queryParam.setPageNo(1); + queryParam.setPageSize(Integer.MAX_VALUE); + TableSelector tableSelector = new TableSelector(); + tableSelector.setColumnList(true); + tableSelector.setIndexList(true); + PageResult tableDTOPageResult = tableService.pageQuery(queryParam, tableSelector); + List tableVOS = rdbWebConverter.tableDto2vo(tableDTOPageResult.getData()); + TableQueryParam param = rdbWebConverter.tableRequest2param(request); + for (TableVO tableVO : tableVOS) { + param.setTableName(tableVO.getName()); + tableVO.setColumnList(tableService.queryColumns(param)); + tableVO.setIndexList(tableService.queryIndexes(param)); + } + Class targetClass = ExportServiceFactory.get(request.getExportType()); + Constructor constructor = targetClass.getDeclaredConstructor(); + DatabaseExportService databaseExportService = (DatabaseExportService) constructor.newInstance(); + // 设置数据集合 + databaseExportService.setExportList(tableVOS); + databaseExportService.generate(request.getDatabaseName(), new FileOutputStream(file), new ExportOptions()); + } catch (Exception e) { + log.error("export error", e); + throw new BusinessException("dataSource.exportError"); + } + } + + + private void removeContext() { + Dbutils.removeSession(); + ContextUtils.removeContext(); + Chat2DBContext.removeContext(); + } + + private void buildContext(LoginUser loginUser, ConnectInfo connectInfo) { + ContextUtils.setContext(Context.builder() + .loginUser(loginUser) + .build()); + Dbutils.setSession(); + Chat2DBContext.putContext(connectInfo); + } + + private DataResult createTask(String tableName, String databaseName, String schemaName, Long datasourceId, String taskName) { + TaskCreateParam param = new TaskCreateParam(); + param.setTaskName("export_" + taskName); + param.setTaskType(TaskTypeEnum.DOWNLOAD_TABLE_DATA.name()); + param.setDatabaseName(databaseName); + param.setSchemaName(schemaName); + param.setTableName(tableName); + param.setDataSourceId(datasourceId); + param.setUserId(ContextUtils.getUserId()); + param.setTaskProgress("0.1"); + return taskService.create(param); + } + + private void updateStatus(Long id, File file, Throwable throwable) { + TaskUpdateParam updateParam = new TaskUpdateParam(); + updateParam.setId(id); + updateParam.setTaskProgress("1"); + updateParam.setDownloadUrl(file.getAbsolutePath()); + if (throwable != null) { + log.error("export error", throwable); + updateParam.setTaskStatus(TaskStatusEnum.ERROR.name()); + } else { + updateParam.setTaskStatus(TaskStatusEnum.FINISH.name()); + } + taskService.updateStatus(updateParam); + } + + private void doExport(String sql, File file, DbType dbType, String tableName, String exportType) { + try { + if (ExportTypeEnum.CSV.getCode().equals(exportType)) { + doExportCsv(sql, file); + } else { + doExportInsert(sql, file, dbType, tableName); + } + } catch (Exception e) { + log.error("export error", e); + throw new BusinessException("dataSource.exportError"); + } + } + + + private File createTempFile(String tableName, String exportType) { + String fileName = URLEncoder.encode( + tableName + "_" + LocalDateTime.now().format(DatePattern.PURE_DATETIME_FORMATTER), + StandardCharsets.UTF_8) + .replaceAll("\\+", "%20"); + + if (ExportTypeEnum.CSV.getCode().equals(exportType)) { + return FileUtil.createTempFile(fileName, ".csv", true); + } else if (ExportTypeEnum.INSERT.getCode().equals(exportType)) { + return FileUtil.createTempFile(fileName, ".sql", true); + } else if (ExportTypeEnum.EXCEL.getCode().equals(exportType)) { + return FileUtil.createTempFile(fileName, ExportFileSuffix.EXCEL.getSuffix(), true); + } else if (ExportTypeEnum.MARKDOWN.getCode().equals(exportType)) { + return FileUtil.createTempFile(fileName, ExportFileSuffix.MARKDOWN.getSuffix(), true); + } else if (ExportTypeEnum.WORD.getCode().equals(exportType)) { + return FileUtil.createTempFile(fileName, ExportFileSuffix.WORD.getSuffix(), true); + } else if (ExportTypeEnum.PDF.getCode().equals(exportType)) { + return FileUtil.createTempFile(fileName, ExportFileSuffix.PDF.getSuffix(), true); + } else if (ExportTypeEnum.HTML.getCode().equals(exportType)) { + return FileUtil.createTempFile(fileName, ExportFileSuffix.HTML.getSuffix(), true); + } + return FileUtil.createTempFile(fileName, ".txt", true); + } + + private String getTableName(DataExportRequest request, String sql, DbType dbType) { + String tableName = null; + if (dbType != null) { + SQLStatement sqlStatement = SQLUtils.parseSingleStatement(sql, dbType); + if (!(sqlStatement instanceof SQLSelectStatement)) { + throw new BusinessException("dataSource.sqlAnalysisError"); + } + tableName = SqlUtils.getTableName(sql, dbType); + } else { + tableName = StringUtils.join(Lists.newArrayList(request.getDatabaseName(), request.getSchemaName()), "_"); + } + return tableName; + } + + private void doExportCsv(String sql, File file) { + RdbDmlExportController.ExcelWrapper excelWrapper = new RdbDmlExportController.ExcelWrapper(); + try { + ExcelWriterBuilder excelWriterBuilder = EasyExcel.write(file) + .charset(StandardCharsets.UTF_8) + .excelType(ExcelTypeEnum.CSV); + excelWrapper.setExcelWriterBuilder(excelWriterBuilder); + SQLExecutor.getInstance().execute(Chat2DBContext.getConnection(), sql, headerList -> { + excelWriterBuilder.head( + EasyCollectionUtils.toList(headerList, header -> Lists.newArrayList(header.getName()))); + excelWrapper.setExcelWriter(excelWriterBuilder.build()); + excelWrapper.setWriteSheet(EasyExcel.writerSheet(0).build()); + }, dataList -> { + List> writeDataList = Lists.newArrayList(); + writeDataList.add(dataList); + excelWrapper.getExcelWriter().write(writeDataList, excelWrapper.getWriteSheet()); + }, false, new DefaultValueHandler()); + } finally { + if (excelWrapper.getExcelWriter() != null) { + excelWrapper.getExcelWriter().finish(); + } + } + } + + private void doExportInsert(String sql, File file, DbType dbType, + String tableName) + throws IOException { + try (PrintWriter printWriter = new PrintWriter(file, StandardCharsets.UTF_8.name())) { + RdbDmlExportController.InsertWrapper insertWrapper = new RdbDmlExportController.InsertWrapper(); + SQLExecutor.getInstance().execute(Chat2DBContext.getConnection(), sql, + headerList -> insertWrapper.setHeaderList( + EasyCollectionUtils.toList(headerList, header -> new SQLIdentifierExpr(header.getName()))) + , dataList -> { + SQLInsertStatement sqlInsertStatement = new SQLInsertStatement(); + sqlInsertStatement.setDbType(dbType); + sqlInsertStatement.setTableSource(new SQLExprTableSource(tableName)); + sqlInsertStatement.getColumns().addAll(insertWrapper.getHeaderList()); + SQLInsertStatement.ValuesClause valuesClause = new SQLInsertStatement.ValuesClause(); + for (String s : dataList) { + valuesClause.addValue(s); + } + sqlInsertStatement.setValues(valuesClause); + + printWriter.println(SQLUtils.toSQLString(sqlInsertStatement, dbType, INSERT_FORMAT_OPTION) + ";"); + }, false, new DefaultValueHandler()); + } + } + + @Data + @SuperBuilder + @NoArgsConstructor + @AllArgsConstructor + public static class InsertWrapper { + private List headerList; + } + + @Data + @SuperBuilder + @NoArgsConstructor + @AllArgsConstructor + public static class ExcelWrapper { + private ExcelWriterBuilder excelWriterBuilder; + private ExcelWriter excelWriter; + private WriteSheet writeSheet; + } +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/http/GatewayClientService.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/http/GatewayClientService.java index 78b72e96..14abe51a 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/http/GatewayClientService.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/http/GatewayClientService.java @@ -215,14 +215,5 @@ public class GatewayClientService { // return result; } - public ActionResult addOperationLog(SqlExecuteHistoryCreateRequest request) { - ActionResult result = Forest.post(chat2dbProperties.getGateway().getBaseUrl() + "/api/client/sql/execute/history") - .connectTimeout(Duration.ofMillis(5000)) - .readTimeout(Duration.ofMillis(10000)) - .contentType("application/json") - .addBody(request) - .execute(new TypeReference<>() { - }); - return result; - } + } diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/ws/WsService.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/ws/WsService.java index e2f575ab..a9c9b4eb 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/ws/WsService.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/ws/WsService.java @@ -16,17 +16,13 @@ import ai.chat2db.server.web.api.controller.ai.chat2db.client.Chat2dbAIClient; import ai.chat2db.server.web.api.controller.rdb.converter.RdbWebConverter; import ai.chat2db.server.web.api.controller.rdb.request.DmlRequest; import ai.chat2db.server.web.api.controller.rdb.vo.ExecuteResultVO; -import ai.chat2db.server.web.api.http.GatewayClientService; -import ai.chat2db.server.web.api.http.request.SqlExecuteHistoryCreateRequest; import ai.chat2db.server.web.api.util.ApplicationContextUtil; import ai.chat2db.spi.config.DriverConfig; import ai.chat2db.spi.model.ExecuteResult; -import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.ConnectInfo; import org.apache.commons.lang3.StringUtils; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Component; -import org.springframework.web.bind.annotation.RequestBody; import java.util.List; import java.util.Objects; @@ -53,26 +49,12 @@ public class WsService { @Autowired private DlTemplateService dlTemplateService; - @Autowired - private GatewayClientService gatewayClientService; - - public static ExecutorService executorService = Executors.newFixedThreadPool(10); public ListResult execute(DmlRequest request) { DlExecuteParam param = rdbWebConverter.request2param(request); ListResult resultDTOListResult = dlTemplateService.execute(param); List resultVOS = rdbWebConverter.dto2vo(resultDTOListResult.getData()); - String type = Chat2DBContext.getConnectInfo().getDbType(); - String clientId = getApiKey(); - String sqlContent = request.getSql(); - executorService.submit(() -> { - try { - addOperationLog(clientId, type, sqlContent, resultDTOListResult.getErrorMessage(), resultDTOListResult.getSuccess(), resultVOS); - } catch (Exception e) { - // do nothing - } - }); return ListResult.of(resultVOS); } @@ -143,20 +125,5 @@ public class WsService { return keyConfig.getContent(); } - private void addOperationLog(String clientId, String sqlType, String sqlContent, String errorMessage, Boolean isSuccess, List executeResultVOS) { - SqlExecuteHistoryCreateRequest createRequest = new SqlExecuteHistoryCreateRequest(); - createRequest.setClientId(clientId); - createRequest.setErrorMessage(errorMessage); - createRequest.setDatabaseType(sqlType); - createRequest.setSqlContent(sqlContent); - createRequest.setExecuteStatus(isSuccess ? "success" : "fail"); - executeResultVOS.forEach(executeResultVO -> { - createRequest.setSqlType(executeResultVO.getSqlType()); - createRequest.setDuration(executeResultVO.getDuration()); - createRequest.setTableName(executeResultVO.getTableName()); - gatewayClientService.addOperationLog(createRequest); - }); - } - } diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/ConnectInfo.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/ConnectInfo.java index eb18d766..41e18ed5 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/ConnectInfo.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/ConnectInfo.java @@ -522,4 +522,30 @@ public class ConnectInfo { public void setSchemaName(String schemaName) { this.schemaName = schemaName; } + + public ConnectInfo copy(){ + ConnectInfo copy = new ConnectInfo(); + copy.setDbVersion(this.getDbVersion()); + copy.setDbType(this.getDbType()); + copy.setHost(this.getHost()); + copy.setPort(this.getPort()); + copy.setDatabaseName(this.getDatabaseName()); + copy.setSchemaName(this.getSchemaName()); + copy.setUser(this.getUser()); + copy.setPassword(this.getPassword()); + copy.setUrl(this.getUrl()); + copy.setAlias(this.getAlias()); + copy.setDataSourceId(this.getDataSourceId()); + copy.setConsoleId(this.getConsoleId()); + copy.setConsoleOwn(this.getConsoleOwn()); + copy.setDriver(this.getDriver()); + copy.setSsh(this.getSsh()); + copy.setSsl(this.getSsl()); + copy.setJdbc(this.getJdbc()); + copy.setExtendInfo(this.getExtendInfo()); + copy.setDriverConfig(this.getDriverConfig()); + copy.setSid(this.getSid()); + copy.setUrlWithOutDatabase(this.getUrlWithOutDatabase()); + return copy; + } } \ No newline at end of file