Skip to content

fix(nl2sql): 修复表不支持联合主键的问题 #2206

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Sep 13, 2025
Merged

fix(nl2sql): 修复表不支持联合主键的问题 #2206

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
227c27d
fix(nl2sql): mul primary keys
flyingbird2 Aug 21, 2025
f3217cd
fix(nl2sql): Run `spring-javaformat:apply` to fix.
flyingbird2 Aug 25, 2025
440acc6
fix(studio): spring-javaformat:apply
flyingbird2 Aug 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 15 additions & 7 deletions ...2sql-chat/src/main/java/com/alibaba/cloud/ai/service/simple/SimpleVectorStoreService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import com.alibaba.cloud.ai.request.SearchRequest;
import com.alibaba.cloud.ai.service.base.BaseVectorStoreService;
import com.google.gson.Gson;
import org.apache.commons.collections.CollectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
Expand Down Expand Up @@ -94,6 +95,7 @@ protected EmbeddingModel getEmbeddingModel() {
* @param schemaInitRequest schema initialization request
* @throws Exception if an error occurs
*/
@Override
public Boolean schema(SchemaInitRequest schemaInitRequest) throws Exception {
log.info("Starting schema initialization for database: {}, schema: {}, tables: {}",
schemaInitRequest.getDbConfig().getUrl(), schemaInitRequest.getDbConfig().getSchema(),
Expand Down Expand Up @@ -173,12 +175,18 @@ private void processTable(TableInfoBO tableInfoBO, DbQueryParameter dqp, DbConfi
columnInfoBO.setSamples(gson.toJson(sampleColumn));
}

ColumnInfoBO primaryColumnDO = columnInfoBOS.stream()
List<ColumnInfoBO> targetPrimaryList = columnInfoBOS.stream()
.filter(ColumnInfoBO::isPrimary)
.findFirst()
.orElse(new ColumnInfoBO());

tableInfoBO.setPrimaryKey(primaryColumnDO.getName());
.collect(Collectors.toList());
if (CollectionUtils.isNotEmpty(targetPrimaryList)) {
List<String> columnNames = targetPrimaryList.stream()
.map(ColumnInfoBO::getName)
.collect(Collectors.toList());
tableInfoBO.setPrimaryKeys(columnNames);
}
else {
tableInfoBO.setPrimaryKeys(new ArrayList<>());
}
tableInfoBO
.setForeignKey(String.join("、", foreignKeyMap.getOrDefault(tableInfoBO.getName(), new ArrayList<>())));
}
Expand Down Expand Up @@ -216,7 +224,7 @@ public Document convertTableToDocument(TableInfoBO tableInfoBO) {
metadata.put("name", tableInfoBO.getName());
metadata.put("description", Optional.ofNullable(tableInfoBO.getDescription()).orElse(""));
metadata.put("foreignKey", Optional.ofNullable(tableInfoBO.getForeignKey()).orElse(""));
metadata.put("primaryKey", Optional.ofNullable(tableInfoBO.getPrimaryKey()).orElse(""));
metadata.put("primaryKey", Optional.ofNullable(tableInfoBO.getPrimaryKeys()).orElse(new ArrayList<>()));
metadata.put("vectorType", "table");
Document document = new Document(tableInfoBO.getName(), text, metadata);
log.debug("Created table document with ID: {}", tableInfoBO.getName());
Expand Down Expand Up @@ -467,7 +475,7 @@ private Document convertTableToDocumentForAgent(String agentId, TableInfoBO tabl
metadata.put("name", tableInfoBO.getName());
metadata.put("description", Optional.ofNullable(tableInfoBO.getDescription()).orElse(""));
metadata.put("foreignKey", Optional.ofNullable(tableInfoBO.getForeignKey()).orElse(""));
metadata.put("primaryKey", Optional.ofNullable(tableInfoBO.getPrimaryKey()).orElse(""));
metadata.put("primaryKey", Optional.ofNullable(tableInfoBO.getPrimaryKeys()).orElse(new ArrayList<>()));
metadata.put("vectorType", "table");

Document document = new Document(id, text, metadata);
Expand Down
3 changes: 2 additions & 1 deletion ...-chat/src/test/java/com/alibaba/cloud/ai/service/simple/SimpleVectorStoreServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import com.alibaba.cloud.ai.request.DeleteRequest;
import com.alibaba.cloud.ai.request.SchemaInitRequest;
import com.alibaba.cloud.ai.request.SearchRequest;
import com.google.common.collect.Lists;
import com.google.gson.Gson;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Disabled;
Expand Down Expand Up @@ -435,7 +436,7 @@ private TableInfoBO createMockTableInfo() {
.name("test_table")
.description("Test table")
.schema("test_schema")
.primaryKey("id")
.primaryKeys(Lists.newArrayList("id"))
.foreignKey("foreign_key_info")
.build();
}
Expand Down
37 changes: 20 additions & 17 deletions ...ai-alibaba-nl2sql-common/src/main/java/com/alibaba/cloud/ai/connector/bo/TableInfoBO.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
*/
package com.alibaba.cloud.ai.connector.bo;

import java.util.List;
import java.util.Objects;

public class TableInfoBO extends DdlBaseBO {
Expand All @@ -29,19 +30,19 @@ public class TableInfoBO extends DdlBaseBO {

private String foreignKey;

private String primaryKey;
private List<String> primaryKeys;

public TableInfoBO() {
}

public TableInfoBO(String schema, String name, String description, String type, String foreignKey,
String primaryKey) {
List<String> primaryKeys) {
this.schema = schema;
this.name = name;
this.description = description;
this.type = type;
this.foreignKey = foreignKey;
this.primaryKey = primaryKey;
this.primaryKeys = primaryKeys;
}

public String getSchema() {
Expand Down Expand Up @@ -84,36 +85,38 @@ public void setForeignKey(String foreignKey) {
this.foreignKey = foreignKey;
}

public String getPrimaryKey() {
return primaryKey;
public List<String> getPrimaryKeys() {
return primaryKeys;
}

public void setPrimaryKey(String primaryKey) {
this.primaryKey = primaryKey;
public void setPrimaryKeys(List<String> primaryKeys) {
this.primaryKeys = primaryKeys;
}

@Override
public String toString() {
return "TableInfoBO{" + "schema='" + schema + '\'' + ", name='" + name + '\'' + ", description='" + description
+ '\'' + ", type='" + type + '\'' + ", foreignKey='" + foreignKey + '\'' + ", primaryKey='" + primaryKey
+ '\'' + '}';
+ '\'' + ", type='" + type + '\'' + ", foreignKey='" + foreignKey + '\'' + ", primaryKeys="
+ primaryKeys + '}';
}

@Override
public boolean equals(Object o) {
if (this == o)
if (this == o) {
return true;
if (o == null || getClass() != o.getClass())
}
if (o == null || getClass() != o.getClass()) {
return false;
}
TableInfoBO that = (TableInfoBO) o;
return Objects.equals(schema, that.schema) && Objects.equals(name, that.name)
&& Objects.equals(description, that.description) && Objects.equals(type, that.type)
&& Objects.equals(foreignKey, that.foreignKey) && Objects.equals(primaryKey, that.primaryKey);
&& Objects.equals(foreignKey, that.foreignKey) && Objects.equals(primaryKeys, that.primaryKeys);
}

@Override
public int hashCode() {
return Objects.hash(schema, name, description, type, foreignKey, primaryKey);
return Objects.hash(schema, name, description, type, foreignKey, primaryKeys);
}

public static TableInfoBOBuilder builder() {
Expand All @@ -132,7 +135,7 @@ public static final class TableInfoBOBuilder {

private String foreignKey;

private String primaryKey;
private List<String> primaryKeys;

private TableInfoBOBuilder() {
}
Expand Down Expand Up @@ -166,8 +169,8 @@ public TableInfoBOBuilder foreignKey(String foreignKey) {
return this;
}

public TableInfoBOBuilder primaryKey(String primaryKey) {
this.primaryKey = primaryKey;
public TableInfoBOBuilder primaryKeys(List<String> primaryKeys) {
this.primaryKeys = primaryKeys;
return this;
}

Expand All @@ -178,7 +181,7 @@ public TableInfoBO build() {
tableInfoBO.setDescription(description);
tableInfoBO.setType(type);
tableInfoBO.setForeignKey(foreignKey);
tableInfoBO.setPrimaryKey(primaryKey);
tableInfoBO.setPrimaryKeys(primaryKeys);
return tableInfoBO;
}

Expand Down
18 changes: 13 additions & 5 deletions ...nt/src/main/java/com/alibaba/cloud/ai/service/AnalyticDbVectorStoreManagementService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.gson.Gson;
import org.apache.commons.collections.CollectionUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
Expand Down Expand Up @@ -271,12 +272,19 @@ private void processTable(TableInfoBO tableInfoBO, DbQueryParameter dqp, DbConfi
columnInfoBO.setSamples(gson.toJson(sampleColumn));
}

ColumnInfoBO primaryColumnDO = columnInfoBOS.stream()
List<ColumnInfoBO> targetPrimaryList = columnInfoBOS.stream()
.filter(ColumnInfoBO::isPrimary)
.findFirst()
.orElse(new ColumnInfoBO());
.collect(Collectors.toList());
if (CollectionUtils.isNotEmpty(targetPrimaryList)) {
List<String> columnNames = targetPrimaryList.stream()
.map(ColumnInfoBO::getName)
.collect(Collectors.toList());
tableInfoBO.setPrimaryKeys(columnNames);
}
else {
tableInfoBO.setPrimaryKeys(new ArrayList<>());
}

tableInfoBO.setPrimaryKey(primaryColumnDO.getName());
tableInfoBO.setForeignKey(String.join("、", buildForeignKeyList(tableInfoBO.getName())));
}

Expand Down Expand Up @@ -313,7 +321,7 @@ public Document convertTableToDocument(TableInfoBO tableInfoBO) {
Map<String, Object> metadata = Map.of("schema", Optional.ofNullable(tableInfoBO.getSchema()).orElse(""), "name",
tableInfoBO.getName(), "description", Optional.ofNullable(tableInfoBO.getDescription()).orElse(""),
"foreignKey", Optional.ofNullable(tableInfoBO.getForeignKey()).orElse(""), "primaryKey",
Optional.ofNullable(tableInfoBO.getPrimaryKey()).orElse(""), "vectorType", "table");
Optional.ofNullable(tableInfoBO.getPrimaryKeys()).orElse(new ArrayList<>()), "vectorType", "table");
return new Document(tableInfoBO.getName(), text, metadata);
}

Expand Down
19 changes: 13 additions & 6 deletions ...gement/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreManagementService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import com.alibaba.cloud.ai.request.EvidenceRequest;
import com.alibaba.cloud.ai.request.SchemaInitRequest;
import com.google.gson.Gson;
import org.apache.commons.collections.CollectionUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.vectorstore.SearchRequest;
Expand Down Expand Up @@ -150,12 +151,18 @@ private void processTable(TableInfoBO tableInfoBO, DbQueryParameter dqp, DbConfi
columnInfoBO.setSamples(gson.toJson(sampleColumn));
}

ColumnInfoBO primaryColumnDO = columnInfoBOS.stream()
List<ColumnInfoBO> targetPrimaryList = columnInfoBOS.stream()
.filter(ColumnInfoBO::isPrimary)
.findFirst()
.orElse(new ColumnInfoBO());

tableInfoBO.setPrimaryKey(primaryColumnDO.getName());
.collect(Collectors.toList());
if (CollectionUtils.isNotEmpty(targetPrimaryList)) {
List<String> columnNames = targetPrimaryList.stream()
.map(ColumnInfoBO::getName)
.collect(Collectors.toList());
tableInfoBO.setPrimaryKeys(columnNames);
}
else {
tableInfoBO.setPrimaryKeys(new ArrayList<>());
}
tableInfoBO.setForeignKey(String.join("、", buildForeignKeyList(tableInfoBO.getName())));
}

Expand All @@ -180,7 +187,7 @@ public Document convertTableToDocument(TableInfoBO tableInfoBO) {
Map<String, Object> metadata = Map.of("schema", Optional.ofNullable(tableInfoBO.getSchema()).orElse(""), "name",
tableInfoBO.getName(), "description", Optional.ofNullable(tableInfoBO.getDescription()).orElse(""),
"foreignKey", Optional.ofNullable(tableInfoBO.getForeignKey()).orElse(""), "primaryKey",
Optional.ofNullable(tableInfoBO.getPrimaryKey()).orElse(""), "vectorType", "table");
Optional.ofNullable(tableInfoBO.getPrimaryKeys()).orElse(new ArrayList<>()), "vectorType", "table");
return new Document(tableInfoBO.getName(), text, metadata);
}

Expand Down
Loading