diff --git a/internal/dao/connector.go b/internal/dao/connector.go index 66f6dd3607..79e326d881 100644 --- a/internal/dao/connector.go +++ b/internal/dao/connector.go @@ -18,6 +18,7 @@ package dao import ( "errors" + "fmt" "ragflow/internal/entity" "ragflow/internal/utility" "time" @@ -28,6 +29,13 @@ import ( // ConnectorDAO connector data access object type ConnectorDAO struct{} +var errConnectorNotAccessible = errors.New("connector is not accessible for this tenant") + +// IsConnectorNotAccessibleErr reports connector ownership mismatches. +func IsConnectorNotAccessibleErr(err error) bool { + return errors.Is(err, errConnectorNotAccessible) +} + // NewConnectorDAO create connector DAO func NewConnectorDAO() *ConnectorDAO { return &ConnectorDAO{} @@ -69,9 +77,14 @@ func (dao *ConnectorDAO) ListByTenantID(tenantID string) ([]*ConnectorListItem, // ListByDatasetID lists connectors linked to a dataset. func (dao *ConnectorDAO) ListByDatasetID(datasetID string) ([]*ConnectorDatasetListItem, error) { + return dao.ListByDatasetIDTx(DB, datasetID) +} + +// ListByDatasetIDTx lists connectors linked to a dataset using the caller's DB handle. +func (dao *ConnectorDAO) ListByDatasetIDTx(db *gorm.DB, datasetID string) ([]*ConnectorDatasetListItem, error) { var connectors []*ConnectorDatasetListItem - err := DB.Model(&entity.Connector2Kb{}). + err := db.Model(&entity.Connector2Kb{}). Select("connector.id, connector.source, connector.name, connector2kb.auto_parse, connector.status"). Joins("JOIN connector ON connector2kb.connector_id = connector.id"). Where("connector2kb.kb_id = ?", datasetID). @@ -93,74 +106,87 @@ type DatasetConnectorLink struct { // LinkDatasetConnectors syncs connector2kb rows for a dataset. func (dao *ConnectorDAO) LinkDatasetConnectors(kbID string, connectors []DatasetConnectorLink) error { return DB.Transaction(func(tx *gorm.DB) error { - var existing []entity.Connector2Kb - if err := tx.Where("kb_id = ?", kbID).Find(&existing).Error; err != nil { + var kb entity.Knowledgebase + if err := tx.Select("tenant_id").Where("id = ? AND status = ?", kbID, string(entity.StatusValid)).First(&kb).Error; err != nil { + return err + } + return dao.LinkDatasetConnectorsTx(tx, kbID, kb.TenantID, connectors) + }) +} + +// LinkDatasetConnectorsTx syncs connector2kb rows using the caller's transaction. +func (dao *ConnectorDAO) LinkDatasetConnectorsTx(tx *gorm.DB, kbID, tenantID string, connectors []DatasetConnectorLink) error { + var existing []entity.Connector2Kb + if err := tx.Where("kb_id = ?", kbID).Find(&existing).Error; err != nil { + return err + } + + oldConnectorIDs := make(map[string]entity.Connector2Kb, len(existing)) + for _, row := range existing { + oldConnectorIDs[row.ConnectorID] = row + } + + nextConnectorIDs := make(map[string]struct{}, len(connectors)) + for _, connector := range connectors { + nextConnectorIDs[connector.ID] = struct{}{} + autoParse := connector.AutoParse + if autoParse == "" { + autoParse = "1" + } + + var fullConnector entity.Connector + if err := tx.Where("id = ? AND tenant_id = ?", connector.ID, tenantID).First(&fullConnector).Error; err != nil { + if IsNotFoundErr(err) { + return fmt.Errorf("%w: %s", errConnectorNotAccessible, connector.ID) + } return err } - oldConnectorIDs := make(map[string]entity.Connector2Kb, len(existing)) - for _, row := range existing { - oldConnectorIDs[row.ConnectorID] = row + if _, ok := oldConnectorIDs[connector.ID]; ok { + if err := tx.Model(&entity.Connector2Kb{}). + Where("connector_id = ? AND kb_id = ?", connector.ID, kbID). + Update("auto_parse", autoParse).Error; err != nil { + return err + } + continue } - nextConnectorIDs := make(map[string]struct{}, len(connectors)) - for _, connector := range connectors { - nextConnectorIDs[connector.ID] = struct{}{} - autoParse := connector.AutoParse - if autoParse == "" { - autoParse = "1" - } - - if _, ok := oldConnectorIDs[connector.ID]; ok { - if err := tx.Model(&entity.Connector2Kb{}). - Where("connector_id = ? AND kb_id = ?", connector.ID, kbID). - Update("auto_parse", autoParse).Error; err != nil { - return err - } - continue - } - - if err := tx.Create(&entity.Connector2Kb{ - ID: utility.GenerateUUID(), - ConnectorID: connector.ID, - KbID: kbID, - AutoParse: autoParse, - }).Error; err != nil { - return err - } - - if err := scheduleConnectorTask(tx, connector.ID, kbID, connectorTaskTypeSync, true); err != nil { - return err - } - - var fullConnector entity.Connector - if err := tx.Where("id = ?", connector.ID).First(&fullConnector).Error; err != nil { - return err - } - if connectorConfigBool(fullConnector.Config, "sync_deleted_files") { - if err := scheduleConnectorTask(tx, connector.ID, kbID, connectorTaskTypePrune, false); err != nil { - return err - } - } + if err := tx.Create(&entity.Connector2Kb{ + ID: utility.GenerateUUID(), + ConnectorID: connector.ID, + KbID: kbID, + AutoParse: autoParse, + }).Error; err != nil { + return err } - for connectorID := range oldConnectorIDs { - if _, ok := nextConnectorIDs[connectorID]; ok { - continue - } - if err := tx.Where("kb_id = ? AND connector_id = ?", kbID, connectorID). - Delete(&entity.Connector2Kb{}).Error; err != nil { - return err - } - if err := tx.Model(&entity.SyncLogs{}). - Where("connector_id = ? AND kb_id = ? AND status IN ?", connectorID, kbID, []string{string(entity.TaskStatusSchedule), string(entity.TaskStatusRunning)}). - Update("status", string(entity.TaskStatusCancel)).Error; err != nil { + if err := scheduleConnectorTask(tx, connector.ID, kbID, connectorTaskTypeSync, true); err != nil { + return err + } + + if connectorConfigBool(fullConnector.Config, "sync_deleted_files") { + if err := scheduleConnectorTask(tx, connector.ID, kbID, connectorTaskTypePrune, false); err != nil { return err } } + } - return nil - }) + for connectorID := range oldConnectorIDs { + if _, ok := nextConnectorIDs[connectorID]; ok { + continue + } + if err := tx.Where("kb_id = ? AND connector_id = ?", kbID, connectorID). + Delete(&entity.Connector2Kb{}).Error; err != nil { + return err + } + if err := tx.Model(&entity.SyncLogs{}). + Where("connector_id = ? AND kb_id = ? AND status IN ?", connectorID, kbID, []string{string(entity.TaskStatusSchedule), string(entity.TaskStatusRunning)}). + Update("status", string(entity.TaskStatusCancel)).Error; err != nil { + return err + } + } + + return nil } // GetByID get connector by ID diff --git a/internal/dao/kb.go b/internal/dao/kb.go index a46c76f16b..d9dafcc2c1 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -291,7 +291,7 @@ func (dao *KnowledgebaseDAO) Accessible(datasetID, userID string) bool { var count int64 err = DB.Table("user_tenant"). - Where("tenant_id = ? AND user_id = ?", kb.TenantID, userID). + Where("tenant_id = ? AND user_id = ? AND status = ?", kb.TenantID, userID, "1"). Count(&count).Error if err != nil { diff --git a/internal/engine/elasticsearch/chunk.go b/internal/engine/elasticsearch/chunk.go index dd87bee73f..74ab265d55 100644 --- a/internal/engine/elasticsearch/chunk.go +++ b/internal/engine/elasticsearch/chunk.go @@ -699,7 +699,7 @@ func (e *elasticsearchEngine) updateChunksByQuery(ctx context.Context, indexName sanitized := sanitizeString(val) params[fmt.Sprintf("pp_%s", k)] = sanitized scripts = append(scripts, fmt.Sprintf("ctx._source.%s=params.pp_%s;", k, k)) - case int, float64: + case int, int8, int16, int32, int64, float32, float64: scripts = append(scripts, fmt.Sprintf("ctx._source.%s=%v;", k, val)) case []interface{}: params[fmt.Sprintf("pp_%s", k)] = val @@ -708,6 +708,9 @@ func (e *elasticsearchEngine) updateChunksByQuery(ctx context.Context, indexName } scriptSource := strings.Join(scripts, "") + if scriptSource == "" { + return fmt.Errorf("no supported update fields for update by query") + } // Build update by query body updateBody := map[string]interface{}{ diff --git a/internal/service/dataset/metadata.go b/internal/service/dataset/metadata.go index c5744c58a6..f3fef6000f 100644 --- a/internal/service/dataset/metadata.go +++ b/internal/service/dataset/metadata.go @@ -13,11 +13,10 @@ import ( // UpdateDocumentMetadataConfig updates the metadata config for a document in a dataset. func (d *DatasetService) UpdateDocumentMetadataConfig(userID, datasetID, documentID string, req map[string]interface{}) (*entity.Document, common.ErrorCode, error) { - if _, err := d.kbDAO.GetByIDAndTenantID(datasetID, userID); err != nil { - if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("You don't own the dataset.") - } - return nil, common.CodeServerError, errors.New("Database operation failed") + userID = strings.TrimSpace(userID) + datasetID = strings.TrimSpace(datasetID) + if !d.Accessible(datasetID, userID) { + return nil, common.CodeDataError, errors.New("You don't own the dataset.") } doc, err := d.documentDAO.GetByDocumentIDAndDatasetID(documentID, datasetID) @@ -52,10 +51,16 @@ func (d *DatasetService) UpdateDocumentMetadataConfig(userID, datasetID, documen // GetMetadataConfig gets the auto-metadata configuration for a dataset. func (d *DatasetService) GetMetadataConfig(datasetID, tenantID string) (map[string]interface{}, common.ErrorCode, error) { - kb, err := d.kbDAO.GetByIDAndTenantID(datasetID, tenantID) + datasetID = strings.TrimSpace(datasetID) + tenantID = strings.TrimSpace(tenantID) + if !d.Accessible(datasetID, tenantID) { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + } + + kb, err := d.kbDAO.GetByID(datasetID) if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + return nil, common.CodeDataError, errors.New("Dataset not found") } return nil, common.CodeServerError, errors.New("Database operation failed") } @@ -74,10 +79,14 @@ func (d *DatasetService) UpdateMetadataConfig(datasetID, tenantID string, req *s datasetID = strings.TrimSpace(datasetID) tenantID = strings.TrimSpace(tenantID) - kb, err := d.kbDAO.GetByIDAndTenantID(datasetID, tenantID) + if !d.Accessible(datasetID, tenantID) { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + } + + kb, err := d.kbDAO.GetByID(datasetID) if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + return nil, common.CodeDataError, errors.New("Dataset not found") } return nil, common.CodeServerError, errors.New("Database operation failed") } diff --git a/internal/service/dataset/metadata_config_test.go b/internal/service/dataset/metadata_config_test.go index 129471d079..fbb6d1b35c 100644 --- a/internal/service/dataset/metadata_config_test.go +++ b/internal/service/dataset/metadata_config_test.go @@ -50,6 +50,20 @@ func insertDatasetMetadataConfigKB(t *testing.T, datasetID, tenantID string) { } } +func insertDatasetMetadataConfigTeamMember(t *testing.T, userID, tenantID string) { + t.Helper() + if err := dao.DB.Create(&entity.UserTenant{ + ID: userID + "-" + tenantID, + UserID: userID, + TenantID: tenantID, + Role: "normal", + InvitedBy: tenantID, + Status: sptr("1"), + }).Error; err != nil { + t.Fatalf("insert user tenant: %v", err) + } +} + func insertDatasetMetadataConfigDoc(t *testing.T, docID, datasetID string, parserConfig entity.JSONMap) { t.Helper() doc := &entity.Document{ @@ -156,3 +170,32 @@ func TestDatasetServiceUpdateDocumentMetadataConfigRejectsNonOwner(t *testing.T) t.Fatalf("unexpected error: %v", err) } } + +func TestDatasetServiceUpdateDocumentMetadataConfigAllowsTeamMember(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertDatasetMetadataConfigKB(t, "kb-1", "owner-1") + if err := dao.DB.Model(&entity.Knowledgebase{}). + Where("id = ?", "kb-1"). + Update("permission", string(entity.TenantPermissionTeam)).Error; err != nil { + t.Fatalf("update kb permission: %v", err) + } + insertDatasetMetadataConfigTeamMember(t, "user-1", "owner-1") + insertDatasetMetadataConfigDoc(t, "doc-1", "kb-1", entity.JSONMap{}) + + doc, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig( + "user-1", + "kb-1", + "doc-1", + map[string]interface{}{"metadata": map[string]interface{}{"author": "Alice"}}, + ) + if err != nil { + t.Fatalf("UpdateDocumentMetadataConfig failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + if doc.ParserConfig["metadata"] == nil { + t.Fatalf("metadata was not updated: %#v", doc.ParserConfig) + } +} diff --git a/internal/service/dataset/setup_test.go b/internal/service/dataset/setup_test.go index bcbadc1199..cded99d7e6 100644 --- a/internal/service/dataset/setup_test.go +++ b/internal/service/dataset/setup_test.go @@ -1,6 +1,7 @@ package dataset import ( + "net/url" "testing" "ragflow/internal/dao" @@ -13,12 +14,19 @@ import ( // setupServiceTestDB initializes an in-memory SQLite database for tests. func setupServiceTestDB(t *testing.T) *gorm.DB { t.Helper() - db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ + 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) } + sqlDB, err := db.DB() + if err != nil { + t.Fatalf("failed to access sqlite db: %v", err) + } + t.Cleanup(func() { + _ = sqlDB.Close() + }) if err = db.AutoMigrate( &entity.Document{}, &entity.Knowledgebase{}, @@ -37,6 +45,8 @@ func setupServiceTestDB(t *testing.T) *gorm.DB { &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}, + &entity.TenantModelGroup{}, + &entity.TenantModelGroupMapping{}, &entity.UserCanvas{}, ); err != nil { t.Fatalf("failed to migrate: %v", err) diff --git a/internal/service/dataset/update.go b/internal/service/dataset/update.go index 3c923e148e..96c60d1fa6 100644 --- a/internal/service/dataset/update.go +++ b/internal/service/dataset/update.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "strings" + "time" "ragflow/internal/common" "ragflow/internal/dao" @@ -13,28 +14,49 @@ import ( "ragflow/internal/service" "go.uber.org/zap" + "gorm.io/gorm" + "gorm.io/gorm/clause" ) +const datasetPagerankUpdateTimeout = 30 * time.Second + +type datasetPagerankUpdate struct { + value int64 + index string + datasetID string +} + func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) { - kb, err := d.kbDAO.GetByID(datasetID) - if err != nil { + datasetID = strings.TrimSpace(datasetID) + tenantID = strings.TrimSpace(tenantID) + if _, err := d.kbDAO.GetByID(datasetID); err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("Dataset not found") } return nil, common.CodeServerError, errors.New("Database operation failed") } - if kb == nil || kb.TenantID != tenantID { - return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) - } - connectorsProvided := req.Connectors != nil connectors := make([]service.DatasetConnectorRequest, 0) if req.Connectors != nil { connectors = *req.Connectors } - updates := make(map[string]interface{}) + connectorLinks := make([]dao.DatasetConnectorLink, 0, len(connectors)) + if connectorsProvided { + for _, connector := range connectors { + connectorID := strings.TrimSpace(connector.ID) + if connectorID == "" { + return nil, common.CodeDataError, errors.New("connector id is required") + } + connectorLinks = append(connectorLinks, dao.DatasetConnectorLink{ + ID: connectorID, + AutoParse: connector.AutoParse, + }) + } + } + + simpleUpdates := make(map[string]interface{}) if req.Name != nil { name := strings.TrimSpace(*req.Name) @@ -44,7 +66,7 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U if len(name) > 128 { return nil, common.CodeDataError, errors.New("String should have at most 128 characters") } - updates["name"] = name + simpleUpdates["name"] = name } if req.Avatar != nil { if len(*req.Avatar) > 65535 { @@ -53,27 +75,27 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U if err := validateDatasetAvatar(*req.Avatar); err != nil { return nil, common.CodeDataError, err } - updates["avatar"] = *req.Avatar + simpleUpdates["avatar"] = *req.Avatar } if req.Description != nil { if len(*req.Description) > 65535 { return nil, common.CodeDataError, errors.New("String should have at most 65535 characters") } - updates["description"] = *req.Description + simpleUpdates["description"] = *req.Description } if req.Language != nil { language := strings.TrimSpace(*req.Language) if len(language) > 32 { return nil, common.CodeDataError, errors.New("String should have at most 32 characters") } - updates["language"] = language + simpleUpdates["language"] = language } if req.Permission != nil { permission := strings.TrimSpace(*req.Permission) if permission != "me" && permission != "team" { return nil, common.CodeDataError, errors.New("Input should be 'me' or 'team'") } - updates["permission"] = permission + simpleUpdates["permission"] = permission } isPipelineMode := req.ParseType != nil && *req.ParseType == 2 @@ -86,23 +108,19 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U req.ParserID = nil } + var pipelineID *string if req.PipelineID != nil { - pipelineID, err := normalizeDatasetPipelineID(*req.PipelineID) + normalizedPipelineID, err := normalizeDatasetPipelineID(*req.PipelineID) if err != nil { return nil, common.CodeDataError, err } - if pipelineID != nil { - updates["pipeline_id"] = *pipelineID - } + pipelineID = normalizedPipelineID } parserID, parserIDProvided, err := datasetUpdateParserID(req) if err != nil { return nil, common.CodeDataError, err } - if parserIDProvided { - updates["parser_id"] = parserID - } if req.ParseType == nil && parserIDProvided && req.PipelineID != nil { return nil, common.CodeDataError, errors.New("parser_id and pipeline_id are mutually exclusive") @@ -112,45 +130,108 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U if err != nil { return nil, common.CodeDataError, err } - if embdIDProvided { - tenantEmbdID := ptrStringValue(kb.TenantEmbdID) - if embdID == "" { - embdID = kb.EmbdID - } else { - tenantEmbdID = "" - } - ok, message := d.verifyEmbeddingAvailability(embdID, tenantID) - if !ok { - return nil, common.CodeDataError, errors.New(message) - } - if embdID != "" && tenantEmbdID == "" { - resolvedID, err := service.NewModelProviderService().ResolveModelID(tenantID, entity.ModelTypeEmbedding, embdID) - if err == nil { - tenantEmbdID = resolvedID - } - } - updates["embd_id"] = embdID - updates["tenant_embd_id"] = stringPtrIfNotEmpty(tenantEmbdID) - } if req.ParserConfig != nil { if err := validateDatasetParserConfigSize(req.ParserConfig); err != nil { return nil, common.CodeDataError, err } - if len(req.ParserConfig) > 0 { - effectiveParserID := kb.ParserID + } + + var requestedPagerank int64 + pagerankRequested := req.Pagerank != nil + if pagerankRequested { + requestedPagerank = *req.Pagerank + if *req.Pagerank < 0 || *req.Pagerank > 100 { + return nil, common.CodeDataError, errors.New("Input should be less than or equal to 100") + } + if d.docEngine == nil { + return nil, common.CodeServerError, errors.New("Document engine is not initialized") + } + if !d.docEngine.SupportsPageRank() { + return nil, common.CodeDataError, errors.New("'pagerank' can only be set when doc_engine is elasticsearch") + } + } + + requestedAnyUpdate := len(simpleUpdates) > 0 || connectorsProvided || parserIDProvided || + pipelineID != nil || embdIDProvided || pagerankRequested || req.ParserConfig != nil || req.ParserConfigProvided + if !requestedAnyUpdate { + return nil, common.CodeDataError, errors.New("No properties were modified") + } + + txCode := common.CodeSuccess + var updatedKB *entity.Knowledgebase + var linkedConnectors []*dao.ConnectorDatasetListItem + var pagerankUpdate *datasetPagerankUpdate + err = dao.DB.Transaction(func(tx *gorm.DB) error { + lockedKB, code, authErr := d.lockAccessibleDatasetForUpdate(tx, datasetID, tenantID) + if authErr != nil { + txCode = code + return authErr + } + + if req.Permission != nil && lockedKB.TenantID != tenantID { + txCode = common.CodeDataError + return errors.New("Only dataset owner can change permission") + } + + updates := make(map[string]interface{}, len(simpleUpdates)+6) + for key, value := range simpleUpdates { + updates[key] = value + } + + if nameValue, ok := updates["name"].(string); ok && strings.ToLower(nameValue) != strings.ToLower(lockedKB.Name) { + var existing entity.Knowledgebase + lookupErr := tx.Where("LOWER(name) = LOWER(?) AND tenant_id = ? AND status = ?", nameValue, tenantID, string(entity.StatusValid)).First(&existing).Error + if lookupErr != nil && !dao.IsNotFoundErr(lookupErr) { + txCode = common.CodeServerError + return errors.New("Database operation failed") + } + if lookupErr == nil { + txCode = common.CodeDataError + return fmt.Errorf("Dataset name '%s' already exists", nameValue) + } + } + + if pipelineID != nil { + updates["pipeline_id"] = *pipelineID + } + if parserIDProvided { + updates["parser_id"] = parserID + } + if embdIDProvided { + effectiveEmbdID := embdID + tenantEmbdID := ptrStringValue(lockedKB.TenantEmbdID) + if effectiveEmbdID == "" { + effectiveEmbdID = lockedKB.EmbdID + } else { + tenantEmbdID = "" + } + ok, message := d.verifyEmbeddingAvailability(effectiveEmbdID, tenantID) + if !ok { + txCode = common.CodeDataError + return errors.New(message) + } + if effectiveEmbdID != "" && tenantEmbdID == "" { + resolvedID, err := service.NewModelProviderService().ResolveModelID(tenantID, entity.ModelTypeEmbedding, effectiveEmbdID) + if err == nil { + tenantEmbdID = resolvedID + } + } + updates["embd_id"] = effectiveEmbdID + updates["tenant_embd_id"] = stringPtrIfNotEmpty(tenantEmbdID) + } + + if req.ParserConfig != nil && len(req.ParserConfig) > 0 { + effectiveParserID := lockedKB.ParserID if parserIDProvided { effectiveParserID = parserID } - effectivePipelineID := kb.PipelineID - if req.PipelineID != nil { - if normalized, err := normalizeDatasetPipelineID(*req.PipelineID); err == nil { - effectivePipelineID = normalized - } - } else if parserIDProvided && kb.PipelineID != nil { + effectivePipelineID := lockedKB.PipelineID + if pipelineID != nil { + effectivePipelineID = pipelineID + } else if parserIDProvided && lockedKB.PipelineID != nil { effectivePipelineID = nil } - isCanvas := effectivePipelineID != nil && strings.TrimSpace(*effectivePipelineID) != "" dslJSON, dslErr := service.LoadPipelineDSL(isCanvas, effectiveParserID, effectivePipelineID) if dslErr != nil { @@ -161,115 +242,141 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U updates["parser_config"] = pipelinepkg.BuildParserConfig(dslJSON, map[string]interface{}(req.ParserConfig)) } } + if pagerankRequested { + pagerankUpdate = &datasetPagerankUpdate{ + value: requestedPagerank, + index: fmt.Sprintf("ragflow_%s", lockedKB.TenantID), + datasetID: lockedKB.ID, + } + if requestedPagerank != lockedKB.Pagerank { + updates["pagerank"] = requestedPagerank + } + } + if parserIDProvided && parserID != lockedKB.ParserID { + if _, ok := updates["parser_config"]; !ok { + if resolved, cpErr := service.ResolveComponentParamsDefaults(parserID, nil); cpErr != nil { + common.Warn("failed to resolve component params defaults on parser_id switch", + zap.String("parserID", parserID), zap.Error(cpErr)) + } else if resolved != nil { + updates["parser_config"] = resolved + } + } + } + if lockedKB.PipelineID != nil && parserIDProvided { + if _, ok := updates["pipeline_id"]; !ok { + updates["pipeline_id"] = nil + } + } + + pipelineChanged := pipelineID != nil && (lockedKB.PipelineID == nil || *pipelineID != *lockedKB.PipelineID) + if pipelineChanged { + cfgParserID := lockedKB.ParserID + if parserIDProvided { + cfgParserID = parserID + } + if cpDefaults, cpErr := service.ResolveComponentParamsDefaults(cfgParserID, pipelineID); cpErr != nil { + common.Warn("failed to resolve component params defaults on pipeline change", + zap.String("parserID", cfgParserID), zap.Error(cpErr)) + } else if cpDefaults != nil { + updates["parser_config"] = cpDefaults + } + } + if len(updates) > 0 { + if err = tx.Model(&entity.Knowledgebase{}).Where("id = ?", lockedKB.ID).Updates(updates).Error; err != nil { + if dao.IsDuplicateKeyErr(err) { + if nameValue, ok := updates["name"].(string); ok { + txCode = common.CodeDataError + return fmt.Errorf("Dataset name '%s' already exists", nameValue) + } + txCode = common.CodeDataError + return errors.New("Dataset name already exists") + } + txCode = common.CodeServerError + return errors.New("Update dataset error.(Database error)") + } + } + + if connectorsProvided { + if err = d.connectorDAO.LinkDatasetConnectorsTx(tx, lockedKB.ID, lockedKB.TenantID, connectorLinks); err != nil { + if dao.IsConnectorNotAccessibleErr(err) { + txCode = common.CodeDataError + return err + } + txCode = common.CodeServerError + return errors.New("Database operation failed") + } + } + + updatedKB = &entity.Knowledgebase{} + if err = tx.Where("id = ? AND status = ?", lockedKB.ID, string(entity.StatusValid)).First(updatedKB).Error; err != nil { + txCode = common.CodeDataError + return errors.New("Dataset updated failed") + } + + linkedConnectors, err = d.connectorDAO.ListByDatasetIDTx(tx, lockedKB.ID) + if err != nil { + txCode = common.CodeServerError + return errors.New("Database operation failed") + } + + return nil + }) + if err != nil { + if txCode == common.CodeSuccess { + txCode = common.CodeServerError + } + return nil, txCode, err } - if req.Pagerank != nil && *req.Pagerank != kb.Pagerank { - if *req.Pagerank < 0 || *req.Pagerank > 100 { - return nil, common.CodeDataError, errors.New("Input should be less than or equal to 100") - } - if !d.docEngine.SupportsPageRank() { - return nil, common.CodeDataError, errors.New("'pagerank' can only be set when doc_engine is elasticsearch") - } - indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) - if *req.Pagerank > 0 { - err = d.docEngine.UpdateChunks(context.Background(), map[string]interface{}{"kb_id": kb.ID}, map[string]interface{}{common.PAGERANK_FLD: *req.Pagerank}, indexName, kb.ID) - } else { - err = d.docEngine.UpdateChunks(context.Background(), map[string]interface{}{"exists": common.PAGERANK_FLD}, map[string]interface{}{"remove": common.PAGERANK_FLD}, indexName, kb.ID) - } - if err != nil { + if pagerankUpdate != nil { + if err = d.updateDatasetPagerankChunks(*pagerankUpdate); err != nil { return nil, common.CodeServerError, err } - updates["pagerank"] = *req.Pagerank - } - - if parserIDProvided && parserID != kb.ParserID { - if _, ok := updates["parser_config"]; !ok { - if resolved, cpErr := service.ResolveComponentParamsDefaults(parserID, nil); cpErr != nil { - common.Warn("failed to resolve component params defaults on parser_id switch", - zap.String("parserID", parserID), zap.Error(cpErr)) - } else if resolved != nil { - updates["parser_config"] = resolved - } - } - } - if kb.PipelineID != nil && parserIDProvided { - if _, ok := updates["pipeline_id"]; !ok { - updates["pipeline_id"] = nil - } - } - - pipelineChanged := req.PipelineID != nil && (kb.PipelineID == nil || *req.PipelineID != *kb.PipelineID) - if pipelineChanged { - cfgParserID := kb.ParserID - if parserIDProvided { - cfgParserID = parserID - } - cfgPipelineID, _ := updates["pipeline_id"].(string) - var cpPipelineID *string - if cfgPipelineID != "" { - cpPipelineID = &cfgPipelineID - } - if cpDefaults, cpErr := service.ResolveComponentParamsDefaults(cfgParserID, cpPipelineID); cpErr != nil { - common.Warn("failed to resolve component params defaults on pipeline change", - zap.String("parserID", cfgParserID), zap.Error(cpErr)) - } else if cpDefaults != nil { - updates["parser_config"] = cpDefaults - } - } - - if nameValue, ok := updates["name"].(string); ok && strings.ToLower(nameValue) != strings.ToLower(kb.Name) { - existing, lookupErr := d.kbDAO.GetByName(nameValue, tenantID) - if lookupErr != nil && !dao.IsNotFoundErr(lookupErr) { - return nil, common.CodeServerError, errors.New("Database operation failed") - } - if existing != nil { - return nil, common.CodeDataError, fmt.Errorf("Dataset name '%s' already exists", nameValue) - } - } - - if len(updates) == 0 && !connectorsProvided && !req.ParserConfigProvided { - return nil, common.CodeDataError, errors.New("No properties were modified") - } - - if len(updates) > 0 { - if err = d.kbDAO.UpdateByID(kb.ID, updates); err != nil { - if dao.IsDuplicateKeyErr(err) { - if nameValue, ok := updates["name"].(string); ok { - return nil, common.CodeDataError, fmt.Errorf("Dataset name '%s' already exists", nameValue) - } - return nil, common.CodeDataError, errors.New("Dataset name already exists") - } - return nil, common.CodeServerError, errors.New("Update dataset error.(Database error)") - } - } - - if connectorsProvided { - connectorLinks := make([]dao.DatasetConnectorLink, 0, len(connectors)) - for _, connector := range connectors { - connectorID := strings.TrimSpace(connector.ID) - if connectorID == "" { - return nil, common.CodeDataError, errors.New("connector id is required") - } - connectorLinks = append(connectorLinks, dao.DatasetConnectorLink{ - ID: connectorID, - AutoParse: connector.AutoParse, - }) - } - if err = d.connectorDAO.LinkDatasetConnectors(kb.ID, connectorLinks); err != nil { - return nil, common.CodeServerError, errors.New("Database operation failed") - } - } - - updatedKB, err := d.kbDAO.GetByID(kb.ID) - if err != nil { - return nil, common.CodeDataError, errors.New("Dataset updated failed") } data := datasetToMap(updatedKB) - linkedConnectors, err := d.connectorDAO.ListByDatasetID(kb.ID) - if err != nil { - return nil, common.CodeServerError, errors.New("Database operation failed") - } data["connectors"] = datasetConnectorsOrEmpty(linkedConnectors) return data, common.CodeSuccess, nil } + +func (d *DatasetService) updateDatasetPagerankChunks(update datasetPagerankUpdate) error { + ctx, cancel := context.WithTimeout(context.Background(), datasetPagerankUpdateTimeout) + defer cancel() + if update.value > 0 { + return d.docEngine.UpdateChunks(ctx, map[string]interface{}{"kb_id": update.datasetID}, map[string]interface{}{common.PAGERANK_FLD: update.value}, update.index, update.datasetID) + } + return d.docEngine.UpdateChunks(ctx, map[string]interface{}{"exists": common.PAGERANK_FLD}, map[string]interface{}{"remove": common.PAGERANK_FLD}, update.index, update.datasetID) +} + +func (d *DatasetService) lockAccessibleDatasetForUpdate(tx *gorm.DB, datasetID, userID string) (*entity.Knowledgebase, common.ErrorCode, error) { + var kb entity.Knowledgebase + err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("id = ? AND status = ?", datasetID, string(entity.StatusValid)). + First(&kb).Error + if err != nil { + if dao.IsNotFoundErr(err) { + return nil, common.CodeDataError, errors.New("Dataset not found") + } + return nil, common.CodeServerError, errors.New("Database operation failed") + } + + if kb.TenantID == userID { + return &kb, common.CodeSuccess, nil + } + if kb.Permission != string(entity.TenantPermissionTeam) { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID) + } + + var relation entity.UserTenant + err = tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("tenant_id = ? AND user_id = ? AND status = ?", kb.TenantID, userID, "1"). + First(&relation).Error + if err != nil { + if dao.IsNotFoundErr(err) { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID) + } + return nil, common.CodeServerError, errors.New("Database operation failed") + } + + return &kb, common.CodeSuccess, nil +} diff --git a/internal/service/dataset/update_test.go b/internal/service/dataset/update_test.go index 3d98f86eed..687d79c3a2 100644 --- a/internal/service/dataset/update_test.go +++ b/internal/service/dataset/update_test.go @@ -237,6 +237,40 @@ func TestDatasetServiceUpdateDatasetRejectsNonOwner(t *testing.T) { } } +func TestDatasetServiceUpdateDatasetRejectsTeamMemberPermissionChange(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "owner-1", "Original") + if err := dao.DB.Model(&entity.Knowledgebase{}). + Where("id = ?", "kb-1"). + Update("permission", string(entity.TenantPermissionTeam)).Error; err != nil { + t.Fatalf("update kb permission: %v", err) + } + insertDatasetUpdateTeamMember(t, "user-1", "owner-1") + + permission := string(entity.TenantPermissionMe) + _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "user-1", service.UpdateDatasetRequest{ + Permission: &permission, + }) + if err == nil { + t.Fatal("expected permission change error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if err.Error() != "Only dataset owner can change permission" { + t.Fatalf("unexpected error: %v", err) + } + + persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + if err != nil { + t.Fatalf("get dataset: %v", err) + } + if persisted.Permission != string(entity.TenantPermissionTeam) { + t.Fatalf("expected permission unchanged, got %q", persisted.Permission) + } +} + func TestDatasetServiceUpdateDatasetValidatesName(t *testing.T) { db := setupDatasetUpdateTestDB(t) pushServiceDB(t, db) @@ -327,6 +361,37 @@ func TestDatasetServiceUpdateDatasetLinksConnectors(t *testing.T) { } } +func TestDatasetServiceUpdateDatasetRejectsCrossTenantConnector(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + insertDatasetUpdateConnector(t, "connector-1", "tenant-2") + + autoParse := "0" + _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + Connectors: &[]service.DatasetConnectorRequest{{ID: "connector-1", AutoParse: autoParse}}, + }) + if err == nil { + t.Fatal("expected connector tenant mismatch error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if !dao.IsConnectorNotAccessibleErr(err) { + t.Fatalf("unexpected error: %v", err) + } + + var count int64 + if err := dao.DB.Model(&entity.Connector2Kb{}). + Where("kb_id = ? AND connector_id = ?", "kb-1", "connector-1"). + Count(&count).Error; err != nil { + t.Fatalf("count connector link: %v", err) + } + if count != 0 { + t.Fatalf("expected no connector link, got %d", count) + } +} + func TestDatasetServiceUpdateDatasetAcceptsProviderInstanceEmbedding(t *testing.T) { db := setupDatasetUpdateTestDB(t) pushServiceDB(t, db) @@ -589,6 +654,13 @@ func setupDatasetUpdateTestDB(t *testing.T) *gorm.DB { t.Helper() db := setupServiceTestDB(t) + migrateDatasetUpdateTestTables(t, db) + return db +} + +func migrateDatasetUpdateTestTables(t *testing.T, db *gorm.DB) { + t.Helper() + if err := db.AutoMigrate( &entity.Connector{}, &entity.Connector2Kb{}, @@ -596,11 +668,12 @@ func setupDatasetUpdateTestDB(t *testing.T) *gorm.DB { &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}, + &entity.TenantModelGroup{}, + &entity.TenantModelGroupMapping{}, &entity.UserCanvas{}, ); err != nil { t.Fatalf("failed to migrate dataset update tables: %v", err) } - return db } func testDatasetUpdateService(t *testing.T) *DatasetService { @@ -644,6 +717,20 @@ func insertDatasetUpdateCanvas(t *testing.T, id, userID string) { } } +func insertDatasetUpdateTeamMember(t *testing.T, userID, tenantID string) { + t.Helper() + if err := dao.DB.Create(&entity.UserTenant{ + ID: userID + "-" + tenantID, + UserID: userID, + TenantID: tenantID, + Role: "normal", + InvitedBy: tenantID, + Status: sptr("1"), + }).Error; err != nil { + t.Fatalf("insert user tenant: %v", err) + } +} + func insertDatasetUpdateConnector(t *testing.T, id, tenantID string) { t.Helper() @@ -707,13 +794,9 @@ func insertDatasetUpdateTenantModel(t *testing.T, id, providerID, instanceID, mo } } -// seedDatasetUpdateCanvas migrates user_canvas on the active test DB and -// inserts a canvas row with the given DSL. +// seedDatasetUpdateCanvas inserts a canvas row with the given DSL. func seedDatasetUpdateCanvas(t *testing.T, id, userID string, dslJSON []byte) { t.Helper() - if err := dao.DB.AutoMigrate(&entity.UserCanvas{}); err != nil { - t.Fatalf("migrate user_canvas: %v", err) - } var dslMap map[string]any if err := json.Unmarshal(dslJSON, &dslMap); err != nil { t.Fatalf("unmarshal seed dsl: %v", err)