Skip to content

Commit 96e6340

Browse files
Chat2DB-ProChat2DB-Pro
authored
Merge pull request #1432 from openai0229/data-import
data export
2 parents 56b7666 + cd16555 commit 96e6340

File tree

38 files changed

+1133
-370
lines changed

38 files changed

+1133
-370
lines changed

‎chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilder.java

+7-7
Original file line numberDiff line numberDiff line change
@@ -385,11 +385,11 @@ private static String[] moveElement(String[] originalArray, int from, int to, St
385385
}
386386

387387

388-
// @Override
389-
// protected void buildTableName(String databaseName, String schemaName, String tableName, StringBuilder script) {
390-
// if (StringUtils.isNotBlank(databaseName)) {
391-
// script.append("`").append(databaseName).append("`").append('.');
392-
// }
393-
// script.append("`").append(tableName).append("`");
394-
// }
388+
@Override
389+
protected void buildTableName(String databaseName, String schemaName, String tableName, StringBuilder script) {
390+
if (StringUtils.isNotBlank(databaseName)) {
391+
script.append("`").append(databaseName).append("`").append('.');
392+
}
393+
script.append("`").append(tableName).append("`");
394+
}
395395
}

‎chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/ExportFileSuffix.java

+1
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ public enum ExportFileSuffix {
1414
WORD(".docx"),
1515
//excel
1616
EXCEL(".xlsx"),
17+
XLS(".xls"),
1718
//markdown
1819
MARKDOWN(".md"),
1920
//html

‎chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/datasource/DatabaseExportDataParam.java

+7
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
import lombok.Data;
55
import lombok.NoArgsConstructor;
66

7+
import java.util.List;
8+
79
/**
810
* @author: zgq
911
* @date: 2024年03月24日 13:17
@@ -12,7 +14,12 @@
1214
@AllArgsConstructor
1315
@NoArgsConstructor
1416
public class DatabaseExportDataParam {
17+
private Long dataSourceId;
1518
private String databaseName;
1619
private String schemaName;
1720
private String exportType;
21+
private List<String> tableNames;
22+
private String sqyType;
23+
private Boolean containsHeader;
24+
1825
}

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

+1
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ public ActionResult updateStatus(TaskUpdateParam param) {
4444
taskDO.setId(param.getId());
4545
taskDO.setTaskStatus(param.getTaskStatus());
4646
taskDO.setContent(param.getContent());
47+
taskDO.setDownloadUrl(param.getDownloadUrl());
4748
MapperUtils.getTaskMapper().updateById(taskDO);
4849
return ActionResult.isSuccess();
4950
}

‎chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyStringUtils.java

+6-6
Original file line numberDiff line numberDiff line change
@@ -193,17 +193,17 @@ public static String escapeAndQuoteString(String value) {
193193
}
194194

195195
/**
196-
* @param value "abcd"
197-
* @param quoteChar '%'
198-
* @return "%abcd%"
196+
* @param value "abcd"
197+
* @param quoteChar '%'
198+
* @return "%abcd%"
199199
*/
200200
public static String quoteString(String value, char quoteChar) {
201-
return StringUtils.wrap(value, quoteChar);
201+
return quoteChar + value + quoteChar;
202202
}
203203

204204
/**
205-
* @param value "abcd"
206-
* @return "'abcd'"
205+
* @param value "abcd"
206+
* @return "'abcd'"
207207
*/
208208
public static String quoteString(String value) {
209209
// (char)39 -> '

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/DatabaseController.java

+29-22
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package ai.chat2db.server.web.api.controller.rdb;
22

3+
import ai.chat2db.server.domain.api.enums.TaskStatusEnum;
34
import ai.chat2db.server.domain.api.param.MetaDataQueryParam;
45
import ai.chat2db.server.domain.api.param.datasource.DatabaseCreateParam;
56
import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam;
@@ -14,8 +15,9 @@
1415
import ai.chat2db.server.web.api.controller.data.source.vo.DatabaseVO;
1516
import ai.chat2db.server.web.api.controller.rdb.converter.DatabaseConverter;
1617
import ai.chat2db.server.web.api.controller.rdb.converter.RdbWebConverter;
17-
import ai.chat2db.server.web.api.controller.rdb.data.export.strategy.ExportDBDataStrategy;
18-
import ai.chat2db.server.web.api.controller.rdb.factory.ExportDBDataStrategyFactory;
18+
import ai.chat2db.server.web.api.controller.rdb.data.service.DatabaseDataService;
19+
import ai.chat2db.server.web.api.controller.rdb.data.task.TaskManager;
20+
import ai.chat2db.server.web.api.controller.rdb.data.task.TaskState;
1921
import ai.chat2db.server.web.api.controller.rdb.request.DatabaseCreateRequest;
2022
import ai.chat2db.server.web.api.controller.rdb.request.DatabaseExportDataRequest;
2123
import ai.chat2db.server.web.api.controller.rdb.request.DatabaseExportRequest;
@@ -26,12 +28,12 @@
2628
import ai.chat2db.spi.model.Sql;
2729
import jakarta.servlet.http.HttpServletResponse;
2830
import jakarta.validation.Valid;
31+
import lombok.extern.slf4j.Slf4j;
2932
import org.apache.commons.lang3.StringUtils;
3033
import org.springframework.beans.factory.annotation.Autowired;
3134
import org.springframework.web.bind.annotation.*;
3235

3336
import java.io.PrintWriter;
34-
import java.lang.reflect.Constructor;
3537
import java.util.Objects;
3638

3739
/**
@@ -40,6 +42,7 @@
4042
@ConnectionInfoAspect
4143
@RequestMapping("/api/rdb/database")
4244
@RestController
45+
@Slf4j
4346
public class DatabaseController {
4447
@Autowired
4548
private RdbWebConverter rdbWebConverter;
@@ -49,6 +52,8 @@ public class DatabaseController {
4952

5053
@Autowired
5154
public DatabaseConverter databaseConverter;
55+
@Autowired
56+
private DatabaseDataService databaseDataService;
5257

5358
/**
5459
* Query the database_schema_list contained in the database
@@ -59,8 +64,8 @@ public class DatabaseController {
5964
@GetMapping("/database_schema_list")
6065
public DataResult<MetaSchemaVO> databaseSchemaList(@Valid DataSourceBaseRequest request) {
6166
MetaDataQueryParam queryParam = MetaDataQueryParam.builder().dataSourceId(request.getDataSourceId())
62-
.refresh(
63-
request.isRefresh()).build();
67+
.refresh(
68+
request.isRefresh()).build();
6469
DataResult<MetaSchema> result = databaseService.queryDatabaseSchema(queryParam);
6570
MetaSchemaVO schemaDto2vo = rdbWebConverter.metaSchemaDto2vo(result.getData());
6671
return DataResult.of(schemaDto2vo);
@@ -69,8 +74,8 @@ public DataResult<MetaSchemaVO> databaseSchemaList(@Valid DataSourceBaseRequest
6974
@GetMapping("list")
7075
public ListResult<DatabaseVO> databaseList(@Valid DataSourceBaseRequest request) {
7176
DatabaseQueryAllParam queryParam = DatabaseQueryAllParam.builder().dataSourceId(request.getDataSourceId())
72-
.refresh(
73-
request.isRefresh()).build();
77+
.refresh(
78+
request.isRefresh()).build();
7479
ListResult<Database> result = databaseService.queryAll(queryParam);
7580
return ListResult.of(rdbWebConverter.databaseDto2vo(result.getData()));
7681
}
@@ -95,7 +100,7 @@ public ActionResult deleteDatabase(@Valid @RequestBody DataSourceBaseRequest req
95100
*/
96101
@PostMapping("/create_database_sql")
97102
public DataResult<Sql> createDatabase(@Valid @RequestBody DatabaseCreateRequest request) {
98-
if(StringUtils.isBlank(request.getName())){
103+
if (StringUtils.isBlank(request.getName())) {
99104
request.setName(request.getDatabaseName());
100105
}
101106
Database database = databaseConverter.request2param(request);
@@ -111,12 +116,13 @@ public DataResult<Sql> createDatabase(@Valid @RequestBody DatabaseCreateRequest
111116
@PostMapping("/modify_database")
112117
public ActionResult modifyDatabase(@Valid @RequestBody UpdateDatabaseRequest request) {
113118
DatabaseCreateParam param = DatabaseCreateParam.builder().name(request.getDatabaseName())
114-
.name(request.getNewDatabaseName()).build();
119+
.name(request.getNewDatabaseName()).build();
115120
return databaseService.modifyDatabase(param);
116121
}
122+
117123
@PostMapping("/export")
118-
public void exportDatabase(@Valid @RequestBody DatabaseExportRequest request, HttpServletResponse response){
119-
String fileName = Objects.isNull(request.getSchemaName())?request.getDatabaseName() : request.getSchemaName();
124+
public void exportDatabase(@Valid @RequestBody DatabaseExportRequest request, HttpServletResponse response) {
125+
String fileName = Objects.isNull(request.getSchemaName()) ? request.getDatabaseName() : request.getSchemaName();
120126
response.setContentType("text/sql");
121127
response.setHeader("Content-disposition", "attachment;filename*=utf-8''" + fileName + ".sql");
122128
response.setCharacterEncoding("utf-8");
@@ -130,17 +136,18 @@ public void exportDatabase(@Valid @RequestBody DatabaseExportRequest request, Ht
130136
}
131137

132138
@PostMapping("/export_data")
133-
public void exportData(@Valid @RequestBody DatabaseExportDataRequest request, HttpServletResponse response) {
134-
Class<?> targetClass = ExportDBDataStrategyFactory.get(request.getExportType());
135-
response.setCharacterEncoding("utf-8");
136-
DatabaseExportDataParam param = databaseConverter.request2param(request);
137-
try {
138-
Constructor<?> constructor = targetClass.getDeclaredConstructor();
139-
ExportDBDataStrategy service = (ExportDBDataStrategy) constructor.newInstance();
140-
service.doExport(param, response);
141-
} catch (Exception e) {
142-
throw new RuntimeException(e);
143-
}
139+
public DataResult<Long> exportData(@Valid @RequestBody DatabaseExportDataRequest request) {
140+
DatabaseExportDataParam databaseExportDataParam = databaseConverter.request2param(request);
141+
return databaseDataService.doExportAsync(databaseExportDataParam);
142+
}
144143

144+
@GetMapping("/export_data_status/{taskId}")
145+
public DataResult<String> exportDataStatus(@PathVariable("taskId") Long taskId) {
146+
TaskState task = TaskManager.getTask(taskId);
147+
String state = task.getState();
148+
if (Objects.equals(state, TaskStatusEnum.FINISH.name()) || Objects.equals(state, TaskStatusEnum.ERROR.name())) {
149+
TaskManager.removeTask(taskId);
150+
}
151+
return DataResult.of(task.getExportStatus());
145152
}
146153
}

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataExporter.java

+70
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package ai.chat2db.server.web.api.controller.rdb.data;
2+
3+
import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam;
4+
import ai.chat2db.server.web.api.util.StringUtils;
5+
import ai.chat2db.spi.sql.Chat2DBContext;
6+
import lombok.extern.slf4j.Slf4j;
7+
import org.apache.commons.collections4.CollectionUtils;
8+
9+
import java.io.*;
10+
import java.sql.Connection;
11+
import java.sql.SQLException;
12+
import java.util.List;
13+
import java.util.zip.ZipEntry;
14+
import java.util.zip.ZipOutputStream;
15+
16+
/**
17+
* @author: zgq
18+
* @date: 2024年06月04日 10:51
19+
*/
20+
@Slf4j
21+
public abstract class BaseDataExporter implements DataExportStrategy {
22+
23+
protected String contentType;
24+
protected String suffix;
25+
26+
public static int BATCH_SIZE = 1000;
27+
28+
@Override
29+
public void doExport(DatabaseExportDataParam databaseExportDataParam, File file) throws IOException, SQLException {
30+
List<String> tableNames = databaseExportDataParam.getTableNames();
31+
if (CollectionUtils.isEmpty(tableNames)) {
32+
throw new IllegalArgumentException("tableNames should not be null or empty");
33+
}
34+
try (Connection connection = Chat2DBContext.getConnection()) {
35+
if (tableNames.size() == 1) {
36+
String tableName = tableNames.get(0);
37+
if (StringUtils.isEmpty(tableName)) {
38+
throw new IllegalArgumentException("tableName should not be null or empty");
39+
}
40+
singleExport(connection, databaseExportDataParam,file);
41+
} else {
42+
multiExport(databaseExportDataParam, connection, file);
43+
}
44+
}
45+
46+
}
47+
48+
49+
private void multiExport(DatabaseExportDataParam databaseExportDataParam,
50+
Connection connection, File file) throws IOException {
51+
try (OutputStream outputStream = new FileOutputStream(file);
52+
ZipOutputStream zipOutputStream = new ZipOutputStream(outputStream)) {
53+
List<String> tableNames = databaseExportDataParam.getTableNames();
54+
for (String tableName : tableNames) {
55+
String fileName = tableName + suffix;
56+
zipOutputStream.putNextEntry(new ZipEntry(fileName));
57+
try (ByteArrayOutputStream byteArrayOutputStream = multiExport(connection, databaseExportDataParam, tableName)) {
58+
byteArrayOutputStream.writeTo(zipOutputStream);
59+
zipOutputStream.closeEntry();
60+
}
61+
}
62+
}
63+
}
64+
65+
66+
protected abstract void singleExport(Connection connectionInfo, DatabaseExportDataParam databaseExportDataParam, File file) throws IOException, SQLException;
67+
68+
69+
protected abstract ByteArrayOutputStream multiExport(Connection connection, DatabaseExportDataParam databaseExportDataParam, String tableName);
70+
}

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseDataImporter.java

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
package ai.chat2db.server.web.api.controller.rdb.data;
2+
3+
/**
4+
* @author: zgq
5+
* @date: 2024年06月04日 10:52
6+
*/
7+
public abstract class BaseDataImporter implements DataImportStrategy{
8+
}

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/data/BaseExcelExporter.java

+105
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
package ai.chat2db.server.web.api.controller.rdb.data;
2+
3+
import ai.chat2db.server.domain.api.enums.TaskStatusEnum;
4+
import ai.chat2db.server.domain.api.param.datasource.DatabaseExportDataParam;
5+
import ai.chat2db.server.web.api.controller.rdb.data.task.TaskManager;
6+
import ai.chat2db.spi.ValueProcessor;
7+
import ai.chat2db.spi.model.JDBCDataValue;
8+
import ai.chat2db.spi.sql.Chat2DBContext;
9+
import ai.chat2db.spi.sql.SQLExecutor;
10+
import ai.chat2db.spi.util.ResultSetUtils;
11+
import com.alibaba.excel.EasyExcel;
12+
import com.alibaba.excel.support.ExcelTypeEnum;
13+
import com.alibaba.excel.write.builder.ExcelWriterSheetBuilder;
14+
import lombok.extern.slf4j.Slf4j;
15+
16+
import java.io.*;
17+
import java.sql.Connection;
18+
import java.sql.ResultSet;
19+
import java.sql.ResultSetMetaData;
20+
import java.sql.SQLException;
21+
import java.util.ArrayList;
22+
import java.util.Collections;
23+
import java.util.List;
24+
import java.util.stream.Collectors;
25+
26+
/**
27+
* @author: zgq
28+
* @date: 2024年06月04日 10:56
29+
*/
30+
@Slf4j
31+
public abstract class BaseExcelExporter extends BaseDataExporter {
32+
@Override
33+
protected void singleExport(Connection connection, DatabaseExportDataParam exportParam, File outputFile) {
34+
ExcelTypeEnum excelType = getExcelType();
35+
try (OutputStream outputStream = new FileOutputStream(outputFile)) {
36+
37+
38+
String tableName = exportParam.getTableNames().get(0);
39+
String querySql = getQuerySql(exportParam, tableName);
40+
41+
log.info("开始导出：{}表数据，导出类型：{}", tableName, excelType);
42+
43+
SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet ->
44+
writeExcelData(resultSet, excelType, outputStream, tableName, exportParam.getContainsHeader()));
45+
46+
} catch (IOException e) {
47+
TaskManager.updateStatus(TaskStatusEnum.ERROR);
48+
throw new RuntimeException(e);
49+
}
50+
}
51+
52+
@Override
53+
protected ByteArrayOutputStream multiExport(Connection connection, DatabaseExportDataParam databaseExportDataParam, String tableName) {
54+
ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
55+
ExcelTypeEnum excelType = getExcelType();
56+
57+
log.info("开始导出：{}表数据，导出类型：{}", tableName, excelType);
58+
59+
String querySql = getQuerySql(databaseExportDataParam, tableName);
60+
SQLExecutor.getInstance().execute(connection, querySql, BATCH_SIZE, resultSet -> {
61+
writeExcelData(resultSet, excelType, byteArrayOutputStream, tableName, databaseExportDataParam.getContainsHeader());
62+
});
63+
return byteArrayOutputStream;
64+
}
65+
66+
private void writeExcelData(ResultSet resultSet, ExcelTypeEnum excelType, OutputStream outputStream, String sheetName, Boolean containsHeader) {
67+
try {
68+
ExcelWriterSheetBuilder excelWriterSheetBuilder = EasyExcel.write(outputStream).excelType(excelType).sheet(sheetName);
69+
ResultSetMetaData metaData = resultSet.getMetaData();
70+
int columnCount = metaData.getColumnCount();
71+
ValueProcessor valueProcessor = Chat2DBContext.getMetaData().getValueProcessor();
72+
List<List<Object>> dataList = new ArrayList<>();
73+
74+
if (containsHeader) {
75+
List<String> header = ResultSetUtils.getRsHeader(resultSet);
76+
excelWriterSheetBuilder.head(header.stream().map(Collections::singletonList).collect(Collectors.toList()));
77+
}
78+
79+
while (resultSet.next()) {
80+
List<Object> rowDataList = new ArrayList<>();
81+
for (int i = 1; i <= columnCount; i++) {
82+
JDBCDataValue jdbcDataValue = new JDBCDataValue(resultSet, metaData, i, false);
83+
rowDataList.add(valueProcessor.getJdbcValue(jdbcDataValue));
84+
}
85+
dataList.add(rowDataList);
86+
}
87+
88+
excelWriterSheetBuilder.doWrite(dataList);
89+
TaskManager.increaseCurrent();
90+
} catch (SQLException e) {
91+
TaskManager.updateStatus(TaskStatusEnum.ERROR);
92+
log.error("Error writing Excel data", e);
93+
throw new RuntimeException(e);
94+
}
95+
}
96+
97+
private String getQuerySql(DatabaseExportDataParam databaseExportDataParam, String tableName) {
98+
String databaseName = databaseExportDataParam.getDatabaseName();
99+
String schemaName = databaseExportDataParam.getSchemaName();
100+
return Chat2DBContext.getSqlBuilder().buildTableQuerySql(databaseName, schemaName, tableName);
101+
}
102+
103+
protected abstract ExcelTypeEnum getExcelType();
104+
}
105+

0 commit comments

Comments
 (0)