opensource/Chat2DB
1
0
Fork 0
mirror of https://github.com/CodePhiliaX/Chat2DB.git synced 2025-07-30 19:22:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

embedding query

Browse Source
This commit is contained in:
robin
2023-10-17 14:39:39 +08:00
parent 64c6fda770
commit 2beccd5cc4
1 changed files with 28 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/ChatController.java
View File

@ -13,6 +13,7 @@ import java.util.stream.Collectors;
import ai.chat2db.server.domain.api.enums.AiSqlSourceEnum; import ai.chat2db.server.domain.api.enums.AiSqlSourceEnum;
import ai.chat2db.server.domain.api.model.Config; import ai.chat2db.server.domain.api.model.Config;
import ai.chat2db.server.domain.api.model.DataSource; import ai.chat2db.server.domain.api.model.DataSource;
import ai.chat2db.server.domain.api.param.ShowCreateTableParam;
import ai.chat2db.server.domain.api.param.TableQueryParam; import ai.chat2db.server.domain.api.param.TableQueryParam;
import ai.chat2db.server.domain.api.service.ConfigService; import ai.chat2db.server.domain.api.service.ConfigService;
import ai.chat2db.server.domain.api.service.DataSourceService; import ai.chat2db.server.domain.api.service.DataSourceService;
@ -444,22 +445,39 @@ public class ChatController {
* @param tableNames * @param tableNames
* @return * @return
*/ */
private Map<String, List<TableColumn>> buildTableColumn(TableQueryParam tableQueryParam, private String buildTableColumn(TableQueryParam tableQueryParam,
List<String> tableNames) { List<String> tableNames) {
if (CollectionUtils.isEmpty(tableNames)) { if (CollectionUtils.isEmpty(tableNames)) {
return Maps.newHashMap(); return "";
} }
List<TableColumn> tableColumns = Lists.newArrayList(); List<String> schemaContent = Lists.newArrayList();
try { try {
tableColumns = tableService.queryColumns(tableQueryParam); schemaContent = tableNames.stream().map(tableName -> {
tableQueryParam.setTableName(tableName);
return queryTableDdl(tableName, tableQueryParam);
}).collect(Collectors.toList());
} catch (Exception exception) { } catch (Exception exception) {
log.error("query table error, do nothing"); log.error("query table error, do nothing");
} }
if (CollectionUtils.isEmpty(tableColumns)) {
return Maps.newHashMap(); return JSON.toJSONString(schemaContent);
} }
return tableColumns.stream().filter(tableColumn -> tableNames.contains(tableColumn.getTableName())).collect(
Collectors.groupingBy(TableColumn::getTableName, Collectors.toList())); /**
* query table schema
*
* @param tableName
* @param request
* @return
*/
private String queryTableDdl(String tableName, TableQueryParam request) {
ShowCreateTableParam param = new ShowCreateTableParam();
param.setTableName(tableName);
param.setDataSourceId(request.getDataSourceId());
param.setDatabaseName(request.getDatabaseName());
param.setSchemaName(request.getSchemaName());
DataResult<String> tableSchema = tableService.showCreateTable(param);
return tableSchema.getData();
} }
/** /**
@ -478,12 +496,7 @@ public class ChatController {
String properties = ""; String properties = "";
if (CollectionUtils.isNotEmpty(queryRequest.getTableNames())) { if (CollectionUtils.isNotEmpty(queryRequest.getTableNames())) {
TableQueryParam queryParam = chatConverter.chat2tableQuery(queryRequest); TableQueryParam queryParam = chatConverter.chat2tableQuery(queryRequest);
Map<String, List<TableColumn>> tableColumns = buildTableColumn(queryParam, queryRequest.getTableNames()); properties = buildTableColumn(queryParam, queryRequest.getTableNames());
List<String> tableSchemas = tableColumns.entrySet().stream().map(
entry -> String.format("%s(%s)", entry.getKey(),
entry.getValue().stream().map(TableColumn::getName).collect(
Collectors.joining(", ")))).collect(Collectors.toList());
properties = String.join("\n#", tableSchemas);
} else { } else {
properties = queryDatabaseSchema(queryRequest); properties = queryDatabaseSchema(queryRequest);
} }