diff --git a/internal/dao/kb.go b/internal/dao/kb.go index 1eda211460..e10ac12002 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -173,7 +173,7 @@ func (dao *KnowledgebaseDAO) Count(ctx context.Context, db *gorm.DB, filters map // GetByTenantIDs retrieves knowledge bases by tenant IDs with pagination // This matches the Python get_by_tenant_ids method -func (dao *KnowledgebaseDAO) GetByTenantIDs(ctx context.Context, db *gorm.DB, tenantIDs []string, userID string, pageNumber, itemsPerPage int, orderby string, desc bool, keywords, parserID, id, name string) ([]*entity.KnowledgebaseListItem, int64, error) { +func (dao *KnowledgebaseDAO) GetByTenantIDs(ctx context.Context, db *gorm.DB, tenantIDs []string, userID string, pageNumber, itemsPerPage int, orderby string, desc bool, keywords, parserID, id, name string, ids []string) ([]*entity.KnowledgebaseListItem, int64, error) { var kbs []*entity.KnowledgebaseListItem var total int64 @@ -193,6 +193,10 @@ func (dao *KnowledgebaseDAO) GetByTenantIDs(ctx context.Context, db *gorm.DB, te query = query.Where("knowledgebase.id = ?", id) } + if len(ids) > 0 { + query = query.Where("knowledgebase.id IN ?", ids) + } + if name != "" { query = query.Where("knowledgebase.name = ?", name) } @@ -318,6 +322,21 @@ func (dao *KnowledgebaseDAO) Accessible(ctx context.Context, db *gorm.DB, datase return count > 0 } +// GetAccessibleIDs returns the subset of ids that are visible to the user: +// team-permission KBs owned by any joined tenant, plus the user's own KBs. +// This matches the Python get_accessible_ids method. +func (dao *KnowledgebaseDAO) GetAccessibleIDs(ctx context.Context, db *gorm.DB, joinedTenantIDs []string, userID string, ids []string) ([]string, error) { + accessibleIDs := make([]string, 0, len(ids)) + err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). + Where("id IN ? AND ((tenant_id IN ? AND permission = ?) OR tenant_id = ?) AND status = ?", + ids, joinedTenantIDs, string(entity.TenantPermissionTeam), userID, string(entity.StatusValid)). + Pluck("id", &accessibleIDs).Error + if err != nil { + return nil, err + } + return accessibleIDs, nil +} + // Accessible4Deletion checks if a knowledge base can be deleted by a user // This matches the Python accessible4deletion method func (dao *KnowledgebaseDAO) Accessible4Deletion(ctx context.Context, db *gorm.DB, kbID, userID string) bool { diff --git a/internal/handler/dataset.go b/internal/handler/dataset.go index cb2748c86b..37feaec5e4 100644 --- a/internal/handler/dataset.go +++ b/internal/handler/dataset.go @@ -168,6 +168,39 @@ func (h *DatasetsHandler) ListDatasets(c *gin.Context) { } } + // Mirror pydantic's ListDatasetReq.ids: each occurrence is comma-split, + // every value must be a valid UUID, and duplicates are rejected. + var ids []string + if rawIDs, exists := c.Request.URL.Query()["ids"]; exists { + seen := make(map[string]int) + for _, item := range rawIDs { + for _, value := range strings.Split(item, ",") { + if value == "" { + continue + } + normalizedID, err := dataset.NormalizeDatasetID(strings.TrimSpace(value)) + if err != nil { + common.ResponseWithCodeData(c, common.CodeArgumentError, nil, err.Error()) + return + } + seen[normalizedID]++ + ids = append(ids, normalizedID) + } + } + duplicates := make([]string, 0, len(ids)) + reported := make(map[string]bool) + for _, normalizedID := range ids { + if seen[normalizedID] > 1 && !reported[normalizedID] { + reported[normalizedID] = true + duplicates = append(duplicates, normalizedID) + } + } + if len(duplicates) > 0 { + common.ResponseWithCodeData(c, common.CodeArgumentError, nil, fmt.Sprintf("Duplicate ids: '%s'", strings.Join(duplicates, ", "))) + return + } + } + ctx := c.Request.Context() data, total, code, err := h.datasetsService.ListDatasets( ctx, @@ -181,6 +214,7 @@ func (h *DatasetsHandler) ListDatasets(c *gin.Context) { ownerIDs, parserID, user.ID, + ids, ) if err != nil { common.ErrorWithCode(c, code, err.Error()) @@ -326,7 +360,7 @@ func pythonJSONTypeName(v interface{}) string { // ListDatasetReq (BaseListReq + include_parsing_status/ext; `type` is handled // before validation in the Python endpoint). var listDatasetsAllowedParams = map[string]bool{ - "id": true, "name": true, "page": true, "page_size": true, + "id": true, "ids": true, "name": true, "page": true, "page_size": true, "orderby": true, "desc": true, "include_parsing_status": true, "ext": true, "type": true, } diff --git a/internal/handler/dataset_list_test.go b/internal/handler/dataset_list_test.go new file mode 100644 index 0000000000..6792300946 --- /dev/null +++ b/internal/handler/dataset_list_test.go @@ -0,0 +1,147 @@ +package handler + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "gorm.io/gorm" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + dataset "ragflow/internal/service/dataset" +) + +const ( + listDatasetsTestKBID = "123e4567e89b12d3a456426614174000" + listDatasetsTestKBIDDashed = "123e4567-e89b-12d3-a456-426614174000" +) + +func setupListDatasetsTestDB(t *testing.T) *gorm.DB { + t.Helper() + + db, err := gorm.Open(sqlite.Open("file:"+url.QueryEscape(t.Name())+"?mode=memory&cache=shared"), &gorm.Config{ + TranslateError: true, + }) + if err != nil { + t.Fatalf("failed to open sqlite: %v", err) + } + + if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.User{}, &entity.UserTenant{}, &entity.Tenant{}); err != nil { + t.Fatalf("failed to migrate test schema: %v", err) + } + + origDB := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = origDB }) + + return db +} + +func insertListDatasetsTestKB(t *testing.T, id, tenantID, name string) { + t.Helper() + + status := string(entity.StatusValid) + kb := &entity.Knowledgebase{ + ID: id, + TenantID: tenantID, + Name: name, + EmbdID: "BAAI/bge-large-zh-v1.5@Builtin", + CreatedBy: tenantID, + Permission: string(entity.TenantPermissionMe), + ParserID: string(entity.ParserTypeNaive), + ParserConfig: entity.JSONMap{"chunk_token_num": float64(128)}, + Status: &status, + } + if err := dao.DB.Create(kb).Error; err != nil { + t.Fatalf("insert test kb: %v", err) + } +} + +func newListDatasetsTestRouter() *gin.Engine { + gin.SetMode(gin.TestMode) + h := NewDatasetsHandler(dataset.NewDatasetService(), nil) + r := gin.New() + r.GET("/api/v1/datasets", func(c *gin.Context) { + c.Set("user", &entity.User{ID: "user-1"}) + h.ListDatasets(c) + }) + return r +} + +type listDatasetsTestResponse struct { + Code int `json:"code"` + Data []map[string]interface{} `json:"data"` + Message string `json:"message"` + TotalDatasets int64 `json:"total_datasets"` +} + +func getListDatasets(t *testing.T, r *gin.Engine, rawQuery string) listDatasetsTestResponse { + t.Helper() + + resp := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/datasets?"+rawQuery, nil) + r.ServeHTTP(resp, req) + + if resp.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", resp.Code, resp.Body.String()) + } + + var body listDatasetsTestResponse + if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v body=%s", err, resp.Body.String()) + } + return body +} + +func TestDatasetsHandlerListDatasetsFiltersByIDs(t *testing.T) { + setupListDatasetsTestDB(t) + insertListDatasetsTestKB(t, listDatasetsTestKBID, "user-1", "Alpha") + + body := getListDatasets(t, newListDatasetsTestRouter(), + fmt.Sprintf("ids=%s&page_size=1", listDatasetsTestKBIDDashed)) + + if body.Code != int(common.CodeSuccess) { + t.Fatalf("code=%d message=%q", body.Code, body.Message) + } + if body.TotalDatasets != 1 || len(body.Data) != 1 { + t.Fatalf("expected exactly one dataset, got total=%d len=%d", body.TotalDatasets, len(body.Data)) + } + if body.Data[0]["id"] != listDatasetsTestKBID { + t.Fatalf("expected dataset id %q, got %#v", listDatasetsTestKBID, body.Data[0]["id"]) + } +} + +func TestDatasetsHandlerListDatasetsRejectsInvalidIDInIDs(t *testing.T) { + setupListDatasetsTestDB(t) + + body := getListDatasets(t, newListDatasetsTestRouter(), "ids=not-a-uuid") + + if body.Code != int(common.CodeArgumentError) { + t.Fatalf("code=%d want=%d", body.Code, common.CodeArgumentError) + } + if body.Message != "Invalid UUID format" { + t.Fatalf("message=%q want=%q", body.Message, "Invalid UUID format") + } +} + +func TestDatasetsHandlerListDatasetsRejectsDuplicateIDs(t *testing.T) { + setupListDatasetsTestDB(t) + + rawQuery := fmt.Sprintf("ids=%s,%s", listDatasetsTestKBIDDashed, listDatasetsTestKBIDDashed) + body := getListDatasets(t, newListDatasetsTestRouter(), rawQuery) + + if body.Code != int(common.CodeArgumentError) { + t.Fatalf("code=%d want=%d", body.Code, common.CodeArgumentError) + } + expected := fmt.Sprintf("Duplicate ids: '%s'", listDatasetsTestKBID) + if body.Message != expected { + t.Fatalf("message=%q want=%q", body.Message, expected) + } +} diff --git a/internal/handler/mcp_server.go b/internal/handler/mcp_server.go index fe0334ce80..d7a8c509d7 100644 --- a/internal/handler/mcp_server.go +++ b/internal/handler/mcp_server.go @@ -35,7 +35,7 @@ import ( // by the MCP server handler. type MCPRetrievalService interface { SearchDatasets(req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) - ListDatasets(id, name string, page, pageSize int, orderby string, desc bool, keywords string, ownerIDs []string, parserID, userID string) ([]map[string]interface{}, int64, common.ErrorCode, error) + ListDatasets(id, name string, page, pageSize int, orderby string, desc bool, keywords string, ownerIDs []string, parserID, userID string, ids []string) ([]map[string]interface{}, int64, common.ErrorCode, error) } // MCPServerHandler handles MCP protocol requests (JSON-RPC over HTTP). @@ -118,7 +118,7 @@ func (h *MCPServerHandler) HandleMCP(c *gin.Context) { func MCPListDatasets(ctx context.Context, ds *dataset.DatasetService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { data, total, _, err := ds.ListDatasets(ctx, "", "", page, pageSize, orderby, desc, - "", nil, "", userID, + "", nil, "", userID, nil, ) return data, total, err } @@ -154,7 +154,7 @@ func MCPRetrieval(ctx context.Context, ds *dataset.DatasetService, userID string for { data, _, _, err := ds.ListDatasets(ctx, "", "", page, maxPageSize, "create_time", true, - "", nil, "", userID, + "", nil, "", userID, nil, ) if err != nil { return "", fmt.Errorf("cannot resolve accessible datasets: %w", err) diff --git a/internal/service/dataset/crud.go b/internal/service/dataset/crud.go index 7f88d34ed7..959b2356b2 100644 --- a/internal/service/dataset/crud.go +++ b/internal/service/dataset/crud.go @@ -328,8 +328,11 @@ func (d *DatasetService) deleteDataset(tenantID string, kb *entity.Knowledgebase }) } -func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page, pageSize int, orderby string, desc bool, keywords string, ownerIDs []string, parserID, userID string) ([]map[string]interface{}, int64, common.ErrorCode, error) { +func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page, pageSize int, orderby string, desc bool, keywords string, ownerIDs []string, parserID, userID string, ids []string) ([]map[string]interface{}, int64, common.ErrorCode, error) { id = strings.TrimSpace(id) + if id != "" && len(ids) > 0 { + return nil, 0, common.CodeDataError, fmt.Errorf("Should not provide both 'id':%s and 'ids'%s", id, pythonStringListRepr(ids)) + } if id != "" { normalizedID, err := normalizeDatasetID(id) if err != nil { @@ -380,6 +383,7 @@ func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page } } queryUserID := userID + var joinedTenantIDs []string if len(tenantIDs) > 0 { joinedTenants, err := d.tenantDAO.GetJoinedTenantsByUserID(ctx, dao.DB, userID) if err != nil { @@ -391,6 +395,7 @@ func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page continue } allowedTenantIDs[joinedTenant.TenantID] = struct{}{} + joinedTenantIDs = append(joinedTenantIDs, joinedTenant.TenantID) } filteredTenantIDs := tenantIDs[:0] queryUserID = "" @@ -414,10 +419,33 @@ func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page continue } tenantIDs = append(tenantIDs, joinedTenant.TenantID) + joinedTenantIDs = append(joinedTenantIDs, joinedTenant.TenantID) } } - kbs, total, err := d.kbDAO.GetByTenantIDs(ctx, dao.DB, tenantIDs, queryUserID, page, pageSize, orderby, desc, keywords, parserID, id, name) + // Mirror Python: ids are checked for accessibility against the joined + // tenants (not the owner-filtered tenant list) before filtering. + if len(ids) > 0 { + accessibleIDs, err := d.kbDAO.GetAccessibleIDs(ctx, dao.DB, joinedTenantIDs, userID, ids) + if err != nil { + return nil, 0, common.CodeServerError, errors.New("database operation failed") + } + accessible := make(map[string]struct{}, len(accessibleIDs)) + for _, accessibleID := range accessibleIDs { + accessible[accessibleID] = struct{}{} + } + deniedIDs := make([]string, 0, len(ids)) + for _, datasetID := range ids { + if _, ok := accessible[datasetID]; !ok { + deniedIDs = append(deniedIDs, datasetID) + } + } + if len(deniedIDs) > 0 { + return nil, 0, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for datasets: '%s'", userID, strings.Join(deniedIDs, ", ")) + } + } + + kbs, total, err := d.kbDAO.GetByTenantIDs(ctx, dao.DB, tenantIDs, queryUserID, page, pageSize, orderby, desc, keywords, parserID, id, name, ids) if err != nil { return nil, 0, common.CodeServerError, errors.New("database operation failed") } diff --git a/internal/service/dataset/helpers.go b/internal/service/dataset/helpers.go index 36966489c2..d1559b594d 100644 --- a/internal/service/dataset/helpers.go +++ b/internal/service/dataset/helpers.go @@ -180,6 +180,16 @@ func normalizeDatasetID(id string) (string, error) { return strings.ReplaceAll(parsedUUID.String(), "-", ""), nil } +// pythonStringListRepr renders a string slice the way Python prints a list of +// strings, e.g. ['a', 'b'], for error messages that mirror the Python API. +func pythonStringListRepr(items []string) string { + quoted := make([]string, 0, len(items)) + for _, item := range items { + quoted = append(quoted, "'"+item+"'") + } + return "[" + strings.Join(quoted, ", ") + "]" +} + func canvasAccessibleForUser(ctx context.Context, userID, canvasID string) (bool, error) { tenantIDs, _ := dao.NewUserTenantDAO().GetTenantIDsByUserID(ctx, dao.DB, userID) return dao.NewUserCanvasDAO().Accessible(ctx, dao.DB, canvasID, userID, tenantIDs), nil diff --git a/internal/service/dataset/list_test.go b/internal/service/dataset/list_test.go new file mode 100644 index 0000000000..7c2e072c9c --- /dev/null +++ b/internal/service/dataset/list_test.go @@ -0,0 +1,119 @@ +package dataset + +import ( + "strings" + "testing" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" +) + +func testDatasetListService(t *testing.T) *DatasetService { + t.Helper() + + return &DatasetService{ + kbDAO: dao.NewKnowledgebaseDAO(), + documentDAO: dao.NewDocumentDAO(), + tenantDAO: dao.NewTenantDAO(), + } +} + +func TestDatasetServiceListDatasetsFiltersByIDs(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Alpha") + insertDatasetUpdateKB(t, "kb-2", "tenant-1", "Beta") + + ctx := t.Context() + data, total, code, err := testDatasetListService(t).ListDatasets(ctx, + "", "", 1, 30, "create_time", true, + "", nil, "", "tenant-1", []string{"kb-1"}, + ) + if err != nil { + t.Fatalf("ListDatasets failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + if total != 1 || len(data) != 1 { + t.Fatalf("expected exactly one dataset, got total=%d len=%d", total, len(data)) + } + if data[0]["id"] != "kb-1" { + t.Fatalf("expected kb-1, got %#v", data[0]["id"]) + } +} + +func TestDatasetServiceListDatasetsIDsAccessibleViaTeamTenant(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-team", "owner-1", "Shared") + if err := dao.DB.Create(&entity.Tenant{ID: "owner-1", Name: sptr("owner"), Status: sptr("1")}).Error; err != nil { + t.Fatalf("insert owner tenant: %v", err) + } + insertDatasetUpdateTeamMember(t, "user-1", "owner-1") + if err := dao.DB.Model(&entity.Knowledgebase{}). + Where("id = ?", "kb-team"). + Update("permission", string(entity.TenantPermissionTeam)).Error; err != nil { + t.Fatalf("update kb permission: %v", err) + } + + ctx := t.Context() + data, total, code, err := testDatasetListService(t).ListDatasets(ctx, + "", "", 1, 30, "create_time", true, + "", nil, "", "user-1", []string{"kb-team"}, + ) + if err != nil { + t.Fatalf("ListDatasets failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + if total != 1 || len(data) != 1 || data[0]["id"] != "kb-team" { + t.Fatalf("expected the shared dataset, got total=%d data=%#v", total, data) + } +} + +func TestDatasetServiceListDatasetsRejectsIDAndIDsTogether(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Alpha") + + ctx := t.Context() + _, _, code, err := testDatasetListService(t).ListDatasets(ctx, + "kb-1", "", 1, 30, "create_time", true, + "", nil, "", "tenant-1", []string{"kb-1"}, + ) + if err == nil { + t.Fatal("expected id/ids conflict error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + expected := "Should not provide both 'id':kb-1 and 'ids'['kb-1']" + if err.Error() != expected { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestDatasetServiceListDatasetsRejectsDeniedIDs(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-private", "owner-1", "Private") + + ctx := t.Context() + _, _, code, err := testDatasetListService(t).ListDatasets(ctx, + "", "", 1, 30, "create_time", true, + "", nil, "", "user-1", []string{"kb-private"}, + ) + if err == nil { + t.Fatal("expected permission error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + expected := "User 'user-1' lacks permission for datasets: 'kb-private'" + if !strings.Contains(err.Error(), expected) { + t.Fatalf("unexpected error: %v", err) + } +}