From cd16555c45afb948e27663fba68e91935a05d21a Mon Sep 17 00:00:00 2001 From: zgq <203083679@qq.com> Date: Tue, 18 Jun 2024 13:56:25 +0800 Subject: [PATCH] data export --- .../plugin/mysql/builder/MysqlSqlBuilder.java | 14 +- .../domain/api/enums/ExportFileSuffix.java | 1 + .../datasource/DatabaseExportDataParam.java | 7 + .../domain/core/impl/TaskServiceImpl.java | 1 + .../tools/common/util/EasyStringUtils.java | 12 +- .../controller/rdb/DatabaseController.java | 51 ++--- .../controller/rdb/data/BaseDataExporter.java | 70 +++++++ .../controller/rdb/data/BaseDataImporter.java | 8 + .../rdb/data/BaseExcelExporter.java | 105 ++++++++++ .../rdb/data/BaseExcelImporter.java | 10 + .../rdb/data/DataExportStrategy.java | 13 ++ .../rdb/data/DataImportStrategy.java | 4 + .../rdb/data/csv/CsvDataExporter.java | 26 +++ .../rdb/data/csv/CsvDataImporter.java | 12 ++ .../strategy/ExportDBData2CsvStrategy.java | 54 ----- .../strategy/ExportDBData2ExcelStrategy.java | 53 ----- .../strategy/ExportDBData2JsonStrategy.java | 61 ------ .../strategy/ExportDBData2SqlStrategy.java | 33 ---- .../export/strategy/ExportDBDataStrategy.java | 65 ------- .../rdb/data/factory/DataExportFactory.java | 33 ++++ .../rdb/data/factory/DataImportFactory.java | 35 ++++ .../rdb/data/json/JsonDataExporter.java | 112 +++++++++++ .../rdb/data/json/JsonDataImporter.java | 12 ++ .../rdb/data/service/DatabaseDataService.java | 13 ++ .../data/service/impl/DatabaseDataImpl.java | 150 ++++++++++++++ .../rdb/data/sql/SqlDataExporter.java | 184 ++++++++++++++++++ .../controller/rdb/data/task/TaskManager.java | 69 +++++++ .../controller/rdb/data/task/TaskState.java | 28 +++ .../rdb/data/xls/XlsDataExporter.java | 24 +++ .../rdb/data/xls/XlsDataImporter.java | 13 ++ .../rdb/data/xlsx/XlsxDataExporter.java | 25 +++ .../rdb/data/xlsx/XlsxDataImporter.java | 12 ++ .../factory/ExportDBDataStrategyFactory.java | 31 --- .../request/DatabaseExportDataRequest.java | 10 + .../api/controller/task/TaskController.java | 47 +++-- .../main/java/ai/chat2db/spi/SqlBuilder.java | 19 +- .../chat2db/spi/jdbc/DefaultSqlBuilder.java | 70 +++++-- .../java/ai/chat2db/spi/sql/SQLExecutor.java | 16 ++ 38 files changed, 1133 insertions(+), 370 deletions(-) create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataExporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataImporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelExporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelImporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataExportStrategy.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataImportStrategy.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataExporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataImporter.java delete mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2CsvStrategy.java delete mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2ExcelStrategy.java delete mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2JsonStrategy.java delete mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2SqlStrategy.java delete mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBDataStrategy.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataExportFactory.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataImportFactory.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataExporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataImporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/DatabaseDataService.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/impl/DatabaseDataImpl.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/sql/SqlDataExporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskManager.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskState.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataExporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataImporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataExporter.java create mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataImporter.java delete mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/factory/ExportDBDataStrategyFactory.java diff --git a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilder.java b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilder.java index 06a2ad37a..9a2ad6c7c 100644 --- a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilder.java +++ b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilder.java @@ -385,11 +385,11 @@ private static String[] moveElement(String[] originalArray, int from, int to, St } -// @Override -// protected void buildTableName(String databaseName, String schemaName, String tableName, StringBuilder script) { -// if (StringUtils.isNotBlank(databaseName)) { -// script.append("`").append(databaseName).append("`").append('.'); -// } -// script.append("`").append(tableName).append("`"); -// } + @Override + protected void buildTableName(String databaseName, String schemaName, String tableName, StringBuilder script) { + if (StringUtils.isNotBlank(databaseName)) { + script.append("`").append(databaseName).append("`").append('.'); + } + script.append("`").append(tableName).append("`"); + } } diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/ExportFileSuffix.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/ExportFileSuffix.java index 9b5d854ee..4b55204d3 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/ExportFileSuffix.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/ExportFileSuffix.java @@ -14,6 +14,7 @@ public enum ExportFileSuffix { WORD(".docx"), //excel EXCEL(".xlsx"), + XLS(".xls"), //markdown MARKDOWN(".md"), //html diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/datasource/DatabaseExportDataParam.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/datasource/DatabaseExportDataParam.java index 4e562cedf..8f60eca06 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/datasource/DatabaseExportDataParam.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/datasource/DatabaseExportDataParam.java @@ -4,6 +4,8 @@ import lombok.Data; import lombok.NoArgsConstructor; +import java.util.List; + /** * @author: zgq * @date: 2024年03月24日 13:17 @@ -12,7 +14,12 @@ @AllArgsConstructor @NoArgsConstructor public class DatabaseExportDataParam { + private Long dataSourceId; private String databaseName; private String schemaName; private String exportType; + private List tableNames; + private String sqyType; + private Boolean containsHeader; + } \ No newline at end of file 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 index 25ba9a9fe..42f3f8fcc 100644 --- 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 @@ -44,6 +44,7 @@ public ActionResult updateStatus(TaskUpdateParam param) { taskDO.setId(param.getId()); taskDO.setTaskStatus(param.getTaskStatus()); taskDO.setContent(param.getContent()); + taskDO.setDownloadUrl(param.getDownloadUrl()); MapperUtils.getTaskMapper().updateById(taskDO); return ActionResult.isSuccess(); } diff --git a/chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyStringUtils.java b/chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyStringUtils.java index 8033c61e2..7ae6d796b 100644 --- a/chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyStringUtils.java +++ b/chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyStringUtils.java @@ -193,17 +193,17 @@ public static String escapeAndQuoteString(String value) { } /** - * @param value "abcd" - * @param quoteChar '%' - * @return "%abcd%" + * @param value "abcd" + * @param quoteChar '%' + * @return "%abcd%" */ public static String quoteString(String value, char quoteChar) { - return StringUtils.wrap(value, quoteChar); + return quoteChar + value + quoteChar; } /** - * @param value "abcd" - * @return "'abcd'" + * @param value "abcd" + * @return "'abcd'" */ public static String quoteString(String value) { // (char)39 -> ' diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/DatabaseController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/DatabaseController.java index 6348db871..6e4571151 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/DatabaseController.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/DatabaseController.java @@ -1,5 +1,6 @@ package ai.chat2db.server.web.api.controller.rdb; +import ai.chat2db.server.domain.api.enums.TaskStatusEnum; import ai.chat2db.server.domain.api.param.MetaDataQueryParam; import ai.chat2db.server.domain.api.param.datasource.DatabaseCreateParam; import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; @@ -14,8 +15,9 @@ import ai.chat2db.server.web.api.controller.data.source.vo.DatabaseVO; import ai.chat2db.server.web.api.controller.rdb.converter.DatabaseConverter; import ai.chat2db.server.web.api.controller.rdb.converter.RdbWebConverter; -import ai.chat2db.server.web.api.controller.rdb.data.export.strategy.ExportDBDataStrategy; -import ai.chat2db.server.web.api.controller.rdb.factory.ExportDBDataStrategyFactory; +import ai.chat2db.server.web.api.controller.rdb.data.service.DatabaseDataService; +import ai.chat2db.server.web.api.controller.rdb.data.task.TaskManager; +import ai.chat2db.server.web.api.controller.rdb.data.task.TaskState; import ai.chat2db.server.web.api.controller.rdb.request.DatabaseCreateRequest; import ai.chat2db.server.web.api.controller.rdb.request.DatabaseExportDataRequest; import ai.chat2db.server.web.api.controller.rdb.request.DatabaseExportRequest; @@ -26,12 +28,12 @@ import ai.chat2db.spi.model.Sql; import jakarta.servlet.http.HttpServletResponse; import jakarta.validation.Valid; +import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.web.bind.annotation.*; import java.io.PrintWriter; -import java.lang.reflect.Constructor; import java.util.Objects; /** @@ -40,6 +42,7 @@ @ConnectionInfoAspect @RequestMapping("/api/rdb/database") @RestController +@Slf4j public class DatabaseController { @Autowired private RdbWebConverter rdbWebConverter; @@ -49,6 +52,8 @@ public class DatabaseController { @Autowired public DatabaseConverter databaseConverter; + @Autowired + private DatabaseDataService databaseDataService; /** * Query the database_schema_list contained in the database @@ -59,8 +64,8 @@ public class DatabaseController { @GetMapping("/database_schema_list") public DataResult databaseSchemaList(@Valid DataSourceBaseRequest request) { MetaDataQueryParam queryParam = MetaDataQueryParam.builder().dataSourceId(request.getDataSourceId()) - .refresh( - request.isRefresh()).build(); + .refresh( + request.isRefresh()).build(); DataResult result = databaseService.queryDatabaseSchema(queryParam); MetaSchemaVO schemaDto2vo = rdbWebConverter.metaSchemaDto2vo(result.getData()); return DataResult.of(schemaDto2vo); @@ -69,8 +74,8 @@ public DataResult databaseSchemaList(@Valid DataSourceBaseRequest @GetMapping("list") public ListResult databaseList(@Valid DataSourceBaseRequest request) { DatabaseQueryAllParam queryParam = DatabaseQueryAllParam.builder().dataSourceId(request.getDataSourceId()) - .refresh( - request.isRefresh()).build(); + .refresh( + request.isRefresh()).build(); ListResult result = databaseService.queryAll(queryParam); return ListResult.of(rdbWebConverter.databaseDto2vo(result.getData())); } @@ -95,7 +100,7 @@ public ActionResult deleteDatabase(@Valid @RequestBody DataSourceBaseRequest req */ @PostMapping("/create_database_sql") public DataResult createDatabase(@Valid @RequestBody DatabaseCreateRequest request) { - if(StringUtils.isBlank(request.getName())){ + if (StringUtils.isBlank(request.getName())) { request.setName(request.getDatabaseName()); } Database database = databaseConverter.request2param(request); @@ -111,12 +116,13 @@ public DataResult createDatabase(@Valid @RequestBody DatabaseCreateRequest @PostMapping("/modify_database") public ActionResult modifyDatabase(@Valid @RequestBody UpdateDatabaseRequest request) { DatabaseCreateParam param = DatabaseCreateParam.builder().name(request.getDatabaseName()) - .name(request.getNewDatabaseName()).build(); + .name(request.getNewDatabaseName()).build(); return databaseService.modifyDatabase(param); } + @PostMapping("/export") - public void exportDatabase(@Valid @RequestBody DatabaseExportRequest request, HttpServletResponse response){ - String fileName = Objects.isNull(request.getSchemaName())?request.getDatabaseName() : request.getSchemaName(); + public void exportDatabase(@Valid @RequestBody DatabaseExportRequest request, HttpServletResponse response) { + String fileName = Objects.isNull(request.getSchemaName()) ? request.getDatabaseName() : request.getSchemaName(); response.setContentType("text/sql"); response.setHeader("Content-disposition", "attachment;filename*=utf-8''" + fileName + ".sql"); response.setCharacterEncoding("utf-8"); @@ -130,17 +136,18 @@ public void exportDatabase(@Valid @RequestBody DatabaseExportRequest request, Ht } @PostMapping("/export_data") - public void exportData(@Valid @RequestBody DatabaseExportDataRequest request, HttpServletResponse response) { - Class targetClass = ExportDBDataStrategyFactory.get(request.getExportType()); - response.setCharacterEncoding("utf-8"); - DatabaseExportDataParam param = databaseConverter.request2param(request); - try { - Constructor constructor = targetClass.getDeclaredConstructor(); - ExportDBDataStrategy service = (ExportDBDataStrategy) constructor.newInstance(); - service.doExport(param, response); - } catch (Exception e) { - throw new RuntimeException(e); - } + public DataResult exportData(@Valid @RequestBody DatabaseExportDataRequest request) { + DatabaseExportDataParam databaseExportDataParam = databaseConverter.request2param(request); + return databaseDataService.doExportAsync(databaseExportDataParam); + } + @GetMapping("/export_data_status/{taskId}") + public DataResult exportDataStatus(@PathVariable("taskId") Long taskId) { + TaskState task = TaskManager.getTask(taskId); + String state = task.getState(); + if (Objects.equals(state, TaskStatusEnum.FINISH.name()) || Objects.equals(state, TaskStatusEnum.ERROR.name())) { + TaskManager.removeTask(taskId); + } + return DataResult.of(task.getExportStatus()); } } diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataExporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataExporter.java new file mode 100644 index 000000000..dc5fa1a43 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataExporter.java @@ -0,0 +1,70 @@ +package ai.chat2db.server.web.api.controller.rdb.data; + +import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; +import ai.chat2db.server.web.api.util.StringUtils; +import ai.chat2db.spi.sql.Chat2DBContext; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.collections4.CollectionUtils; + +import java.io.*; +import java.sql.Connection; +import java.sql.SQLException; +import java.util.List; +import java.util.zip.ZipEntry; +import java.util.zip.ZipOutputStream; + +/** + * @author: zgq + * @date: 2024年06月04日 10:51 + */ +@Slf4j +public abstract class BaseDataExporter implements DataExportStrategy { + + protected String contentType; + protected String suffix; + + public static int BATCH_SIZE = 1000; + + @Override + public void doExport(DatabaseExportDataParam databaseExportDataParam, File file) throws IOException, SQLException { + List tableNames = databaseExportDataParam.getTableNames(); + if (CollectionUtils.isEmpty(tableNames)) { + throw new IllegalArgumentException("tableNames should not be null or empty"); + } + try (Connection connection = Chat2DBContext.getConnection()) { + if (tableNames.size() == 1) { + String tableName = tableNames.get(0); + if (StringUtils.isEmpty(tableName)) { + throw new IllegalArgumentException("tableName should not be null or empty"); + } + singleExport(connection, databaseExportDataParam,file); + } else { + multiExport(databaseExportDataParam, connection, file); + } + } + + } + + + private void multiExport(DatabaseExportDataParam databaseExportDataParam, + Connection connection, File file) throws IOException { + try (OutputStream outputStream = new FileOutputStream(file); + ZipOutputStream zipOutputStream = new ZipOutputStream(outputStream)) { + List tableNames = databaseExportDataParam.getTableNames(); + for (String tableName : tableNames) { + String fileName = tableName + suffix; + zipOutputStream.putNextEntry(new ZipEntry(fileName)); + try (ByteArrayOutputStream byteArrayOutputStream = multiExport(connection, databaseExportDataParam, tableName)) { + byteArrayOutputStream.writeTo(zipOutputStream); + zipOutputStream.closeEntry(); + } + } + } + } + + + protected abstract void singleExport(Connection connectionInfo, DatabaseExportDataParam databaseExportDataParam, File file) throws IOException, SQLException; + + + protected abstract ByteArrayOutputStream multiExport(Connection connection, DatabaseExportDataParam databaseExportDataParam, String tableName); +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataImporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataImporter.java new file mode 100644 index 000000000..99fac6299 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataImporter.java @@ -0,0 +1,8 @@ +package ai.chat2db.server.web.api.controller.rdb.data; + +/** + * @author: zgq + * @date: 2024年06月04日 10:52 + */ +public abstract class BaseDataImporter implements DataImportStrategy{ +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelExporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelExporter.java new file mode 100644 index 000000000..2adc541c3 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelExporter.java @@ -0,0 +1,105 @@ +package ai.chat2db.server.web.api.controller.rdb.data; + +import ai.chat2db.server.domain.api.enums.TaskStatusEnum; +import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; +import ai.chat2db.server.web.api.controller.rdb.data.task.TaskManager; +import ai.chat2db.spi.ValueProcessor; +import ai.chat2db.spi.model.JDBCDataValue; +import ai.chat2db.spi.sql.Chat2DBContext; +import ai.chat2db.spi.sql.SQLExecutor; +import ai.chat2db.spi.util.ResultSetUtils; +import com.alibaba.excel.EasyExcel; +import com.alibaba.excel.support.ExcelTypeEnum; +import com.alibaba.excel.write.builder.ExcelWriterSheetBuilder; +import lombok.extern.slf4j.Slf4j; + +import java.io.*; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +/** + * @author: zgq + * @date: 2024年06月04日 10:56 + */ +@Slf4j +public abstract class BaseExcelExporter extends BaseDataExporter { + @Override + protected void singleExport(Connection connection, DatabaseExportDataParam exportParam, File outputFile) { + ExcelTypeEnum excelType = getExcelType(); + try (OutputStream outputStream = new FileOutputStream(outputFile)) { + + + String tableName = exportParam.getTableNames().get(0); + String querySql = getQuerySql(exportParam, tableName); + + log.info("开始导出：{}表数据，导出类型：{}", tableName, excelType); + + SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> + writeExcelData(resultSet, excelType, outputStream, tableName, exportParam.getContainsHeader())); + + } catch (IOException e) { + TaskManager.updateStatus(TaskStatusEnum.ERROR); + throw new RuntimeException(e); + } + } + + @Override + protected ByteArrayOutputStream multiExport(Connection connection, DatabaseExportDataParam databaseExportDataParam, String tableName) { + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + ExcelTypeEnum excelType = getExcelType(); + + log.info("开始导出：{}表数据，导出类型：{}", tableName, excelType); + + String querySql = getQuerySql(databaseExportDataParam, tableName); + SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> { + writeExcelData(resultSet, excelType, byteArrayOutputStream, tableName, databaseExportDataParam.getContainsHeader()); + }); + return byteArrayOutputStream; + } + + private void writeExcelData(ResultSet resultSet, ExcelTypeEnum excelType, OutputStream outputStream, String sheetName, Boolean containsHeader) { + try { + ExcelWriterSheetBuilder excelWriterSheetBuilder = EasyExcel.write(outputStream).excelType(excelType).sheet(sheetName); + ResultSetMetaData metaData = resultSet.getMetaData(); + int columnCount = metaData.getColumnCount(); + ValueProcessor valueProcessor = Chat2DBContext.getMetaData().getValueProcessor(); + List> dataList = new ArrayList<>(); + + if (containsHeader) { + List header = ResultSetUtils.getRsHeader(resultSet); + excelWriterSheetBuilder.head(header.stream().map(Collections::singletonList).collect(Collectors.toList())); + } + + while (resultSet.next()) { + List rowDataList = new ArrayList<>(); + for (int i = 1; i <= columnCount; i++) { + JDBCDataValue jdbcDataValue = new JDBCDataValue(resultSet, metaData, i, false); + rowDataList.add(valueProcessor.getJdbcValue(jdbcDataValue)); + } + dataList.add(rowDataList); + } + + excelWriterSheetBuilder.doWrite(dataList); + TaskManager.increaseCurrent(); + } catch (SQLException e) { + TaskManager.updateStatus(TaskStatusEnum.ERROR); + log.error("Error writing Excel data", e); + throw new RuntimeException(e); + } + } + + private String getQuerySql(DatabaseExportDataParam databaseExportDataParam, String tableName) { + String databaseName = databaseExportDataParam.getDatabaseName(); + String schemaName = databaseExportDataParam.getSchemaName(); + return Chat2DBContext.getSqlBuilder().buildTableQuerySql(databaseName, schemaName, tableName); + } + + protected abstract ExcelTypeEnum getExcelType(); +} + diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelImporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelImporter.java new file mode 100644 index 000000000..65ab11e58 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelImporter.java @@ -0,0 +1,10 @@ +package ai.chat2db.server.web.api.controller.rdb.data; + +/** + * 功能描述 + * + * @author: zgq + * @date: 2024年06月04日 10:57 + */ +public abstract class BaseExcelImporter extends BaseDataImporter{ +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataExportStrategy.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataExportStrategy.java new file mode 100644 index 000000000..de651c6fb --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataExportStrategy.java @@ -0,0 +1,13 @@ +package ai.chat2db.server.web.api.controller.rdb.data; + +import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; + +import java.io.File; +import java.io.IOException; +import java.sql.SQLException; + +public interface DataExportStrategy { + + + void doExport(DatabaseExportDataParam databaseExportDataParam, File file) throws IOException, SQLException; +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataImportStrategy.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataImportStrategy.java new file mode 100644 index 000000000..96a453766 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataImportStrategy.java @@ -0,0 +1,4 @@ +package ai.chat2db.server.web.api.controller.rdb.data; + +public interface DataImportStrategy { +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataExporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataExporter.java new file mode 100644 index 000000000..470164757 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataExporter.java @@ -0,0 +1,26 @@ +package ai.chat2db.server.web.api.controller.rdb.data.csv; + +import ai.chat2db.server.domain.api.enums.ExportFileSuffix; +import ai.chat2db.server.web.api.controller.rdb.data.BaseExcelExporter; +import com.alibaba.excel.support.ExcelTypeEnum; +import org.springframework.stereotype.Component; + +/** + * @author: zgq + * @date: 2024年06月04日 10:05 + */ +@Component("csvExporter") +public class CsvDataExporter extends BaseExcelExporter { + + + public CsvDataExporter() { + this.contentType = "text/csv"; + this.suffix = ExportFileSuffix.CSV.getSuffix(); + } + + + @Override + protected ExcelTypeEnum getExcelType() { + return ExcelTypeEnum.CSV; + } +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataImporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataImporter.java new file mode 100644 index 000000000..28cb5f52a --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataImporter.java @@ -0,0 +1,12 @@ +package ai.chat2db.server.web.api.controller.rdb.data.csv; + +import ai.chat2db.server.web.api.controller.rdb.data.BaseExcelImporter; +import org.springframework.stereotype.Component; + +/** + * @author: zgq + * @date: 2024年06月04日 10:04 + */ +@Component("csvImporter") +public class CsvDataImporter extends BaseExcelImporter { +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2CsvStrategy.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2CsvStrategy.java deleted file mode 100644 index d091bf0ca..000000000 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2CsvStrategy.java +++ /dev/null @@ -1,54 +0,0 @@ -package ai.chat2db.server.web.api.controller.rdb.data.export.strategy; - -import ai.chat2db.server.domain.api.enums.ExportFileSuffix; -import ai.chat2db.spi.util.ResultSetUtils; -import com.alibaba.excel.EasyExcel; -import com.alibaba.excel.support.ExcelTypeEnum; - -import java.io.ByteArrayOutputStream; -import java.sql.Connection; -import java.sql.ResultSet; -import java.sql.ResultSetMetaData; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; -import java.util.Objects; -import java.util.stream.Collectors; - -public class ExportDBData2CsvStrategy extends ExportDBDataStrategy { - - public ExportDBData2CsvStrategy() { - suffix = ExportFileSuffix.CSV.getSuffix(); - contentType = "application/zip"; - } - - @Override - protected ByteArrayOutputStream exportData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException { - String sql; - if (Objects.isNull(schemaName)) { - sql = String.format("select * from %s", tableName); - } else { - sql = String.format("select * from %s.%s", schemaName, tableName); - } - ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); - try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { - ResultSetMetaData metaData = resultSet.getMetaData(); - List> headList = ResultSetUtils.getRsHeader(resultSet) - .stream() - .map(Collections::singletonList) - .collect(Collectors.toList()); - int columnCount = metaData.getColumnCount(); - List> dataList = new ArrayList<>(); - while (resultSet.next()) { - List row = new ArrayList<>(); - for (int i = 1; i <= columnCount; i++) { - row.add(resultSet.getString(i)); - } - dataList.add(row); - } - EasyExcel.write(byteOut).excelType(ExcelTypeEnum.CSV).sheet(tableName).head(headList).doWrite(dataList); - } - return byteOut; - } -} \ No newline at end of file diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2ExcelStrategy.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2ExcelStrategy.java deleted file mode 100644 index b8fc38ee8..000000000 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2ExcelStrategy.java +++ /dev/null @@ -1,53 +0,0 @@ -package ai.chat2db.server.web.api.controller.rdb.data.export.strategy; - -import ai.chat2db.server.domain.api.enums.ExportFileSuffix; -import ai.chat2db.spi.util.ResultSetUtils; -import com.alibaba.excel.EasyExcel; - -import java.io.ByteArrayOutputStream; -import java.sql.Connection; -import java.sql.ResultSet; -import java.sql.ResultSetMetaData; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; -import java.util.Objects; -import java.util.stream.Collectors; - -public class ExportDBData2ExcelStrategy extends ExportDBDataStrategy { - - public ExportDBData2ExcelStrategy() { - suffix = ExportFileSuffix.EXCEL.getSuffix(); - contentType = "application/zip"; - } - - @Override - protected ByteArrayOutputStream exportData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException { - String sql; - if (Objects.isNull(schemaName)) { - sql = String.format("select * from %s", tableName); - } else { - sql = String.format("select * from %s.%s", schemaName, tableName); - } - ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); - try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { - ResultSetMetaData metaData = resultSet.getMetaData(); - int columnCount = metaData.getColumnCount(); - List> headList = ResultSetUtils.getRsHeader(resultSet) - .stream() - .map(Collections::singletonList) - .collect(Collectors.toList()); - List> dataList = new ArrayList<>(); - while (resultSet.next()) { - List row = new ArrayList<>(); - for (int i = 1; i <= columnCount; i++) { - row.add(resultSet.getString(i)); - } - dataList.add(row); - } - EasyExcel.write(byteOut).sheet(tableName).head(headList).doWrite(dataList); - } - return byteOut; - } -} \ No newline at end of file diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2JsonStrategy.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2JsonStrategy.java deleted file mode 100644 index 47c8a59d8..000000000 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2JsonStrategy.java +++ /dev/null @@ -1,61 +0,0 @@ -package ai.chat2db.server.web.api.controller.rdb.data.export.strategy; - -import ai.chat2db.server.domain.api.enums.ExportFileSuffix; -import ai.chat2db.server.tools.base.excption.BusinessException; -import com.fasterxml.jackson.databind.ObjectMapper; - -import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.OutputStreamWriter; -import java.io.PrintWriter; -import java.nio.charset.StandardCharsets; -import java.sql.Connection; -import java.sql.ResultSet; -import java.sql.ResultSetMetaData; -import java.sql.SQLException; -import java.util.*; - -public class ExportDBData2JsonStrategy extends ExportDBDataStrategy { - - public ExportDBData2JsonStrategy() { - suffix = ExportFileSuffix.JSON.getSuffix(); - contentType = "application/zip"; - } - - @Override - protected ByteArrayOutputStream exportData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException { - ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); - try (PrintWriter writer = new PrintWriter(new OutputStreamWriter(byteOut, StandardCharsets.UTF_8))) { - String sql; - if (Objects.isNull(schemaName)) { - sql = String.format("SELECT * FROM %s", tableName); - } else { - sql = String.format("SELECT * FROM %s.%s", schemaName, tableName); - } - try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { - ResultSetMetaData metaData = resultSet.getMetaData(); - int columnCount = metaData.getColumnCount(); - List> data = new ArrayList<>(); - - while (resultSet.next()) { - Map row = new LinkedHashMap<>(); - for (int i = 1; i <= columnCount; i++) { - row.put(metaData.getColumnName(i), resultSet.getObject(i)); - } - data.add(row); - } - - ObjectMapper objectMapper = new ObjectMapper(); - try { - String jsonString = objectMapper.writeValueAsString(data); - writer.println(jsonString); - } catch (IOException e) { - throw new BusinessException("data.export2Json.error",data.toArray(),e); - } - } - } - return byteOut; - } - - -} \ No newline at end of file diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2SqlStrategy.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2SqlStrategy.java deleted file mode 100644 index 1e0577474..000000000 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2SqlStrategy.java +++ /dev/null @@ -1,33 +0,0 @@ -package ai.chat2db.server.web.api.controller.rdb.data.export.strategy; - -import ai.chat2db.server.domain.api.enums.ExportFileSuffix; -import ai.chat2db.spi.sql.Chat2DBContext; - -import java.io.ByteArrayOutputStream; -import java.io.OutputStreamWriter; -import java.io.PrintWriter; -import java.nio.charset.StandardCharsets; -import java.sql.Connection; -import java.sql.SQLException; - -/** - * @author: zgq - * @date: 2024年03月24日 12:50 - */ -public class ExportDBData2SqlStrategy extends ExportDBDataStrategy { - - public ExportDBData2SqlStrategy() { - suffix = ExportFileSuffix.SQL.getSuffix(); - contentType = "application/zip"; - } - - @Override - protected ByteArrayOutputStream exportData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException { - ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); - try (PrintWriter writer = new PrintWriter(new OutputStreamWriter(byteOut, StandardCharsets.UTF_8))) { - String sql = Chat2DBContext.getDBManage().exportDatabaseData(connection, databaseName, schemaName, tableName); - writer.println(sql); - } - return byteOut; - } -} \ No newline at end of file diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBDataStrategy.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBDataStrategy.java deleted file mode 100644 index cb4a18209..000000000 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBDataStrategy.java +++ /dev/null @@ -1,65 +0,0 @@ -package ai.chat2db.server.web.api.controller.rdb.data.export.strategy; - -import ai.chat2db.server.domain.api.enums.ExportFileSuffix; -import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; -import ai.chat2db.spi.sql.Chat2DBContext; -import jakarta.servlet.ServletOutputStream; -import jakarta.servlet.http.HttpServletResponse; -import lombok.AllArgsConstructor; -import lombok.Getter; -import lombok.NoArgsConstructor; - -import java.io.ByteArrayOutputStream; -import java.sql.Connection; -import java.sql.SQLException; -import java.util.List; -import java.util.Objects; -import java.util.zip.ZipEntry; -import java.util.zip.ZipOutputStream; - -/** - * @author: zgq - * @date: 2024年03月24日 12:46 - */ -@Getter -@AllArgsConstructor -@NoArgsConstructor -public abstract class ExportDBDataStrategy { - - public String suffix; - public String contentType; - - public void doExport(DatabaseExportDataParam param, HttpServletResponse response) { - String databaseName = param.getDatabaseName(); - String schemaName = param.getSchemaName(); - setResponseHeaders(param, response); - try (ServletOutputStream outputStream = response.getOutputStream(); - ZipOutputStream zipOut = new ZipOutputStream(outputStream); - Connection connection = Chat2DBContext.getConnection();) { - List tableNames = Chat2DBContext.getMetaData().tableNames(connection, databaseName, schemaName, null); - tableNames.addAll(Chat2DBContext.getMetaData().viewNames(connection, databaseName, schemaName)); - for (String tableName : tableNames) { - String fileName = tableName + getSuffix(); - zipOut.putNextEntry(new ZipEntry(fileName)); - ByteArrayOutputStream byteOut = exportData(connection, databaseName, schemaName, tableName); - byteOut.writeTo(zipOut); - zipOut.closeEntry(); - byteOut.close(); - } - } catch (Exception e) { - throw new RuntimeException(e); - } - } - - private void setResponseHeaders(DatabaseExportDataParam param, HttpServletResponse response) { - response.setContentType(contentType); - response.setHeader("Content-disposition", "attachment;filename*=utf-8''" + getFileName(param) + ExportFileSuffix.ZIP.getSuffix()); - } - - protected String getFileName(DatabaseExportDataParam param) { - return Objects.isNull(param.getSchemaName()) ? param.getDatabaseName() : param.getSchemaName(); - } - - protected abstract ByteArrayOutputStream exportData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException; - -} \ No newline at end of file diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataExportFactory.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataExportFactory.java new file mode 100644 index 000000000..e82412728 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataExportFactory.java @@ -0,0 +1,33 @@ +package ai.chat2db.server.web.api.controller.rdb.data.factory; + +import ai.chat2db.server.tools.common.exception.ParamBusinessException; +import ai.chat2db.server.web.api.controller.rdb.data.DataExportStrategy; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Component; + +import java.util.Map; +import java.util.Objects; + +/** + * @author: zgq + * @date: 2024年06月04日 10:26 + */ +@Component +public class DataExportFactory { + + public static final String BEAN_SUFFIX = "Exporter"; + private final Map exports; + + @Autowired + public DataExportFactory(Map exports) { + this.exports = exports; + } + + public DataExportStrategy getExporter(String type) { + DataExportStrategy dataExportStrategy = exports.get(type.toLowerCase() + BEAN_SUFFIX); + if (Objects.isNull(dataExportStrategy)) { + throw new ParamBusinessException(type); + } + return dataExportStrategy; + } +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataImportFactory.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataImportFactory.java new file mode 100644 index 000000000..8e8b4afa2 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataImportFactory.java @@ -0,0 +1,35 @@ +package ai.chat2db.server.web.api.controller.rdb.data.factory; + +import ai.chat2db.server.tools.common.exception.ParamBusinessException; +import ai.chat2db.server.web.api.controller.rdb.data.DataImportStrategy; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Component; + +import java.util.Map; +import java.util.Objects; + +/** + * @author: zgq + * @date: 2024年06月04日 10:07 + */ +@Component +public class DataImportFactory { + + + private static final String BEAN_SUFFIX = "Importer"; + private final Map imports; + + @Autowired + public DataImportFactory(Map imports) { + this.imports = imports; + } + + public DataImportStrategy getImporter(String type) { + DataImportStrategy dataImportStrategy = imports.get(type.toLowerCase() + BEAN_SUFFIX); + if (Objects.isNull(dataImportStrategy)) { + throw new ParamBusinessException(type); + } + return dataImportStrategy; + } + +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataExporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataExporter.java new file mode 100644 index 000000000..39112c059 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataExporter.java @@ -0,0 +1,112 @@ +package ai.chat2db.server.web.api.controller.rdb.data.json; + +import ai.chat2db.server.domain.api.enums.ExportFileSuffix; +import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; +import ai.chat2db.server.tools.base.excption.BusinessException; +import ai.chat2db.server.web.api.controller.rdb.data.BaseDataExporter; +import ai.chat2db.server.web.api.controller.rdb.data.task.TaskManager; +import ai.chat2db.spi.ValueProcessor; +import ai.chat2db.spi.model.JDBCDataValue; +import ai.chat2db.spi.sql.Chat2DBContext; +import ai.chat2db.spi.sql.SQLExecutor; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializationFeature; +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Component; + +import java.io.*; +import java.nio.charset.StandardCharsets; +import java.sql.Connection; +import java.sql.ResultSetMetaData; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * @author: zgq + * @date: 2024年06月04日 10:33 + */ +@Component("jsonExporter") +@Slf4j +public class JsonDataExporter extends BaseDataExporter { + + public JsonDataExporter() { + this.suffix = ExportFileSuffix.JSON.getSuffix(); + this.contentType = "application/json"; + } + + + @Override + protected void singleExport(Connection connection, DatabaseExportDataParam databaseExportDataParam, File file) { + String tableName = databaseExportDataParam.getTableNames().get(0); + String querySql = getQuerySql(databaseExportDataParam, tableName); + log.info("开始导出：{}表数据，导出类型：json", tableName); + try (PrintWriter writer = new PrintWriter(file, StandardCharsets.UTF_8);) { + writeJsonData(connection, querySql, writer); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Override + protected ByteArrayOutputStream multiExport(Connection connection, DatabaseExportDataParam databaseExportDataParam, String tableName) { + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + log.info("开始导出：{}表数据，导出类型：json", tableName); + try (PrintWriter writer = new PrintWriter(new OutputStreamWriter(byteArrayOutputStream, StandardCharsets.UTF_8))) { + String querySql = getQuerySql(databaseExportDataParam, tableName); + writeJsonData(connection, querySql, writer); + } + return byteArrayOutputStream; + } + + private void writeJsonData(Connection connection, String querySql, PrintWriter writer) { + SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> { + List> dataBatch = new ArrayList<>(); + ResultSetMetaData metaData = resultSet.getMetaData(); + ValueProcessor valueProcessor = Chat2DBContext.getMetaData().getValueProcessor(); + ObjectMapper objectMapper = new ObjectMapper(); + objectMapper.configure(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS, false); + + writer.println("["); + boolean firstBatch = true; + while (resultSet.next()) { + Map row = new LinkedHashMap<>(); + for (int i = 1; i <= metaData.getColumnCount(); i++) { + row.put(metaData.getColumnName(i), valueProcessor.getJdbcValue(new JDBCDataValue(resultSet, metaData, i, false))); + } + dataBatch.add(row); + + if (dataBatch.size() >= BATCH_SIZE || resultSet.isLast()) { + if (!firstBatch) { + writer.println(","); + } + writeBatch(writer, objectMapper, dataBatch); + firstBatch = false; + } + } + writer.println("]"); + }); + TaskManager.increaseCurrent(); + } + + private void writeBatch(PrintWriter writer, ObjectMapper objectMapper, List> dataBatch) { + try { + String jsonBatch = objectMapper.writeValueAsString(dataBatch); + writer.println(jsonBatch.substring(1, jsonBatch.length() - 1)); + writer.flush(); + dataBatch.clear(); + } catch (JsonProcessingException e) { + throw new BusinessException("data.export.json.error", null, e); + } + } + + private String getQuerySql(DatabaseExportDataParam databaseExportDataParam, String tableName) { + String databaseName = databaseExportDataParam.getDatabaseName(); + String schemaName = databaseExportDataParam.getSchemaName(); + return Chat2DBContext.getSqlBuilder().buildTableQuerySql(databaseName, schemaName, tableName); + } + + +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataImporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataImporter.java new file mode 100644 index 000000000..9dabcbf19 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataImporter.java @@ -0,0 +1,12 @@ +package ai.chat2db.server.web.api.controller.rdb.data.json; + +import ai.chat2db.server.web.api.controller.rdb.data.BaseDataImporter; +import org.springframework.stereotype.Component; + +/** + * @author: zgq + * @date: 2024年06月04日 10:33 + */ +@Component("jsonImporter") +public class JsonDataImporter extends BaseDataImporter { +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/DatabaseDataService.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/DatabaseDataService.java new file mode 100644 index 000000000..4a1df4974 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/DatabaseDataService.java @@ -0,0 +1,13 @@ +package ai.chat2db.server.web.api.controller.rdb.data.service; + +import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; +import ai.chat2db.server.tools.base.wrapper.result.DataResult; + +/** + * @author: zgq + * @date: 2024年06月08日 10:32 + */ +public interface DatabaseDataService { + + DataResult doExportAsync(DatabaseExportDataParam databaseExportDataParam); +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/impl/DatabaseDataImpl.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/impl/DatabaseDataImpl.java new file mode 100644 index 000000000..7e37f84c7 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/impl/DatabaseDataImpl.java @@ -0,0 +1,150 @@ +package ai.chat2db.server.web.api.controller.rdb.data.service.impl; + +import ai.chat2db.server.domain.api.enums.TaskStatusEnum; +import ai.chat2db.server.domain.api.enums.TaskTypeEnum; +import ai.chat2db.server.domain.api.param.TaskCreateParam; +import ai.chat2db.server.domain.api.param.TaskUpdateParam; +import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; +import ai.chat2db.server.domain.api.service.TaskService; +import ai.chat2db.server.domain.repository.Dbutils; +import ai.chat2db.server.tools.base.wrapper.result.DataResult; +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.web.api.controller.rdb.data.factory.DataExportFactory; +import ai.chat2db.server.web.api.controller.rdb.data.factory.DataImportFactory; +import ai.chat2db.server.web.api.controller.rdb.data.service.DatabaseDataService; +import ai.chat2db.server.web.api.controller.rdb.data.task.TaskManager; +import ai.chat2db.server.web.api.controller.rdb.data.task.TaskState; +import ai.chat2db.spi.sql.Chat2DBContext; +import ai.chat2db.spi.sql.ConnectInfo; +import cn.hutool.core.date.DatePattern; +import cn.hutool.core.io.FileUtil; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Service; + +import java.io.File; +import java.io.IOException; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.sql.SQLException; +import java.time.LocalDateTime; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +/** + * @author: zgq + * @date: 2024年06月08日 10:32 + */ +@Service +@Slf4j +public class DatabaseDataImpl implements DatabaseDataService { + + public static final String EXPORT_DATA_TASK_TEMPLATE = "export_%s_data"; + public static final String IMPORT_DATA_TASK_TEMPLATE = "import_%s_data"; + @Autowired + private DataExportFactory dataExportFactory; + @Autowired + private DataImportFactory dataImportFactory; + @Autowired + private TaskService taskService; + + @Override + public DataResult doExportAsync(DatabaseExportDataParam databaseExportDataParam) { + List tableNames = databaseExportDataParam.getTableNames(); + String databaseName = databaseExportDataParam.getDatabaseName(); + String schemaName = databaseExportDataParam.getSchemaName(); + Long dataSourceId = databaseExportDataParam.getDataSourceId(); + String taskName = buildTaskName(tableNames, databaseName, schemaName); + String fileName = URLEncoder.encode( + taskName + "_" + LocalDateTime.now().format(DatePattern.PURE_DATETIME_FORMATTER), + StandardCharsets.UTF_8); + String suffix = "."; + int size = tableNames.size(); + if (size > 1) { + suffix += "zip"; + } else { + suffix += databaseExportDataParam.getExportType().toLowerCase(); + } + File file = FileUtil.createTempFile(fileName, suffix, true); + file.deleteOnExit(); + LoginUser loginUser = ContextUtils.getLoginUser(); + ConnectInfo connectInfo = Chat2DBContext.getConnectInfo().copy(); + DataResult dataResult = createTask(tableNames.get(0), databaseName, schemaName, dataSourceId, taskName); + Long taskId = dataResult.getData(); + CompletableFuture.runAsync(() -> { + buildContext(loginUser, connectInfo); + TaskManager.addTask(taskId, TaskState.builder().state(TaskStatusEnum.PROCESSING.name()).total(size) + .current(0).build()); + try { + dataExportFactory.getExporter(databaseExportDataParam.getExportType()).doExport(databaseExportDataParam, file); + } catch (IOException | SQLException e) { + throw new RuntimeException(e); + } + }).whenComplete((v, ex) -> { + updateStatus(taskId, file, ex); + removeContext(); + TaskManager.removeTaskId(); + }); + return dataResult; + + } + + 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 removeContext() { + Dbutils.removeSession(); + ContextUtils.removeContext(); + Chat2DBContext.removeContext(); + } + + private DataResult createTask(String tableName, String databaseName, String schemaName, Long datasourceId, String taskName) { + TaskCreateParam param = new TaskCreateParam(); + param.setTaskName(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 buildContext(LoginUser loginUser, ConnectInfo connectInfo) { + ContextUtils.setContext(Context.builder() + .loginUser(loginUser) + .build()); + Dbutils.setSession(); + Chat2DBContext.putContext(connectInfo); + } + + private String buildTaskName(List tableNames, String databaseName, String schemaName) { + StringBuilder taskNameBuilder = new StringBuilder(); + if (StringUtils.isNotBlank(databaseName)) { + taskNameBuilder.append(databaseName).append("_"); + } + if (StringUtils.isNotBlank(schemaName)) { + taskNameBuilder.append(schemaName).append("_"); + } + if (tableNames.size() == 1) { + taskNameBuilder.append(StringUtils.join(tableNames, "_")); + } + return String.format(EXPORT_DATA_TASK_TEMPLATE, taskNameBuilder); + } + +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/sql/SqlDataExporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/sql/SqlDataExporter.java new file mode 100644 index 000000000..bc2a330ad --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/sql/SqlDataExporter.java @@ -0,0 +1,184 @@ +package ai.chat2db.server.web.api.controller.rdb.data.sql; + +import ai.chat2db.server.domain.api.enums.ExportFileSuffix; +import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam; +import ai.chat2db.server.web.api.controller.rdb.data.BaseDataExporter; +import ai.chat2db.server.web.api.controller.rdb.data.task.TaskManager; +import ai.chat2db.spi.MetaData; +import ai.chat2db.spi.SqlBuilder; +import ai.chat2db.spi.ValueProcessor; +import ai.chat2db.spi.model.JDBCDataValue; +import ai.chat2db.spi.sql.Chat2DBContext; +import ai.chat2db.spi.sql.SQLExecutor; +import ai.chat2db.spi.util.ResultSetUtils; +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Component; + +import java.io.*; +import java.nio.charset.StandardCharsets; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author: zgq + * @date: 2024年06月04日 10:33 + */ +@Component("sqlExporter") +@Slf4j +public class SqlDataExporter extends BaseDataExporter { + + public SqlDataExporter() { + this.suffix = ExportFileSuffix.SQL.getSuffix(); + this.contentType = "text/sql"; + } + + /** + * @param connection + * @param databaseExportDataParam + * @param file + */ + @Override + protected void singleExport(Connection connection, DatabaseExportDataParam databaseExportDataParam, File file) { + String tableName = databaseExportDataParam.getTableNames().get(0); + log.info("开始导出：{}表数据，导出类型：sql", tableName); + try (PrintWriter writer = new PrintWriter(file);) { + exportSql(connection, databaseExportDataParam, tableName, writer); + } catch (FileNotFoundException e) { + throw new RuntimeException(e); + } + } + + @Override + protected ByteArrayOutputStream multiExport(Connection connection, DatabaseExportDataParam databaseExportDataParam, String tableName) { + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + log.info("开始导出：{}表数据，导出类型：sql", tableName); + try (PrintWriter writer = new PrintWriter(new OutputStreamWriter(byteArrayOutputStream, StandardCharsets.UTF_8))) { + exportSql(connection, databaseExportDataParam, tableName, writer); + } + return byteArrayOutputStream; + } + + private void exportSql(Connection connection, DatabaseExportDataParam databaseExportDataParam, String tableName, PrintWriter writer) { + String databaseName = databaseExportDataParam.getDatabaseName(); + String schemaName = databaseExportDataParam.getSchemaName(); + Boolean containsHeader = databaseExportDataParam.getContainsHeader(); + MetaData metaData = Chat2DBContext.getMetaData(); + String querySql = metaData.getSqlBuilder().buildTableQuerySql(databaseName, schemaName, tableName); + SqlBuilder sqlBuilder = metaData.getSqlBuilder(); + ValueProcessor valueProcessor = metaData.getValueProcessor(); + String sqyType = databaseExportDataParam.getSqyType(); + + switch (sqyType) { + case "single" -> exportSingleInsert(connection, querySql, containsHeader, sqlBuilder, + valueProcessor, databaseName, schemaName, tableName, writer); + case "multi" -> exportMultiInsert(connection, querySql, containsHeader, sqlBuilder, + valueProcessor, databaseName, schemaName, tableName, writer); + case "update" -> exportUpdate(connection, querySql, sqlBuilder, valueProcessor, + databaseName, schemaName, tableName, writer); + default -> throw new IllegalArgumentException("Unsupported sqyType: " + sqyType); + } + } + + private void exportSingleInsert(Connection connection, String querySql, Boolean containsHeader, + SqlBuilder sqlBuilder, ValueProcessor valueProcessor, + String databaseName, String schemaName, String tableName, PrintWriter writer) { + List sqlList = new ArrayList<>(BATCH_SIZE); + SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> { + List header = containsHeader ? ResultSetUtils.getRsHeader(resultSet) : null; + while (resultSet.next()) { + List rowData = extractRowData(resultSet, valueProcessor); + String sql = sqlBuilder.buildSingleInsertSql(databaseName, schemaName, tableName, header, rowData); + sqlList.add(sql); + if (sqlList.size() >= BATCH_SIZE || resultSet.isLast()) { + writeSqlList(writer, sqlList); + } + } + }); + TaskManager.increaseCurrent(); + } + + private void exportMultiInsert(Connection connection, String querySql, Boolean containsHeader, + SqlBuilder sqlBuilder, ValueProcessor valueProcessor, + String databaseName, String schemaName, String tableName, PrintWriter writer) { + SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> { + List> dataList = new ArrayList<>(BATCH_SIZE); + List header = containsHeader ? ResultSetUtils.getRsHeader(resultSet) : null; + while (resultSet.next()) { + dataList.add(extractRowData(resultSet, valueProcessor)); + } + String sql = sqlBuilder.buildMultiInsertSql(databaseName, schemaName, tableName, header, dataList); + writer.println(sql); + writer.flush(); + }); + TaskManager.increaseCurrent(); + } + + private void exportUpdate(Connection connection, String querySql, SqlBuilder sqlBuilder, + ValueProcessor valueProcessor, + String databaseName, String schemaName, String tableName, PrintWriter writer) { + List sqlList = new ArrayList<>(BATCH_SIZE); + SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> { + Map primaryKeyMap = getPrimaryKeyMap(connection, databaseName, schemaName, tableName); + while (resultSet.next()) { + Map row = extractRowDataAsMap(resultSet, valueProcessor, primaryKeyMap); + String sql = sqlBuilder.buildUpdateSql(databaseName, schemaName, tableName, row, primaryKeyMap); + sqlList.add(sql); + if (sqlList.size() >= BATCH_SIZE || resultSet.isLast()) { + writeSqlList(writer, sqlList); + } + } + }); + TaskManager.increaseCurrent(); + } + + private List extractRowData(ResultSet resultSet, ValueProcessor valueProcessor) throws SQLException { + ResultSetMetaData metaData = resultSet.getMetaData(); + List rowData = new ArrayList<>(metaData.getColumnCount()); + for (int i = 1; i <= metaData.getColumnCount(); i++) { + JDBCDataValue jdbcDataValue = new JDBCDataValue(resultSet, metaData, i, false); + rowData.add(valueProcessor.getJdbcValueString(jdbcDataValue)); + } + return rowData; + } + + private Map extractRowDataAsMap(ResultSet resultSet, ValueProcessor valueProcessor, + Map primaryKeyMap) throws SQLException { + ResultSetMetaData metaData = resultSet.getMetaData(); + Map row = new HashMap<>(metaData.getColumnCount()); + for (int i = 1; i <= metaData.getColumnCount(); i++) { + JDBCDataValue jdbcDataValue = new JDBCDataValue(resultSet, metaData, i, false); + String columnName = metaData.getColumnName(i); + String jdbcValueString = valueProcessor.getJdbcValueString(jdbcDataValue); + if (primaryKeyMap.containsKey(columnName)) { + primaryKeyMap.put(columnName, jdbcValueString); + } else { + row.put(columnName, jdbcValueString); + } + } + return row; + } + + private Map getPrimaryKeyMap(Connection connection, String databaseName, + String schemaName, String tableName) throws SQLException { + Map primaryKeyMap = new HashMap<>(); + try (ResultSet primaryKeys = connection.getMetaData().getPrimaryKeys(databaseName, schemaName, tableName)) { + while (primaryKeys.next()) { + primaryKeyMap.put(primaryKeys.getString("COLUMN_NAME"), ""); + } + } + return primaryKeyMap; + } + + private void writeSqlList(PrintWriter writer, List sqlList) { + sqlList.forEach(writer::println); + sqlList.clear(); + } + + +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskManager.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskManager.java new file mode 100644 index 000000000..a59e4c434 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskManager.java @@ -0,0 +1,69 @@ +package ai.chat2db.server.web.api.controller.rdb.data.task; + +import ai.chat2db.server.domain.api.enums.TaskStatusEnum; + +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; + + +public class TaskManager { + public static final ThreadLocal TASK_ID = new ThreadLocal<>(); + public static final Map taskMap = new ConcurrentHashMap<>(); + + + public static void increaseCurrent(int current) { + TaskState task = getTask(); + task.setCurrent(task.getCurrent() + current); + if (task.getCurrent() >= task.getTotal()) { + task.setState(TaskStatusEnum.FINISH.name()); + } + } + + public static void increaseCurrent() { + TaskState task = getTask(); + task.setCurrent(task.getCurrent() +1); + if (task.getCurrent() >= task.getTotal()) { + task.setState(TaskStatusEnum.FINISH.name()); + } + } + + public static void updateStatus(TaskStatusEnum status) { + TaskState task = getTask(); + task.setState(status.name()); + } + + + public static void addTask(Long taskId, TaskState taskState) { + setTaskId(taskId); + taskMap.put(taskId, taskState); + } + + public static TaskState getTask(Long taskId) { + TaskState taskState = taskMap.get(taskId); + if (Objects.isNull(taskState)) { + throw new IllegalArgumentException("taskId is not valid"); + } + return taskState; + } + + public static TaskState getTask() { + return getTask(getTaskId()); + } + + public static void removeTask(Long taskId) { + taskMap.remove(taskId); + } + + public static void setTaskId(Long taskId) { + TASK_ID.set(taskId); + } + + public static Long getTaskId() { + return TASK_ID.get(); + } + + public static void removeTaskId() { + TASK_ID.remove(); + } +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskState.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskState.java new file mode 100644 index 000000000..c93bdf81c --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskState.java @@ -0,0 +1,28 @@ +package ai.chat2db.server.web.api.controller.rdb.data.task; + +import lombok.Builder; +import lombok.Data; + +/** + * @author: zgq + * @date: 2024年06月10日 15:51 + */ +@Data +@Builder +public class TaskState { + private String taskId; + private String state; + private int total; + private int current; + + + public String getExportStatus() { + StringBuilder statusBuilder = new StringBuilder(); + statusBuilder.append("导出状态: ").append(state) + .append(" 导出进度: ") + .append(current).append("/") + .append(total); + return statusBuilder.toString(); + } + +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataExporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataExporter.java new file mode 100644 index 000000000..e643aef86 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataExporter.java @@ -0,0 +1,24 @@ +package ai.chat2db.server.web.api.controller.rdb.data.xls; + +import ai.chat2db.server.domain.api.enums.ExportFileSuffix; +import ai.chat2db.server.web.api.controller.rdb.data.BaseExcelExporter; +import com.alibaba.excel.support.ExcelTypeEnum; +import org.springframework.stereotype.Component; + +/** + * @author: zgq + * @date: 2024年06月04日 10:34 + */ +@Component("xlsExporter") +public class XlsDataExporter extends BaseExcelExporter { + + public XlsDataExporter() { + this.suffix = ExportFileSuffix.XLS.getSuffix(); + this.contentType="application/vnd.ms-excel"; + } + + @Override + protected ExcelTypeEnum getExcelType() { + return ExcelTypeEnum.XLS; + } +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataImporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataImporter.java new file mode 100644 index 000000000..bd232b6c7 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataImporter.java @@ -0,0 +1,13 @@ +package ai.chat2db.server.web.api.controller.rdb.data.xls; + +import ai.chat2db.server.web.api.controller.rdb.data.BaseExcelImporter; +import org.springframework.stereotype.Component; + +/** + * @author: zgq + * @date: 2024年06月04日 10:34 + */ +@Component("xlsImporter") +public class XlsDataImporter extends BaseExcelImporter { + +} \ No newline at end of file diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataExporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataExporter.java new file mode 100644 index 000000000..ce4ac61b1 --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataExporter.java @@ -0,0 +1,25 @@ +package ai.chat2db.server.web.api.controller.rdb.data.xlsx; + +import ai.chat2db.server.domain.api.enums.ExportFileSuffix; +import ai.chat2db.server.web.api.controller.rdb.data.BaseExcelExporter; +import com.alibaba.excel.support.ExcelTypeEnum; +import org.springframework.stereotype.Component; + +/** + * @author: zgq + * @date: 2024年06月04日 10:34 + */ +@Component("xlsxExporter") +public class XlsxDataExporter extends BaseExcelExporter { + + public XlsxDataExporter() { + this.suffix = ExportFileSuffix.EXCEL.getSuffix(); + this.contentType="application/vnd.ms-excel"; + } + + + @Override + protected ExcelTypeEnum getExcelType() { + return ExcelTypeEnum.XLSX; + } +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataImporter.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataImporter.java new file mode 100644 index 000000000..43b6dc4cc --- /dev/null +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataImporter.java @@ -0,0 +1,12 @@ +package ai.chat2db.server.web.api.controller.rdb.data.xlsx; + +import ai.chat2db.server.web.api.controller.rdb.data.BaseExcelImporter; +import org.springframework.stereotype.Component; + +/** + * @author: zgq + * @date: 2024年06月04日 10:34 + */ +@Component("xlsxImporter") +public class XlsxDataImporter extends BaseExcelImporter { +} diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/factory/ExportDBDataStrategyFactory.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/factory/ExportDBDataStrategyFactory.java deleted file mode 100644 index 638557b17..000000000 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/factory/ExportDBDataStrategyFactory.java +++ /dev/null @@ -1,31 +0,0 @@ -package ai.chat2db.server.web.api.controller.rdb.factory; - -import ai.chat2db.server.domain.api.enums.ExportTypeEnum; -import ai.chat2db.server.web.api.controller.rdb.data.export.strategy.*; -import lombok.SneakyThrows; - -import java.util.Map; - -/** - * @author: zgq - * @date: 2024年03月24日 12:53 - */ -public class ExportDBDataStrategyFactory { - - public static final Map> SERVICE_MAP = Map.of( - ExportTypeEnum.SQL.getCode(), ExportDBData2SqlStrategy.class, - ExportTypeEnum.CSV.getCode(), ExportDBData2CsvStrategy.class, - ExportTypeEnum.EXCEL.getCode(), ExportDBData2ExcelStrategy.class, - ExportTypeEnum.JSON.getCode(), ExportDBData2JsonStrategy.class - ); - - @SneakyThrows - public static Class get(String type) { - Class dataResult = SERVICE_MAP.get(type); - if (dataResult == null) { - throw new ClassNotFoundException("no ExportUI was found"); - } else { - return dataResult; - } - } -} \ No newline at end of file diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/DatabaseExportDataRequest.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/DatabaseExportDataRequest.java index 41686b503..0f8a03fc1 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/DatabaseExportDataRequest.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/DatabaseExportDataRequest.java @@ -1,11 +1,14 @@ package ai.chat2db.server.web.api.controller.rdb.request; import ai.chat2db.server.web.api.controller.data.source.request.DataSourceBaseRequest; +import jakarta.validation.constraints.NotEmpty; import jakarta.validation.constraints.NotNull; import lombok.AllArgsConstructor; import lombok.Data; import lombok.NoArgsConstructor; +import java.util.List; + /** * @author: zgq * @date: 2024年03月24日 12:36 @@ -16,4 +19,11 @@ public class DatabaseExportDataRequest extends DataSourceBaseRequest { @NotNull private String exportType; + @NotEmpty + private List tableNames; + /** + * single：单行插入，multi：多行插入，update：更新语句 + */ + private String sqyType; + private Boolean containsHeader; } \ No newline at end of file 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 index b6c784117..54d3aa91b 100644 --- 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 @@ -8,19 +8,17 @@ 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 jakarta.servlet.http.HttpServletResponse; 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; +import java.io.*; @ConnectionInfoAspect @RequestMapping("/api/task") @@ -43,35 +41,42 @@ public WebPageResult list() { } @GetMapping("/download/{id}") - public ResponseEntity download(@PathVariable Long id) { + public void download(@PathVariable Long id, HttpServletResponse response) { DataResult task = taskService.get(id); - if(task.getData() == null){ + Task data = task.getData(); + if (data == null) { log.error("task is null"); throw new RuntimeException("task is null"); } - if(ContextUtils.getUserId() != task.getData().getUserId()){ + if (!ContextUtils.getUserId().equals(data.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); - } + File file = new File(data.getDownloadUrl()); - 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!"); + if (!file.exists() || !file.isFile()) { + log.error("File not found or is not a file: {}", file.getAbsolutePath()); + throw new RuntimeException("File not found or accessible"); } - } + response.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + file.getName() + "\""); + response.setContentType(MediaType.APPLICATION_OCTET_STREAM_VALUE); + + try (InputStream inputStream = new FileInputStream(file)) { + OutputStream outputStream = response.getOutputStream(); + byte[] buffer = new byte[8192]; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + outputStream.write(buffer, 0, bytesRead); + outputStream.flush(); + } + } catch (IOException e) { + log.error("Error occurred while processing file download", e); + throw new RuntimeException("Error in file download", e); + } + } } diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/SqlBuilder.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/SqlBuilder.java index 9089634cb..4e59080cf 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/SqlBuilder.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/SqlBuilder.java @@ -1,8 +1,12 @@ package ai.chat2db.spi; -import ai.chat2db.spi.model.*; +import ai.chat2db.spi.model.Database; +import ai.chat2db.spi.model.OrderBy; +import ai.chat2db.spi.model.QueryResult; +import ai.chat2db.spi.model.Schema; import java.util.List; +import java.util.Map; public interface SqlBuilder { @@ -12,7 +16,7 @@ public interface SqlBuilder { * @param table * @return */ - String buildCreateTableSql(T table); + String buildCreateTableSql(T table); /** @@ -83,9 +87,18 @@ public interface SqlBuilder { /** * DML SQL + * * @param table * @param type * @return */ - String getTableDmlSql(T table,String type); + String getTableDmlSql(T table, String type); + + String buildTableQuerySql(String databaseName, String schemaName, String tableName); + + String buildSingleInsertSql(String databaseName, String schemaName, String tableName, List columnList, List valueList); + + String buildMultiInsertSql(String databaseName, String schemaName, String tableName, List columnList, List> valueLists); + + String buildUpdateSql(String databaseName, String schemaName, String tableName, Map row,Map primaryKeyMap); } diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultSqlBuilder.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultSqlBuilder.java index 8f3964c86..33df0264f 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultSqlBuilder.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultSqlBuilder.java @@ -13,16 +13,26 @@ import net.sf.jsqlparser.statement.select.PlainSelect; import net.sf.jsqlparser.statement.select.Select; import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.collections4.MapUtils; import org.apache.commons.lang3.ObjectUtils; import org.apache.commons.lang3.StringUtils; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.stream.Collectors; public class DefaultSqlBuilder implements SqlBuilder { + @Override + public String buildTableQuerySql(String databaseName, String schemaName, String tableName) { + StringBuilder sqlBuilder = new StringBuilder(); + sqlBuilder.append("SELECT * FROM "); + buildTableName(databaseName, schemaName, tableName, sqlBuilder); + return sqlBuilder.toString(); + } + @Override public String buildCreateTableSql(Table table) { return null; @@ -198,19 +208,17 @@ private String getInsertSql(String name, List columnList) { * Generates the base part of the INSERT SQL statement. * Optionally includes column names if provided. * - * @param schemaName Name of the database schema. - * @param tableName Name of the table to insert into. - * @param columnList Optional list of column names. + * @param databaseName + * @param schemaName Name of the database schema. + * @param tableName Name of the table to insert into. + * @param columnList Optional list of column names. * @return The base part of the INSERT SQL statement. */ - private String generateBaseInsertSql(String schemaName, String tableName, List columnList) { + protected String buildBaseInsertSql(String databaseName, String schemaName, String tableName, List columnList) { StringBuilder script = new StringBuilder(); - if (StringUtils.isNotBlank(schemaName)) { - script.append(schemaName).append('.'); - } - - script.append(tableName); + script.append("INSERT INTO "); + buildTableName(databaseName, schemaName, tableName, script); if (CollectionUtils.isNotEmpty(columnList)) { script.append(" (") @@ -218,10 +226,21 @@ private String generateBaseInsertSql(String schemaName, String tableName, List columnList, List valueList) { - String baseSql = generateBaseInsertSql(schemaName, tableName, columnList); + public String buildSingleInsertSql(String databaseName, String schemaName, String tableName, List columnList, List valueList) { + String baseSql = buildBaseInsertSql(databaseName, schemaName, tableName, columnList); return baseSql + "(" + String.join(",", valueList) + ");"; } @@ -245,14 +264,37 @@ public String generateSingleInsertSql(String schemaName, String tableName, List< * @param valueLists List of lists, each inner list represents values for a row. * @return The complete multi-row INSERT SQL statement. */ - public String generateMultiInsertSql(String schemaName, String tableName, List columnList, List> valueLists) { - String baseSql = generateBaseInsertSql(schemaName, tableName, columnList); + public String buildMultiInsertSql(String databaseName, String schemaName, String tableName, List columnList, List> valueLists) { + String baseSql = buildBaseInsertSql(databaseName, schemaName, tableName, columnList); String valuesPart = valueLists.stream() .map(values -> "(" + String.join(",", values) + ")") .collect(Collectors.joining(",\n")); return baseSql + valuesPart + ";"; } + + @Override + public String buildUpdateSql(String databaseName, String schemaName, String tableName, Map row, Map primaryKeyMap) { + StringBuilder script = new StringBuilder(); + script.append("UPDATE "); + buildTableName(databaseName, schemaName, tableName, script); + + script.append(" SET "); + List setClauses = row.entrySet().stream() + .map(entry -> entry.getKey() + " = " + entry.getValue()) + .collect(Collectors.toList()); + script.append(String.join(",", setClauses)); + + if (MapUtils.isNotEmpty(primaryKeyMap)) { + script.append(" WHERE "); + List whereClauses = primaryKeyMap.entrySet().stream() + .map(entry -> entry.getKey() + " = " + entry.getValue()) + .collect(Collectors.toList()); + script.append(String.join(" AND ", whereClauses)); + } + return script + ";"; + } + private List getPrimaryColumns(List
headerList) { if (CollectionUtils.isEmpty(headerList)) { return Lists.newArrayList(); diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java index 8db1999ad..235f1328a 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java @@ -715,4 +715,20 @@ private ExecuteResult execute(String sql, Integer offset, Integer count) { } return executeResult; } + + public void execute(Connection connection, String sql, int batchSize, ResultSetConsumer consumer) { + log.info("execute:{}", sql); + try (Statement stmt = connection.createStatement()) { + stmt.setFetchSize(batchSize); + boolean query = stmt.execute(sql); + // Represents the query + if (query) { + try (ResultSet rs = stmt.getResultSet()) { + consumer.accept(rs); + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } + } }