data export

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("`");
+    }
 }

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

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<String> tableNames;
+    private String sqyType;
+    private Boolean containsHeader;
+
 }

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();
     }

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 -> '

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<MetaSchemaVO> databaseSchemaList(@Valid DataSourceBaseRequest request) {
         MetaDataQueryParam queryParam = MetaDataQueryParam.builder().dataSourceId(request.getDataSourceId())
-            .refresh(
-            request.isRefresh()).build();
+                .refresh(
+                        request.isRefresh()).build();
         DataResult<MetaSchema> result = databaseService.queryDatabaseSchema(queryParam);
         MetaSchemaVO schemaDto2vo = rdbWebConverter.metaSchemaDto2vo(result.getData());
         return DataResult.of(schemaDto2vo);
@@ -69,8 +74,8 @@ public DataResult<MetaSchemaVO> databaseSchemaList(@Valid DataSourceBaseRequest
     @GetMapping("list")
     public ListResult<DatabaseVO> databaseList(@Valid DataSourceBaseRequest request) {
         DatabaseQueryAllParam queryParam = DatabaseQueryAllParam.builder().dataSourceId(request.getDataSourceId())
-            .refresh(
-                request.isRefresh()).build();
+                .refresh(
+                        request.isRefresh()).build();
         ListResult<Database> 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<Sql> 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<Sql> 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<Long> exportData(@Valid @RequestBody DatabaseExportDataRequest request) {
+        DatabaseExportDataParam databaseExportDataParam = databaseConverter.request2param(request);
+        return databaseDataService.doExportAsync(databaseExportDataParam);
+    }
 
+    @GetMapping("/export_data_status/{taskId}")
+    public DataResult<String> 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());
     }
 }

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<String> 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<String> 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);
+}

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{
+}

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<List<Object>> dataList = new ArrayList<>();
+
+            if (containsHeader) {
+                List<String> header = ResultSetUtils.getRsHeader(resultSet);
+                excelWriterSheetBuilder.head(header.stream().map(Collections::singletonList).collect(Collectors.toList()));
+            }
+
+            while (resultSet.next()) {
+                List<Object> 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();
+}
+