opensource/Chat2DB
1
0
Fork 0
mirror of https://github.com/CodePhiliaX/Chat2DB.git synced 2025-07-29 10:43:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

data export

Browse Source
This commit is contained in:
zgq
2024-06-18 13:56:25 +08:00
parent 5ac52afc80
commit cd16555c45
38 changed files with 1133 additions and 370 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/DatabaseController.java
View File

@ -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.request.DataSourceBaseRe
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.MetaSchema;
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 @@ import java.util.Objects;
@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 class DatabaseController {
@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 class DatabaseController {
*/
@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 class DatabaseController {
@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 class DatabaseController {
}
@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());
}
}

70
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataExporter.java Normal file
View File

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

8
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataImporter.java Normal file
View File

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

105
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelExporter.java Normal file
View File

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

10
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelImporter.java Normal file
View File

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

13
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataExportStrategy.java Normal file
View File

@ -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;
}

4
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/DataImportStrategy.java Normal file
View File

@ -0,0 +1,4 @@
package ai.chat2db.server.web.api.controller.rdb.data;
public interface DataImportStrategy {
}

26
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataExporter.java Normal file
View File

@ -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;
}
}

12
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/csv/CsvDataImporter.java Normal file
View File

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

54
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2CsvStrategy.java
View File

@ -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<List<String>> headList = ResultSetUtils.getRsHeader(resultSet)
.stream()
.map(Collections::singletonList)
.collect(Collectors.toList());
int columnCount = metaData.getColumnCount();
List<List<Object>> dataList = new ArrayList<>();
while (resultSet.next()) {
List<Object> 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;
}
}

53
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2ExcelStrategy.java
View File

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

61
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2JsonStrategy.java
View File

@ -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<Map<String, Object>> data = new ArrayList<>();
while (resultSet.next()) {
Map<String, Object> 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;
}
}

33
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBData2SqlStrategy.java
View File

@ -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;
}
}

65
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/export/strategy/ExportDBDataStrategy.java
View File

@ -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<String> 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;
}

33
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataExportFactory.java Normal file
View File

@ -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<String, DataExportStrategy> exports;
@Autowired
public DataExportFactory(Map<String, DataExportStrategy> 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;
}
}

35
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/factory/DataImportFactory.java Normal file
View File

@ -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<String, DataImportStrategy> imports;
@Autowired
public DataImportFactory(Map<String, DataImportStrategy> 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;
}
}

112
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataExporter.java Normal file
View File

@ -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<Map<String, Object>> 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<String, Object> 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<Map<String, Object>> 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);
}
}

12
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/json/JsonDataImporter.java Normal file
View File

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

13
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/DatabaseDataService.java Normal file
View File

@ -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<Long> doExportAsync(DatabaseExportDataParam databaseExportDataParam);
}

150
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/service/impl/DatabaseDataImpl.java Normal file
View File

@ -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<Long> doExportAsync(DatabaseExportDataParam databaseExportDataParam) {
List<String> 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<Long> 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<Long> 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<String> 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);
}
}

184
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/sql/SqlDataExporter.java Normal file
View File

@ -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<String> sqlList = new ArrayList<>(BATCH_SIZE);
SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> {
List<String> header = containsHeader ? ResultSetUtils.getRsHeader(resultSet) : null;
while (resultSet.next()) {
List<String> 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<List<String>> dataList = new ArrayList<>(BATCH_SIZE);
List<String> 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<String> sqlList = new ArrayList<>(BATCH_SIZE);
SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> {
Map<String, String> primaryKeyMap = getPrimaryKeyMap(connection, databaseName, schemaName, tableName);
while (resultSet.next()) {
Map<String, String> 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<String> extractRowData(ResultSet resultSet, ValueProcessor valueProcessor) throws SQLException {
ResultSetMetaData metaData = resultSet.getMetaData();
List<String> 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<String, String> extractRowDataAsMap(ResultSet resultSet, ValueProcessor valueProcessor,
Map<String, String> primaryKeyMap) throws SQLException {
ResultSetMetaData metaData = resultSet.getMetaData();
Map<String, String> 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<String, String> getPrimaryKeyMap(Connection connection, String databaseName,
String schemaName, String tableName) throws SQLException {
Map<String, String> 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<String> sqlList) {
sqlList.forEach(writer::println);
sqlList.clear();
}
}

69
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskManager.java Normal file
View File

@ -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<Long> TASK_ID = new ThreadLocal<>();
public static final Map<Long, TaskState> 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();
}
}

28
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/task/TaskState.java Normal file
View File

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

24
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataExporter.java Normal file
View File

@ -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;
}
}

13
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xls/XlsDataImporter.java Normal file
View File

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

25
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataExporter.java Normal file
View File

@ -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;
}
}

12
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/xlsx/XlsxDataImporter.java Normal file
View File

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

31
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/factory/ExportDBDataStrategyFactory.java
View File

@ -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<String, Class<? extends ExportDBDataStrategy>> 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;
}
}
}

10
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/DatabaseExportDataRequest.java
View File

@ -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 @@ import lombok.NoArgsConstructor;
public class DatabaseExportDataRequest extends DataSourceBaseRequest {
@NotNull
private String exportType;
@NotEmpty
private List<String> tableNames;
/**
* single：单行插入，multi：多行插入，update：更新语句
*/
private String sqyType;
private Boolean containsHeader;
}

47
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/TaskController.java
View File

@ -8,19 +8,17 @@ import ai.chat2db.server.tools.base.wrapper.result.PageResult;
import ai.chat2db.server.tools.base.wrapper.result.web.WebPageResult;
import ai.chat2db.server.tools.common.util.ContextUtils;
import ai.chat2db.server.web.api.aspect.ConnectionInfoAspect;
import 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 class TaskController {
}
@GetMapping("/download/{id}")
public ResponseEntity<Resource> download(@PathVariable Long id) {
public void download(@PathVariable Long id, HttpServletResponse response) {
DataResult<Task> 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 (!file.exists() || !file.isFile()) {
log.error("File not found or is not a file: {}", file.getAbsolutePath());
throw new RuntimeException("File not found or accessible");
}
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!");
}
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);
}
}
}