diff --git a/internal/admin/handler.go b/internal/admin/handler.go index 6b7cd90448..3977834b5c 100644 --- a/internal/admin/handler.go +++ b/internal/admin/handler.go @@ -25,6 +25,7 @@ import ( "ragflow/internal/common" "ragflow/internal/engine" "ragflow/internal/engine/redis" + "ragflow/internal/httputil" "ragflow/internal/server" "ragflow/internal/service" "ragflow/internal/storage" @@ -1053,7 +1054,7 @@ func (h *Handler) RemoveIngestionTasks(c *gin.Context) { if req.Email == nil && req.Status == nil { tasks, err := h.service.RemoveIngestionTasks(req.Tasks) if err != nil { - common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) + common.ErrorWithCode(c, httputil.IngestionTaskErrorCode(err), err.Error()) return } @@ -1061,7 +1062,7 @@ func (h *Handler) RemoveIngestionTasks(c *gin.Context) { } else { tasks, err := h.service.RemoveIngestionTasksByCondition(req.Tasks, req.Email, req.Status) if err != nil { - common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) + common.ErrorWithCode(c, httputil.IngestionTaskErrorCode(err), err.Error()) return } common.SuccessWithData(c, tasks, "Remove tasks successfully") @@ -1084,7 +1085,7 @@ func (h *Handler) StopIngestionTasks(c *gin.Context) { if req.Email == nil && req.Status == nil { tasks, err := h.service.StopIngestionTasks(req.Tasks) if err != nil { - common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) + common.ErrorWithCode(c, httputil.IngestionTaskErrorCode(err), err.Error()) return } var result []map[string]string diff --git a/internal/admin/service.go b/internal/admin/service.go index 71052f63dd..ce7f3194ff 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -32,6 +32,7 @@ import ( "ragflow/internal/entity" modelModule "ragflow/internal/entity/models" "ragflow/internal/server" + servicepkg "ragflow/internal/service" "ragflow/internal/utility" "regexp" "strconv" @@ -61,6 +62,7 @@ type Service struct { llmDAO *dao.LLMDAO ingestionTaskDAO *dao.IngestionTaskDAO ingestionTaskLogDao *dao.IngestionTaskLogDAO + ingestionTaskSvc *servicepkg.IngestionTaskService } // NewService create admin service @@ -85,6 +87,7 @@ func NewService() *Service { llmDAO: dao.NewLLMDAO(), ingestionTaskDAO: dao.NewIngestionTaskDAO(), ingestionTaskLogDao: dao.NewIngestionTaskLogDAO(), + ingestionTaskSvc: servicepkg.NewIngestionTaskService(), } } @@ -100,92 +103,15 @@ func (s *Service) Logout(user interface{}) error { // ListTasks func (s *Service) ListIngestionTasks() ([]map[string]interface{}, error) { - - ingestionTasks, err := s.ingestionTaskDAO.GetAllTasks(0, 0) - if err != nil { - return nil, err - } - - showTasks := []map[string]interface{}{} - for _, task := range ingestionTasks { - var user *entity.User - user, err = s.userDAO.GetByTenantID(task.UserID) - if err != nil { - return nil, err - } - //var document *entity.Document - //document, err = s.documentDAO.GetByID(task.DocumentID) - //if err != nil { - // return nil, err - //} - - var showTask map[string]interface{} - var latestLog *entity.IngestionTaskLog - latestLog, err = s.ingestionTaskLogDao.LatestLogByTaskID(task.ID) - showTask = map[string]interface{}{ - "id": task.ID, - "user_id": task.UserID, - "user": user.Email, - "document_id": task.DocumentID, - "status": task.Status, - } - if err == nil && latestLog != nil && latestLog.Checkpoint != nil { - step, ok := latestLog.Checkpoint["current_step"].(float64) - if !ok { - showTasks = append(showTasks, showTask) - continue - } - showTask = map[string]interface{}{ - "id": task.ID, - "user_id": task.UserID, - "user": user.Email, - "document_id": task.DocumentID, - "status": task.Status, - "step": int(step), - } - } - - showTasks = append(showTasks, showTask) - } - return showTasks, nil + return s.ingestionTaskSvc.ListAllForAdmin() } func (s *Service) RemoveIngestionTasks(tasks []string) ([]map[string]string, error) { - var deletedTasks []map[string]string - for _, taskID := range tasks { - taskRecord := map[string]string{ - "task_id": taskID, - } - _, err := s.ingestionTaskDAO.Delete(taskID, nil) - if err != nil { - taskRecord["remove"] = fmt.Sprintf("fail: %s", err.Error()) - } else { - taskRecord["remove"] = "success" - } - deletedTasks = append(deletedTasks, taskRecord) - } - return deletedTasks, nil + return s.ingestionTaskSvc.RemoveMany(tasks, nil) } func (s *Service) StopIngestionTasks(tasks []string) ([]*entity.IngestionTask, error) { - var taskResponses []*entity.IngestionTask - for _, taskID := range tasks { - task, err := s.ingestionTaskDAO.SetStoppingByAPIServer(taskID) - if err != nil { - return nil, err - } - - if task.Status == common.STOPPING { - msgQueueEngine := engine.GetMessageQueueEngine() - err = msgQueueEngine.PublishTask("tasks.RAGFLOW", []byte(task.ID)) - if err != nil { - return nil, err - } - } - - taskResponses = append(taskResponses, task) - } - return taskResponses, nil + return s.ingestionTaskSvc.RequestStopMany(tasks, nil) } // GetUserByToken get user by access token diff --git a/internal/dao/ingestion.go b/internal/dao/ingestion.go index e8b5f632f3..a41d741ad1 100644 --- a/internal/dao/ingestion.go +++ b/internal/dao/ingestion.go @@ -22,6 +22,8 @@ import ( "ragflow/internal/common" "ragflow/internal/entity" "ragflow/internal/utility" + + "gorm.io/gorm" ) type IngestionTaskDAO struct{} @@ -30,64 +32,40 @@ func NewIngestionTaskDAO() *IngestionTaskDAO { return &IngestionTaskDAO{} } -// Use by api server to create task -// created → running : After the ingestor component assigns the task, it changes the status to running -// running → completed : Task executes successfully -// running → failed : Error occurs during execution -// created → canceling : User cancels before the task is picked up by the ingestor -// running → canceling : User cancels during execution -// completed → canceling : User cancels a completed task (e.g., for cleanup/rollback) -// canceling → canceled : Cancellation completes -// failed → created : Retry (back to start) -// canceled → created : Retry/re-execute (back to start) -func (dao *IngestionTaskDAO) CheckAndCreate(ingestionTask *entity.IngestionTask) (*entity.IngestionTask, error) { - tx := DB.Begin() - if tx.Error != nil { - return nil, tx.Error - } - - defer func() { - if r := recover(); r != nil { - tx.Rollback() - panic(r) - } - }() - - // Check if the task is created - var taskRecord *entity.IngestionTask - err := tx.Where("document_id = ?", ingestionTask.DocumentID).First(&taskRecord).Error - if err == nil { - // found - if taskRecord.Status == common.FAILED || taskRecord.Status == common.STOPPED { - // restart the task - err = tx.Model(&entity.IngestionTask{}).Where("id = ?", taskRecord.ID).Update("status", common.CREATED).Error - if err != nil { - tx.Rollback() - return nil, err - } - } else { - return nil, fmt.Errorf("document id %s already exists, status: %s, task id: %s", ingestionTask.DocumentID, taskRecord.Status, taskRecord.ID) - } - } else { - // create ingestion task - ingestionTask.ID = utility.GenerateUUID() - if err = tx.Create(ingestionTask).Error; err != nil { - tx.Rollback() - return nil, err - } - taskRecord = ingestionTask - } - - if err = tx.Commit().Error; err != nil { +func (dao *IngestionTaskDAO) Create(ingestionTask *entity.IngestionTask) (*entity.IngestionTask, error) { + existing, err := dao.GetByDocumentID(ingestionTask.DocumentID) + if err != nil { return nil, err } - - return taskRecord, nil + if existing != nil { + return nil, fmt.Errorf("document id %s already exists, status: %s, task id: %s", ingestionTask.DocumentID, existing.Status, existing.ID) + } + if ingestionTask.ID == "" { + ingestionTask.ID = utility.GenerateUUID() + } + if err := DB.Create(ingestionTask).Error; err != nil { + if errors.Is(err, gorm.ErrDuplicatedKey) { + existing, getErr := dao.GetByDocumentID(ingestionTask.DocumentID) + if getErr != nil { + return nil, getErr + } + if existing != nil { + return nil, fmt.Errorf("document id %s already exists, status: %s, task id: %s", ingestionTask.DocumentID, existing.Status, existing.ID) + } + } + return nil, err + } + return ingestionTask, nil } -// UpdateStatus Update ingestion task status -func (dao *IngestionTaskDAO) UpdateStatus(taskID, status string) error { - return DB.Model(&entity.IngestionTask{}).Where("id = ?", taskID).Update("status", status).Error +func (dao *IngestionTaskDAO) UpdateStatusIfCurrent(taskID, fromStatus, toStatus string) (bool, error) { + result := DB.Model(&entity.IngestionTask{}). + Where("id = ? AND status = ?", taskID, fromStatus). + Update("status", toStatus) + if result.Error != nil { + return false, result.Error + } + return result.RowsAffected == 1, nil } // UpdateComponentTotal records the number of components in the task's DSL @@ -96,126 +74,6 @@ func (dao *IngestionTaskDAO) UpdateComponentTotal(taskID string, total int) erro return DB.Model(&entity.IngestionTask{}).Where("id = ?", taskID).Update("component_total", total).Error } -// CheckAnd called by ingestor -// if task status is RUNNING, COMPLETED, STOPPED, FAILED, just return without error -// if task status is CREATE, update to RUNNING -// if task status is STOPPING, update to STOPPED -func (dao *IngestionTaskDAO) SetRunningByIngestor(taskID string) (*entity.IngestionTask, error) { - - tx := DB.Begin() - if tx.Error != nil { - return nil, tx.Error - } - var committed bool - - defer func() { - if committed { - tx.Commit() - } else { - tx.Rollback() - if r := recover(); r != nil { - panic(r) - } - } - }() - - var tasks []*entity.IngestionTask - err := tx.Where("id = ?", taskID).Find(&tasks).Error - if err != nil { - return nil, err - } - - if len(tasks) == 0 { - return nil, common.ErrTaskNotFound - } - - if len(tasks) != 1 { - return nil, fmt.Errorf("task %s has multiple records", taskID) - } - - taskStatus := tasks[0].Status - switch taskStatus { - case common.CREATED: - tasks[0].Status = common.RUNNING - err = tx.Model(&entity.IngestionTask{}).Where("id = ?", taskID).Update("status", common.RUNNING).Error - if err != nil { - return nil, err - } - committed = true - return tasks[0], nil - case common.STOPPING: - tasks[0].Status = common.STOPPED - err = tx.Model(&entity.IngestionTask{}).Where("id = ?", taskID).Update("status", common.STOPPED).Error - if err != nil { - return nil, err - } - committed = true - return tasks[0], nil - case common.RUNNING: - // this task was executing before, just return without error - committed = true - return tasks[0], nil - default: - return tasks[0], nil - } -} - -func (dao *IngestionTaskDAO) SetStoppingByAPIServer(taskID string) (*entity.IngestionTask, error) { - - tx := DB.Begin() - if tx.Error != nil { - return nil, tx.Error - } - var committed bool - - defer func() { - if committed { - tx.Commit() - } else { - tx.Rollback() - if r := recover(); r != nil { - panic(r) - } - } - }() - - var tasks []*entity.IngestionTask - err := tx.Where("id = ?", taskID).Find(&tasks).Error - if err != nil { - return nil, err - } - - if len(tasks) == 0 { - return nil, fmt.Errorf("task %s not found", taskID) - } - - if len(tasks) != 1 { - return nil, fmt.Errorf("task %s has multiple records", taskID) - } - - taskStatus := tasks[0].Status - switch taskStatus { - case common.CREATED: - tasks[0].Status = common.STOPPED - err = tx.Model(&entity.IngestionTask{}).Where("id = ?", taskID).Update("status", common.STOPPED).Error - if err != nil { - return nil, err - } - committed = true - return tasks[0], nil - case common.RUNNING: - tasks[0].Status = common.STOPPING - err = tx.Model(&entity.IngestionTask{}).Where("id = ?", taskID).Update("status", common.STOPPING).Error - if err != nil { - return nil, err - } - committed = true - return tasks[0], nil - default: - return tasks[0], nil - } -} - type TaskInfo struct { TaskID string `json:"task_id"` FilesToDelete []string `json:"files_to_delete"` diff --git a/internal/dao/ingestion_task_test.go b/internal/dao/ingestion_task_test.go new file mode 100644 index 0000000000..8378f2cf71 --- /dev/null +++ b/internal/dao/ingestion_task_test.go @@ -0,0 +1,135 @@ +package dao + +import ( + "errors" + "testing" + + "ragflow/internal/common" + "ragflow/internal/entity" + + "gorm.io/gorm" +) + +func TestIngestionTaskDAOUpdateStatusIfCurrentSucceeds(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + task := &entity.IngestionTask{ + ID: "task-1", + UserID: "user-1", + DocumentID: "doc-1", + DatasetID: "kb-1", + Status: common.CREATED, + } + if err := db.Create(task).Error; err != nil { + t.Fatalf("create task: %v", err) + } + + updated, err := NewIngestionTaskDAO().UpdateStatusIfCurrent("task-1", common.CREATED, common.RUNNING) + if err != nil { + t.Fatalf("UpdateStatusIfCurrent failed: %v", err) + } + if !updated { + t.Fatal("expected update to succeed") + } + + reloaded, err := NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("reload task: %v", err) + } + if reloaded.Status != common.RUNNING { + t.Fatalf("status = %q, want %q", reloaded.Status, common.RUNNING) + } +} + +func TestIngestionTaskDAOCreateRejectsExistingTerminalTask(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + testCases := []struct { + name string + status string + }{ + {name: "failed", status: common.FAILED}, + {name: "stopped", status: common.STOPPED}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + if err := db.Where("id = ?", "task-1").Delete(&entity.IngestionTask{}).Error; err != nil { + t.Fatalf("clear task: %v", err) + } + task := &entity.IngestionTask{ID: "task-1", UserID: "user-1", DocumentID: "doc-1", DatasetID: "kb-1", Status: tc.status} + if err := db.Create(task).Error; err != nil { + t.Fatalf("create task: %v", err) + } + _, err := NewIngestionTaskDAO().Create(&entity.IngestionTask{ID: "task-2", UserID: "user-1", DocumentID: "doc-1", DatasetID: "kb-1", Status: common.CREATED}) + if err == nil { + t.Fatal("expected Create to reject duplicate document task") + } + reloaded, err := NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("reload task: %v", err) + } + if reloaded.Status != tc.status { + t.Fatalf("status = %q, want %q", reloaded.Status, tc.status) + } + }) + } +} + +func TestIngestionTaskDAODocumentIDIsUniqueAtDBLevel(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + first := &entity.IngestionTask{ID: "task-1", UserID: "user-1", DocumentID: "doc-1", DatasetID: "kb-1", Status: common.CREATED} + if err := db.Create(first).Error; err != nil { + t.Fatalf("create first task: %v", err) + } + + second := &entity.IngestionTask{ID: "task-2", UserID: "user-1", DocumentID: "doc-1", DatasetID: "kb-1", Status: common.CREATED} + err := db.Create(second).Error + if !errors.Is(err, gorm.ErrDuplicatedKey) { + t.Fatalf("expected duplicated key error, got %v", err) + } +} + +func TestIngestionTaskDAOUpdateStatusIfCurrentRejectsMismatchedStatus(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + task := &entity.IngestionTask{ + ID: "task-1", + UserID: "user-1", + DocumentID: "doc-1", + DatasetID: "kb-1", + Status: common.STOPPING, + } + if err := db.Create(task).Error; err != nil { + t.Fatalf("create task: %v", err) + } + + updated, err := NewIngestionTaskDAO().UpdateStatusIfCurrent("task-1", common.CREATED, common.RUNNING) + if err != nil { + t.Fatalf("UpdateStatusIfCurrent failed: %v", err) + } + if updated { + t.Fatal("expected update to be rejected") + } + + reloaded, err := NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("reload task: %v", err) + } + if reloaded.Status != common.STOPPING { + t.Fatalf("status = %q, want %q", reloaded.Status, common.STOPPING) + } +} diff --git a/internal/dao/migration.go b/internal/dao/migration.go index 23052f6647..6f4b0c08ad 100644 --- a/internal/dao/migration.go +++ b/internal/dao/migration.go @@ -47,6 +47,11 @@ func RunMigrations(db *gorm.DB) error { return fmt.Errorf("failed to add unique index on user.email: %w", err) } + // Add unique index on ingestion_task.document_id + if err := migrateIngestionTaskDocumentIDUnique(db); err != nil { + return fmt.Errorf("failed to add unique index on ingestion_task.document_id: %w", err) + } + // Modify column types that AutoMigrate may not handle correctly if err := modifyColumnTypes(db); err != nil { return fmt.Errorf("failed to modify column types: %w", err) @@ -77,7 +82,7 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { // Check if 'id' column already exists using raw SQL var idColumnExists int64 err := db.Raw(` - SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = 'tenant_llm' AND COLUMN_NAME = 'id' `).Scan(&idColumnExists).Error if err != nil { @@ -88,9 +93,9 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { // Check if id is already a primary key with auto_increment var count int64 err := db.Raw(` - SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS - WHERE TABLE_NAME = 'tenant_llm' - AND COLUMN_NAME = 'id' + SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + WHERE TABLE_NAME = 'tenant_llm' + AND COLUMN_NAME = 'id' AND EXTRA LIKE '%auto_increment%' `).Scan(&count).Error if err != nil { @@ -108,7 +113,7 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { return db.Transaction(func(tx *gorm.DB) error { // Check for temp_id column and drop it if exists var tempIdExists int64 - tx.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + tx.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = 'tenant_llm' AND COLUMN_NAME = 'temp_id'`).Scan(&tempIdExists) if tempIdExists > 0 { if err := tx.Exec("ALTER TABLE tenant_llm DROP COLUMN temp_id").Error; err != nil { @@ -120,7 +125,7 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { if idColumnExists > 0 { // Modify existing id column to be auto_increment primary key if err := tx.Exec(` - ALTER TABLE tenant_llm + ALTER TABLE tenant_llm MODIFY COLUMN id BIGINT NOT NULL AUTO_INCREMENT PRIMARY KEY `).Error; err != nil { return fmt.Errorf("failed to modify id column: %w", err) @@ -128,7 +133,7 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { } else { // Add id column as auto_increment primary key if err := tx.Exec(` - ALTER TABLE tenant_llm + ALTER TABLE tenant_llm ADD COLUMN id BIGINT NOT NULL AUTO_INCREMENT PRIMARY KEY FIRST `).Error; err != nil { return fmt.Errorf("failed to add id column: %w", err) @@ -137,11 +142,11 @@ func migrateTenantLLMPrimaryKey(db *gorm.DB) error { // Add unique index on (tenant_id, llm_factory, llm_name) var idxExists int64 - tx.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + tx.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'tenant_llm' AND INDEX_NAME = 'idx_tenant_llm_unique'`).Scan(&idxExists) if idxExists == 0 { if err := tx.Exec(` - ALTER TABLE tenant_llm + ALTER TABLE tenant_llm ADD UNIQUE INDEX idx_tenant_llm_unique (tenant_id, llm_factory, llm_name) `).Error; err != nil { common.Warn("Failed to add unique index idx_tenant_llm_unique", zap.Error(err)) @@ -161,7 +166,7 @@ func migrateAddUniqueEmail(db *gorm.DB) error { // Check if unique index already exists using raw SQL var count int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'user' AND INDEX_NAME = 'idx_user_email_unique'`).Scan(&count) if count > 0 { return nil @@ -198,12 +203,72 @@ func migrateAddUniqueEmail(db *gorm.DB) error { return nil } +func migrateIngestionTaskDocumentIDUnique(db *gorm.DB) error { + if !db.Migrator().HasTable("ingestion_task") { + return nil + } + + const indexName = "idx_ingestion_task_document_id" + + var uniqueCount int64 + if err := db.Raw(` + SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + WHERE TABLE_NAME = 'ingestion_task' + AND INDEX_NAME = ? + AND COLUMN_NAME = 'document_id' + AND NON_UNIQUE = 0 + `, indexName).Scan(&uniqueCount).Error; err != nil { + return err + } + if uniqueCount > 0 { + return nil + } + + var duplicateCount int64 + if err := db.Raw(` + SELECT COUNT(*) FROM ( + SELECT document_id FROM ingestion_task GROUP BY document_id HAVING COUNT(*) > 1 + ) AS duplicates + `).Scan(&duplicateCount).Error; err != nil { + return err + } + if duplicateCount > 0 { + common.Warn("Found duplicate document_id values in ingestion_task, cannot add unique index", zap.Int64("count", duplicateCount)) + return nil + } + + var existingIndexCount int64 + if err := db.Raw(` + SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + WHERE TABLE_NAME = 'ingestion_task' + AND INDEX_NAME = ? + `, indexName).Scan(&existingIndexCount).Error; err != nil { + return err + } + if existingIndexCount > 0 { + if err := db.Exec(`ALTER TABLE ingestion_task DROP INDEX ` + indexName).Error; err != nil { + return fmt.Errorf("failed to drop existing index %s: %w", indexName, err) + } + } + + if err := db.Exec(`ALTER TABLE ingestion_task ADD UNIQUE INDEX ` + indexName + ` (document_id)`).Error; err != nil { + errStr := err.Error() + if strings.Contains(errStr, "Error 1061") && strings.Contains(errStr, "Duplicate key name") { + common.Info("Index already exists, skipping", zap.String("error", errStr)) + return nil + } + return fmt.Errorf("failed to add unique index on ingestion_task.document_id: %w", err) + } + + return nil +} + // modifyColumnTypes modifies column types that need explicit ALTER statements func modifyColumnTypes(db *gorm.DB) error { // Helper function to check if column exists columnExists := func(table, column string) bool { var count int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = ? AND COLUMN_NAME = ?`, table, column).Scan(&count) return count > 0 } @@ -277,7 +342,7 @@ func renameColumnIfExists(db *gorm.DB, tableName, oldName, newName string) error // Helper to check if column exists columnExists := func(column string) bool { var count int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS + db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = ? AND COLUMN_NAME = ?`, tableName, column).Scan(&count) return count > 0 } @@ -381,7 +446,7 @@ func migrateSkillSearchTables(db *gorm.DB) error { // Drop legacy unique index (tenant_id, embd_id) to allow per-space configs. var legacyIndexExists int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'skill_search_configs' AND INDEX_NAME = 'idx_tenant_embd'`).Scan(&legacyIndexExists) if legacyIndexExists > 0 { common.Info("Dropping legacy unique index idx_tenant_embd from skill_search_configs...") @@ -392,11 +457,11 @@ func migrateSkillSearchTables(db *gorm.DB) error { // Table exists, check if unique index exists var indexExists int64 - db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + db.Raw(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'skill_search_configs' AND INDEX_NAME = 'idx_tenant_space_embd'`).Scan(&indexExists) if indexExists == 0 { common.Info("Adding unique index idx_tenant_space_embd to skill_search_configs...") - if err := db.Exec(`ALTER TABLE skill_search_configs + if err := db.Exec(`ALTER TABLE skill_search_configs ADD UNIQUE INDEX idx_tenant_space_embd (tenant_id, space_id, embd_id)`).Error; err != nil { return fmt.Errorf("failed to add unique index idx_tenant_space_embd: %w", err) } @@ -469,7 +534,7 @@ func migrateSkillSpaceIndex(db *gorm.DB) error { // Check if old index exists and drop it var oldIndexExists int64 db.Raw(` - SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'skill_spaces' AND INDEX_NAME = 'idx_tenant_name' `).Scan(&oldIndexExists) @@ -483,7 +548,7 @@ func migrateSkillSpaceIndex(db *gorm.DB) error { // Check if new index exists var newIndexExists int64 db.Raw(` - SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS + SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_NAME = 'skill_spaces' AND INDEX_NAME = 'idx_tenant_name_status' `).Scan(&newIndexExists) diff --git a/internal/dao/task_test.go b/internal/dao/task_test.go index e4d8093545..882c1411a2 100644 --- a/internal/dao/task_test.go +++ b/internal/dao/task_test.go @@ -145,9 +145,8 @@ func TestDeleteIngestionTasksByDocIDs_Success(t *testing.T) { // Insert ingestion tasks for two different documents task1 := &entity.IngestionTask{ID: "itask-1", DocumentID: "doc-1", UserID: "user-1", DatasetID: "ds-1", Status: "pending"} - task2 := &entity.IngestionTask{ID: "itask-2", DocumentID: "doc-1", UserID: "user-1", DatasetID: "ds-1", Status: "pending"} - task3 := &entity.IngestionTask{ID: "itask-3", DocumentID: "doc-2", UserID: "user-1", DatasetID: "ds-1", Status: "pending"} - for _, tk := range []*entity.IngestionTask{task1, task2, task3} { + task2 := &entity.IngestionTask{ID: "itask-2", DocumentID: "doc-2", UserID: "user-1", DatasetID: "ds-1", Status: "pending"} + for _, tk := range []*entity.IngestionTask{task1, task2} { if err := db.Create(tk).Error; err != nil { t.Fatalf("failed to create ingestion task: %v", err) } @@ -158,8 +157,8 @@ func TestDeleteIngestionTasksByDocIDs_Success(t *testing.T) { if err != nil { t.Fatalf("DeleteIngestionTasksByDocIDs failed: %v", err) } - if rowsAffected != 2 { - t.Fatalf("expected 2 rows affected, got %d", rowsAffected) + if rowsAffected != 1 { + t.Fatalf("expected 1 row affected, got %d", rowsAffected) } // Verify doc-1 tasks are gone, doc-2 remains @@ -170,8 +169,8 @@ func TestDeleteIngestionTasksByDocIDs_Success(t *testing.T) { if len(remaining) != 1 { t.Fatalf("expected 1 task remaining, got %d", len(remaining)) } - if remaining[0].ID != "itask-3" { - t.Fatalf("expected itask-3 to remain, got %s", remaining[0].ID) + if remaining[0].ID != "itask-2" { + t.Fatalf("expected itask-2 to remain, got %s", remaining[0].ID) } } @@ -228,8 +227,7 @@ func TestDeleteIngestionTasksByDocIDs_MultipleIDs(t *testing.T) { {ID: "itask-1", DocumentID: "doc-1", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, {ID: "itask-2", DocumentID: "doc-2", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, {ID: "itask-3", DocumentID: "doc-3", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, - {ID: "itask-4", DocumentID: "doc-2", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, - {ID: "itask-5", DocumentID: "keep", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, + {ID: "itask-4", DocumentID: "keep", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, } for _, tk := range tasks { if err := db.Create(tk).Error; err != nil { @@ -242,8 +240,8 @@ func TestDeleteIngestionTasksByDocIDs_MultipleIDs(t *testing.T) { if err != nil { t.Fatalf("DeleteIngestionTasksByDocIDs failed: %v", err) } - if rowsAffected != 4 { - t.Fatalf("expected 4 rows affected, got %d", rowsAffected) + if rowsAffected != 3 { + t.Fatalf("expected 3 rows affected, got %d", rowsAffected) } // Verify only "keep" remains @@ -251,7 +249,7 @@ func TestDeleteIngestionTasksByDocIDs_MultipleIDs(t *testing.T) { if err := db.Find(&remaining).Error; err != nil { t.Fatalf("failed to find remaining tasks: %v", err) } - if len(remaining) != 1 || remaining[0].ID != "itask-5" { - t.Fatalf("expected only itask-5 to remain, got %d tasks", len(remaining)) + if len(remaining) != 1 || remaining[0].ID != "itask-4" { + t.Fatalf("expected only itask-4 to remain, got %d tasks", len(remaining)) } } diff --git a/internal/deepdoc/parser/pdf/parser_pipeline_integration_test.go b/internal/deepdoc/parser/pdf/parser_pipeline_integration_test.go index 87decd1f0d..1a2cec409c 100644 --- a/internal/deepdoc/parser/pdf/parser_pipeline_integration_test.go +++ b/internal/deepdoc/parser/pdf/parser_pipeline_integration_test.go @@ -11,6 +11,7 @@ import ( _ "image/png" "os" "path/filepath" + "ragflow/internal/common" "strings" "testing" diff --git a/internal/engine/infinity/chunk.go b/internal/engine/infinity/chunk.go index 7e16f2693b..5f824685c9 100644 --- a/internal/engine/infinity/chunk.go +++ b/internal/engine/infinity/chunk.go @@ -2431,9 +2431,14 @@ func transformChunkFields(chunk map[string]interface{}, embeddingCols [][2]inter d["questions"] = utility.ConvertToString(v) } case "kb_id": - if list, ok := v.([]interface{}); ok && len(list) > 0 { + // 1. First check if it's a string + if str, ok := v.(string); ok { + d["kb_id"] = str + } else if list := utility.ConvertToStringSlice(v); len(list) > 0 { + // 2. If it's a list, convert and take the first element d["kb_id"] = list[0] } else { + // 3. Otherwise assign v directly d["kb_id"] = v } case "position_int": diff --git a/internal/entity/ingestion_task.go b/internal/entity/ingestion_task.go index a2e0f44d6e..604388d5e9 100644 --- a/internal/entity/ingestion_task.go +++ b/internal/entity/ingestion_task.go @@ -19,7 +19,7 @@ package entity type IngestionTask struct { ID string `gorm:"column:id;primaryKey;size:32" json:"id"` UserID string `gorm:"column:user_id;size:32;not null" json:"user_id"` - DocumentID string `gorm:"column:document_id;size:32;not null;index" json:"document_id"` + DocumentID string `gorm:"column:document_id;size:32;not null;uniqueIndex:idx_ingestion_task_document_id" json:"document_id"` DatasetID string `gorm:"column:dataset_id;size:32;not null" json:"dataset_id"` Schema JSONMap `gorm:"column:schema;type:longtext" json:"schema"` Status string `gorm:"column:status;size:32;not null;" json:"status"` diff --git a/internal/handler/document.go b/internal/handler/document.go index 06c4851ddb..583a6220b3 100644 --- a/internal/handler/document.go +++ b/internal/handler/document.go @@ -29,6 +29,7 @@ import ( "path/filepath" "ragflow/internal/common" "ragflow/internal/entity" + "ragflow/internal/httputil" "ragflow/internal/utility" "strconv" "strings" @@ -1408,7 +1409,7 @@ func (h *DocumentHandler) ListIngestionTasks(c *gin.Context) { parseResult, err = h.documentService.ListIngestionTasks(userID, req.DatasetID, 0, 0) if err != nil { - common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) + common.ResponseWithCodeData(c, httputil.IngestionTaskErrorCode(err), nil, err.Error()) return } @@ -1458,7 +1459,7 @@ func (h *DocumentHandler) StopIngestionTasks(c *gin.Context) { parseResult, err := h.documentService.StopIngestionTasks(req.Tasks, userID) if err != nil { - common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) + common.ResponseWithCodeData(c, httputil.IngestionTaskErrorCode(err), nil, err.Error()) return } common.SuccessWithData(c, parseResult, "success") @@ -1484,7 +1485,7 @@ func (h *DocumentHandler) RemoveIngestionTasks(c *gin.Context) { deletedTasks, err := h.documentService.RemoveIngestionTasks(req.Tasks, userID) if err != nil { - common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) + common.ResponseWithCodeData(c, httputil.IngestionTaskErrorCode(err), nil, err.Error()) return } common.SuccessWithData(c, deletedTasks, "success") diff --git a/internal/handler/document_test.go b/internal/handler/document_test.go index 06887a0bba..371028ebc4 100644 --- a/internal/handler/document_test.go +++ b/internal/handler/document_test.go @@ -38,42 +38,46 @@ import ( // fakeDocumentService implements documentServiceIface for handler tests. type fakeDocumentService struct { - deleted int - err error - doc *service.DocumentResponse - docErr error - updateCalled bool - updatedID string - deleteCalled bool - deletedID string - stopResult map[string]interface{} - stopErr error - thumbnails map[string]string - thumbnailErr error - thumbnailUserID string - thumbnailDocIDs []string - metadataSummary map[string]interface{} - metadataErr error - metadataKBID string - metadataDocIDs []string - setMetaCalled bool - setMetaDocID string - setMetaValue map[string]interface{} - uploadLocalData []map[string]interface{} - uploadLocalErrs []string - uploadLocalKB *entity.Knowledgebase - uploadLocalPath string - uploadOverride map[string]interface{} - ingestCode common.ErrorCode - ingestErr error - ingestUserID string - ingestReq *service.IngestDocumentRequest - listOpts dao.DocumentListOptions - filterOpts dao.DocumentListOptions - filterResult map[string]interface{} - filterTotal int64 - listIDs []string - metadataByKBs map[string]interface{} + deleted int + err error + doc *service.DocumentResponse + docErr error + updateCalled bool + updatedID string + deleteCalled bool + deletedID string + stopResult map[string]interface{} + stopErr error + stopIngestionTasks []*entity.IngestionTask + stopIngestionTaskErr error + removeIngestionTasks []map[string]string + removeIngestionTaskErr error + thumbnails map[string]string + thumbnailErr error + thumbnailUserID string + thumbnailDocIDs []string + metadataSummary map[string]interface{} + metadataErr error + metadataKBID string + metadataDocIDs []string + setMetaCalled bool + setMetaDocID string + setMetaValue map[string]interface{} + uploadLocalData []map[string]interface{} + uploadLocalErrs []string + uploadLocalKB *entity.Knowledgebase + uploadLocalPath string + uploadOverride map[string]interface{} + ingestCode common.ErrorCode + ingestErr error + ingestUserID string + ingestReq *service.IngestDocumentRequest + listOpts dao.DocumentListOptions + filterOpts dao.DocumentListOptions + filterResult map[string]interface{} + filterTotal int64 + listIDs []string + metadataByKBs map[string]interface{} } func (f *fakeDocumentService) Ingest(userID string, req *service.IngestDocumentRequest) (common.ErrorCode, error) { @@ -246,10 +250,10 @@ func (f *fakeDocumentService) IngestDocuments(datasetID, userID string, docIDs [ return nil, nil } func (f *fakeDocumentService) StopIngestionTasks(tasks []string, userID string) ([]*entity.IngestionTask, error) { - return nil, nil + return f.stopIngestionTasks, f.stopIngestionTaskErr } func (f *fakeDocumentService) RemoveIngestionTasks(tasks []string, userID string) ([]map[string]string, error) { - return nil, nil + return f.removeIngestionTasks, f.removeIngestionTaskErr } func setupGinContextWithUser(method, path, body string) (*gin.Context, *httptest.ResponseRecorder) { @@ -1247,6 +1251,78 @@ func TestStopParseDocumentsHandler_NotAccessible(t *testing.T) { } } +func TestStopIngestionTasksHandler_InvalidTransitionReturnsConflict(t *testing.T) { + gin.SetMode(gin.TestMode) + + fake := &fakeDocumentService{ + stopIngestionTaskErr: &service.InvalidTaskTransitionError{TaskID: "task-1", From: common.CREATED, To: common.COMPLETED}, + } + h := &DocumentHandler{documentService: fake} + + c, w := setupGinContextWithUser("PUT", "/api/v1/datasets/ds-1/ingestion/tasks", `{"tasks":["task-1"]}`) + h.StopIngestionTasks(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["code"] != float64(common.CodeConflict) { + t.Fatalf("expected code %d, got %v", common.CodeConflict, resp["code"]) + } + if !strings.Contains(resp["message"].(string), "task-1") { + t.Fatalf("expected task id in message, got %v", resp["message"]) + } +} + +func TestStopIngestionTasksHandler_TaskNotFoundReturnsNotFound(t *testing.T) { + gin.SetMode(gin.TestMode) + + fake := &fakeDocumentService{ + stopIngestionTaskErr: common.ErrTaskNotFound, + } + h := &DocumentHandler{documentService: fake} + + c, w := setupGinContextWithUser("PUT", "/api/v1/datasets/ds-1/ingestion/tasks", `{"tasks":["task-1"]}`) + h.StopIngestionTasks(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["code"] != float64(common.CodeNotFound) { + t.Fatalf("expected code %d, got %v", common.CodeNotFound, resp["code"]) + } +} + +func TestRemoveIngestionTasksHandler_TaskNotFoundReturnsNotFound(t *testing.T) { + gin.SetMode(gin.TestMode) + + fake := &fakeDocumentService{ + removeIngestionTaskErr: common.ErrTaskNotFound, + } + h := &DocumentHandler{documentService: fake} + + c, w := setupGinContextWithUser("DELETE", "/api/v1/datasets/ds-1/ingestion/tasks", `{"tasks":["task-1"]}`) + h.RemoveIngestionTasks(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["code"] != float64(common.CodeNotFound) { + t.Fatalf("expected code %d, got %v", common.CodeNotFound, resp["code"]) + } +} + func TestMetadataSummaryByDataset_Success(t *testing.T) { db := setupHandlerAccessDB(t) orig := dao.DB diff --git a/internal/httputil/ingestion_task_error.go b/internal/httputil/ingestion_task_error.go new file mode 100644 index 0000000000..784eed0eea --- /dev/null +++ b/internal/httputil/ingestion_task_error.go @@ -0,0 +1,23 @@ +package httputil + +import ( + "errors" + + "ragflow/internal/common" + "ragflow/internal/service" +) + +func IngestionTaskErrorCode(err error) common.ErrorCode { + var transitionErr *service.InvalidTaskTransitionError + if errors.As(err, &transitionErr) { + return common.CodeConflict + } + var conflictErr *service.TaskStatusConflictError + if errors.As(err, &conflictErr) { + return common.CodeConflict + } + if errors.Is(err, common.ErrTaskNotFound) { + return common.CodeNotFound + } + return common.CodeExceptionError +} diff --git a/internal/httputil/ingestion_task_error_test.go b/internal/httputil/ingestion_task_error_test.go new file mode 100644 index 0000000000..f87b35bca6 --- /dev/null +++ b/internal/httputil/ingestion_task_error_test.go @@ -0,0 +1,43 @@ +package httputil + +import ( + "testing" + + "ragflow/internal/common" + "ragflow/internal/service" +) + +func TestIngestionTaskErrorCode(t *testing.T) { + testCases := []struct { + name string + err error + want common.ErrorCode + }{ + { + name: "invalid transition", + err: &service.InvalidTaskTransitionError{TaskID: "task-1", From: common.CREATED, To: common.COMPLETED}, + want: common.CodeConflict, + }, + { + name: "status conflict", + err: &service.TaskStatusConflictError{TaskID: "task-1", ExpectedFrom: common.CREATED, AttemptedTo: common.RUNNING, ActualCurrent: common.STOPPING}, + want: common.CodeConflict, + }, + { + name: "task not found", + err: common.ErrTaskNotFound, + want: common.CodeNotFound, + }, + { + name: "fallback", + err: common.ErrInvalidToken, + want: common.CodeExceptionError, + }, + } + + for _, tc := range testCases { + if got := IngestionTaskErrorCode(tc.err); got != tc.want { + t.Fatalf("%s: got %d, want %d", tc.name, got, tc.want) + } + } +} diff --git a/internal/ingestion/pipeline/real_storage_integration_test.go b/internal/ingestion/pipeline/real_storage_integration_test.go index 9f5b75f976..e66fa4c3e8 100644 --- a/internal/ingestion/pipeline/real_storage_integration_test.go +++ b/internal/ingestion/pipeline/real_storage_integration_test.go @@ -11,6 +11,7 @@ import ( "io" "os" "path/filepath" + "ragflow/internal/common" "sort" "strings" "testing" diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index 71a0b3c916..e1f891a3b2 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -28,6 +28,7 @@ import ( "ragflow/internal/engine" "ragflow/internal/entity" taskpkg "ragflow/internal/ingestion/task" + servicepkg "ragflow/internal/service" "google.golang.org/grpc" "gorm.io/gorm" @@ -61,6 +62,7 @@ type Ingestor struct { ingestionTaskDAO *dao.IngestionTaskDAO ingestionTaskLogDAO *dao.IngestionTaskLogDAO + ingestionTaskSvc *servicepkg.IngestionTaskService // runDocumentTask dispatches to the migrated task handler path. // Tests may override this to verify branch routing without invoking @@ -84,6 +86,7 @@ func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *In ShutdownCh: make(chan struct{}, 1), ingestionTaskDAO: dao.NewIngestionTaskDAO(), ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), + ingestionTaskSvc: servicepkg.NewIngestionTaskService(), } ingestor.runDocumentTask = ingestor.defaultRunDocumentTask return ingestor @@ -124,7 +127,7 @@ func (e *Ingestor) Start() error { continue } var task *entity.IngestionTask - task, err = e.ingestionTaskDAO.SetRunningByIngestor(taskMessage.TaskID) + task, err = e.ingestionTaskSvc.StartRunning(taskMessage.TaskID) if err != nil { if errors.Is(err, common.ErrTaskNotFound) { common.Warn(fmt.Sprintf("task %s not found, skipping", taskMessage.TaskID)) @@ -216,7 +219,7 @@ func (e *Ingestor) executeTask(taskCtx *taskpkg.TaskContext) { if err != nil { if !errors.Is(err, gorm.ErrRecordNotFound) { common.Error(fmt.Sprintf("Failed to get latest task log for task %s", task.ID), err) - if uErr := e.ingestionTaskDAO.UpdateStatus(task.ID, common.FAILED); uErr != nil { + if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) } return @@ -232,7 +235,7 @@ func (e *Ingestor) executeTask(taskCtx *taskpkg.TaskContext) { err = e.ingestionTaskLogDAO.Create(latestLog) if err != nil { common.Error(fmt.Sprintf("Failed to create task log for task %s", task.ID), err) - if uErr := e.ingestionTaskDAO.UpdateStatus(task.ID, common.FAILED); uErr != nil { + if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) } return @@ -244,7 +247,7 @@ func (e *Ingestor) executeTask(taskCtx *taskpkg.TaskContext) { currentStep, ok := common.GetInt(checkpointMap["current_step"]) if !ok { common.Error(fmt.Sprintf("Failed to get current step from task log for task %s", task.ID), nil) - if uErr := e.ingestionTaskDAO.UpdateStatus(task.ID, common.FAILED); uErr != nil { + if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) } return @@ -252,7 +255,7 @@ func (e *Ingestor) executeTask(taskCtx *taskpkg.TaskContext) { totalStep, ok := common.GetInt(checkpointMap["total_step"]) if !ok { common.Error(fmt.Sprintf("Failed to get total step from task log for task %s", task.ID), nil) - if uErr := e.ingestionTaskDAO.UpdateStatus(task.ID, common.FAILED); uErr != nil { + if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) } return @@ -273,13 +276,13 @@ func (e *Ingestor) executeTask(taskCtx *taskpkg.TaskContext) { } if err := e.runDocumentTask(ctx, task); err != nil { common.Error(fmt.Sprintf("Task %s failed", task.ID), err) - if uErr := e.ingestionTaskDAO.UpdateStatus(task.ID, common.FAILED); uErr != nil { + if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) } return } - err = e.ingestionTaskDAO.UpdateStatus(task.ID, common.COMPLETED) + err = e.ingestionTaskSvc.MarkCompleted(task.ID) if err != nil { common.Error(fmt.Sprintf("Task %s update status failed", task.ID), err) return diff --git a/internal/ingestion/task/dataflow_chunk_assets.go b/internal/ingestion/task/dataflow_chunk_assets.go deleted file mode 100644 index f297e65e66..0000000000 --- a/internal/ingestion/task/dataflow_chunk_assets.go +++ /dev/null @@ -1,14 +0,0 @@ -package task - -// PrepareDataflowChunkAssets applies the minimal pre-index cleanup needed by -// the Go dataflow path so real pipeline output can be stored safely. -func PrepareDataflowChunkAssets(chunks []map[string]any) error { - for _, ck := range chunks { - delete(ck, "_pdf_positions") - // FIXME: production parity with Python _prepare_docs_and_upload should - // upload raw image data and set img_id. During the current bring-up phase - // we drop raw image payloads so the main dataflow -> ES path can proceed. - delete(ck, "image") - } - return nil -} diff --git a/internal/ingestion/task/dataflow_e2e_test.go b/internal/ingestion/task/dataflow_e2e_test.go index 5abe220172..a980f0f2de 100644 --- a/internal/ingestion/task/dataflow_e2e_test.go +++ b/internal/ingestion/task/dataflow_e2e_test.go @@ -20,16 +20,17 @@ import ( "context" "fmt" "net" - "ragflow/internal/common" "testing" "time" + "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/engine" "ragflow/internal/engine/elasticsearch" "ragflow/internal/engine/infinity" "ragflow/internal/ingestion/testutil" "ragflow/internal/server" + "ragflow/internal/service" ) // ============================================================================= @@ -184,15 +185,15 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { cleanupDB := testutil.ReplaceDBForTest(t, db) defer cleanupDB() - // Seed test data (lowercase for ES index compatibility) + // Seed test data (lowercase for ES index compatibility, no hyphens for Infinity) lowerName := toLowerSnakeCase(tc.name) tenantID, kbID, _, taskID := testutil.SeedTestData(t, db, - testutil.WithTenantID(fmt.Sprintf("tenant-e2e-%s", lowerName)), - testutil.WithKBID(fmt.Sprintf("kb-e2e-%s", lowerName)), - testutil.WithDocID(fmt.Sprintf("doc-e2e-%s", lowerName)), - testutil.WithTaskID(fmt.Sprintf("task-e2e-%s", lowerName)), - testutil.WithPipelineID(fmt.Sprintf("pipeline-e2e-%s", lowerName)), - testutil.WithDocName(fmt.Sprintf("e2e-test-%s.pdf", lowerName)), + testutil.WithTenantID(fmt.Sprintf("tenant_e2e_%s", lowerName)), + testutil.WithKBID(fmt.Sprintf("kb_e2e_%s", lowerName)), + testutil.WithDocID(fmt.Sprintf("doc_e2e_%s", lowerName)), + testutil.WithTaskID(fmt.Sprintf("task_e2e_%s", lowerName)), + testutil.WithPipelineID(fmt.Sprintf("pipeline_e2e_%s", lowerName)), + testutil.WithDocName(fmt.Sprintf("e2e_test_%s.pdf", lowerName)), ) // Setup DocEngine for this test @@ -235,12 +236,12 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { "chunks": []map[string]any{ { "text": fmt.Sprintf("Hello world from E2E test with %s", tc.name), - "id": fmt.Sprintf("chunk-e2e-%s-1", tc.name), + "id": fmt.Sprintf("chunk_e2e_%s_1", lowerName), "q_2_vec": []float64{0.1, 0.2}, // Pre-vectorized to skip embedding }, { "text": fmt.Sprintf("Second chunk from E2E test with %s", tc.name), - "id": fmt.Sprintf("chunk-e2e-%s-2", tc.name), + "id": fmt.Sprintf("chunk_e2e_%s_2", lowerName), "q_2_vec": []float64{0.3, 0.4}, // Pre-vectorized to skip embedding }, }, @@ -343,18 +344,18 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { t.Fatal("Expected progress to reach 1.0") } - // Verify final task status can be updated to success - ingestionTaskDAO := dao.NewIngestionTaskDAO() - if err = ingestionTaskDAO.UpdateStatus(taskID, "success"); err != nil { - t.Fatalf("UpdateStatus failed: %v", err) + // Verify final task status can be marked completed + ingestSvc := service.NewIngestionTaskService() + if err := ingestSvc.MarkCompleted(taskID); err != nil { + t.Fatalf("MarkCompleted failed: %v", err) } - finalTask, err := ingestionTaskDAO.GetByID(taskID) + finalTask, err := dao.NewIngestionTaskDAO().GetByID(taskID) if err != nil { t.Fatalf("GetByID failed: %v", err) } - if finalTask.Status != "success" { - t.Errorf("Final task status = %q, want %q", finalTask.Status, "success") + if finalTask.Status != common.COMPLETED { + t.Errorf("Final task status = %q, want %q", finalTask.Status, common.COMPLETED) } t.Logf("SUCCESS: Dataflow E2E test passed with %s engine!", tc.name) diff --git a/internal/ingestion/task/dataflow_service.go b/internal/ingestion/task/dataflow_service.go index e3c06bfc8e..af79a8c13d 100644 --- a/internal/ingestion/task/dataflow_service.go +++ b/internal/ingestion/task/dataflow_service.go @@ -22,7 +22,6 @@ import ( "fmt" componentpkg "ragflow/internal/ingestion/component" "ragflow/internal/utility" - "regexp" "sort" "strings" "time" @@ -91,22 +90,6 @@ func (d *defaultChunkCounter) IncrementChunkNum(docID, kbID string, chunkNum, to return service.NewDocumentService().IncrementChunkNum(docID, kbID, chunkNum, tokenConsumption, duration) } -func encodeTexts(model *models.EmbeddingModel, texts []string) ([][]float64, int, error) { - texts = TruncateTexts(texts, model.MaxTokens) - config := &models.EmbeddingConfig{Dimension: 0} - embeds, err := model.ModelDriver.Embed(model.ModelName, texts, model.APIConfig, config) - if err != nil { - return nil, 0, err - } - vecs := make([][]float64, len(embeds)) - totalTokens := 0 - for i, v := range embeds { - vecs[i] = v.Embedding - totalTokens += v.TokenCount - } - return vecs, totalTokens, nil -} - type PipelineExecutor struct { taskCtx *TaskContext dataflowID string @@ -310,9 +293,6 @@ func (s *PipelineExecutor) processOutput(ctx context.Context, pipelineOutput map embeddingTokenConsumption := GetEmbeddingTokenConsumption(pipelineOutput) metadata := s.processChunks(chunks) - if err := s.prepareChunkAssets(chunks); err != nil { - return err - } if len(metadata) > 0 { if err := s.updateDocumentMetadata(s.taskCtx.Doc.ID, metadata); err != nil { @@ -350,10 +330,6 @@ func (s *PipelineExecutor) processChunks(chunks []map[string]any) map[string]any ) } -func (s *PipelineExecutor) prepareChunkAssets(chunks []map[string]any) error { - return PrepareDataflowChunkAssets(chunks) -} - func (s *PipelineExecutor) insertChunks(ctx context.Context, chunks []map[string]any) error { baseName := fmt.Sprintf("ragflow_%s", s.taskCtx.Tenant.ID) if len(chunks) == 0 { @@ -433,19 +409,6 @@ func (s *PipelineExecutor) progress(prog float64, msg string) { } } -func hasVectors(chunks []map[string]any) bool { - for _, ck := range chunks { - for k := range ck { - if matchQVec.MatchString(k) { - return true - } - } - } - return false -} - -var matchQVec = regexp.MustCompile(`^q_\d+_vec$`) - func (s *PipelineExecutor) defaultLoadDSL(ctx context.Context, dataflowID string) (string, string, error) { if s == nil || s.taskCtx == nil { return "", "", fmt.Errorf("dataflow service: nil task context") diff --git a/internal/ingestion/task/embedding_test.go b/internal/ingestion/task/embedding_test.go index 1b3079b2f0..89c362e9bf 100644 --- a/internal/ingestion/task/embedding_test.go +++ b/internal/ingestion/task/embedding_test.go @@ -1,10 +1,6 @@ package task import ( - "errors" - "strings" - "testing" - "ragflow/internal/entity/models" ) @@ -79,122 +75,3 @@ func makeTestEmbeddingModel(stub *stubDriver, maxTokens int) *models.EmbeddingMo MaxTokens: maxTokens, } } - -// ============================================================================= -// encodeTexts — must truncate texts before passing to ModelDriver.Embed -// Python: truncated = truncate_texts(txts, model.max_length); model.encode(truncated) -// ============================================================================= - -func TestEncodeTexts_TruncatesBeforeEmbed(t *testing.T) { - stub := &stubDriver{} - model := makeTestEmbeddingModel(stub, 12) - // "hello world" is 2 tokens, safeMax = 12-10 = 2, fits - texts := []string{"hello world"} - _, _, err := encodeTexts(model, texts) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if len(stub.capturedTexts) != 1 { - t.Fatalf("captured %d texts, want 1", len(stub.capturedTexts)) - } - if stub.capturedTexts[0] != "hello world" { - t.Errorf("texts should pass through when within limit, got %q", stub.capturedTexts[0]) - } -} - -func TestEncodeTexts_TruncatesLongText(t *testing.T) { - stub := &stubDriver{} - model := makeTestEmbeddingModel(stub, 12) - // safeMax = 12-10 = 2 tokens. Text > 2 tokens should be truncated. - longText := strings.Repeat("hello world ", 20) // 40 tokens - texts := []string{longText} - _, _, err := encodeTexts(model, texts) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if len(stub.capturedTexts) != 1 { - t.Fatalf("captured %d texts, want 1", len(stub.capturedTexts)) - } - // The text passed to Embed must be SHORTER than the original - if len(stub.capturedTexts[0]) >= len(longText) { - t.Errorf("text should be truncated before embed: original=%d, truncated=%d", - len(longText), len(stub.capturedTexts[0])) - } -} - -func TestEncodeTexts_TokenCountSum(t *testing.T) { - stub := &stubDriver{} - model := makeTestEmbeddingModel(stub, 100) - texts := []string{"hello", "world"} - _, totalTokens, err := encodeTexts(model, texts) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if totalTokens <= 0 { - t.Errorf("expected positive token count, got %d", totalTokens) - } -} - -func TestEncodeTexts_ReturnsVectors(t *testing.T) { - stub := &stubDriver{} - model := makeTestEmbeddingModel(stub, 100) - texts := []string{"a", "b", "c"} - vecs, _, err := encodeTexts(model, texts) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if len(vecs) != 3 { - t.Fatalf("expected 3 vectors, got %d", len(vecs)) - } -} - -type errDriver struct{} - -func (d *errDriver) Embed(modelName *string, texts []string, apiConfig *models.APIConfig, embeddingConfig *models.EmbeddingConfig) ([]models.EmbeddingData, error) { - return nil, errors.New("embed error") -} -func (d *errDriver) NewInstance(baseURL map[string]string) models.ModelDriver { return d } -func (d *errDriver) Name() string { return "err" } -func (d *errDriver) ChatWithMessages(modelName string, messages []models.Message, apiConfig *models.APIConfig, chatModelConfig *models.ChatConfig) (*models.ChatResponse, error) { - return nil, nil -} -func (d *errDriver) ChatStreamlyWithSender(modelName string, messages []models.Message, apiConfig *models.APIConfig, modelConfig *models.ChatConfig, sender func(*string, *string) error) error { - return nil -} -func (d *errDriver) Rerank(modelName *string, query string, documents []string, apiConfig *models.APIConfig, rerankConfig *models.RerankConfig) (*models.RerankResponse, error) { - return nil, nil -} -func (d *errDriver) TranscribeAudio(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig) (*models.ASRResponse, error) { - return nil, nil -} -func (d *errDriver) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig, sender func(*string, *string) error) error { - return nil -} -func (d *errDriver) AudioSpeech(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig) (*models.TTSResponse, error) { - return nil, nil -} -func (d *errDriver) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig, sender func(*string, *string) error) error { - return nil -} -func (d *errDriver) OCRFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, ocrConfig *models.OCRConfig) (*models.OCRFileResponse, error) { - return nil, nil -} -func (d *errDriver) ParseFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, parseFileConfig *models.ParseFileConfig) (*models.ParseFileResponse, error) { - return nil, nil -} -func (d *errDriver) ListModels(apiConfig *models.APIConfig) ([]models.ListModelResponse, error) { - return nil, nil -} -func (d *errDriver) Balance(apiConfig *models.APIConfig) (map[string]interface{}, error) { - return nil, nil -} -func (d *errDriver) CheckConnection(apiConfig *models.APIConfig) error { return nil } -func (d *errDriver) ListTasks(apiConfig *models.APIConfig) ([]models.ListTaskStatus, error) { - return nil, nil -} -func (d *errDriver) ShowTask(taskID string, apiConfig *models.APIConfig) (*models.TaskResponse, error) { - return nil, nil -} -func (d *errDriver) ToolCall(name string, arguments map[string]interface{}) (string, error) { - return "", nil -} diff --git a/internal/ingestion/task/has_vectors_test.go b/internal/ingestion/task/has_vectors_test.go deleted file mode 100644 index 5f61d9cc38..0000000000 --- a/internal/ingestion/task/has_vectors_test.go +++ /dev/null @@ -1,74 +0,0 @@ -package task - -import ( - "testing" -) - -// ============================================================================= -// hasVectors — checks if any chunk has q_\d+_vec key -// ============================================================================= - -func TestHasVectors_WithVector(t *testing.T) { - chunks := []map[string]any{ - {"text": "hello", "q_768_vec": []float64{0.1, 0.2}}, - } - if !hasVectors(chunks) { - t.Error("expected true for chunk with q_768_vec") - } -} - -func TestHasVectors_WithoutVector(t *testing.T) { - chunks := []map[string]any{ - {"text": "hello"}, - } - if hasVectors(chunks) { - t.Error("expected false for chunk without q_*_vec") - } -} - -func TestHasVectors_MultipleChunksOneHasVector(t *testing.T) { - chunks := []map[string]any{ - {"text": "hello"}, - {"text": "world", "q_3_vec": []float64{0.1, 0.2, 0.3}}, - } - if !hasVectors(chunks) { - t.Error("expected true when one chunk has vector") - } -} - -func TestHasVectors_VectorInFirstChunkOnly(t *testing.T) { - chunks := []map[string]any{ - {"q_128_vec": []float64{0.5}, "text": "a"}, - {"text": "b"}, - } - if !hasVectors(chunks) { - t.Error("expected true when first chunk has vector") - } -} - -func TestHasVectors_EmptyChunks(t *testing.T) { - if hasVectors(nil) { - t.Error("expected false for nil chunks") - } - if hasVectors([]map[string]any{}) { - t.Error("expected false for empty chunks") - } -} - -func TestHasVectors_OtherNumberKeys(t *testing.T) { - chunks := []map[string]any{ - {"text": "hello", "q_abc_vec": "not a vector"}, - } - if hasVectors(chunks) { - t.Error("q_abc_vec should not match q_\\d+_vec") - } -} - -func TestHasVectors_EmptySliceValue(t *testing.T) { - chunks := []map[string]any{ - {"text": "hello", "q_0_vec": []float64{}}, - } - if !hasVectors(chunks) { - t.Error("q_0_vec should match even with empty vector") - } -} diff --git a/internal/ingestion/task/real_consumer_test.go b/internal/ingestion/task/real_consumer_test.go index 37ebd3ee5f..f556880a1f 100644 --- a/internal/ingestion/task/real_consumer_test.go +++ b/internal/ingestion/task/real_consumer_test.go @@ -37,7 +37,7 @@ import ( // TestRealProducerConsumer exercises the project's real producer and consumer code paths: // -// Producer: document.go pattern — CheckAndCreate(IngestionTask) → PublishTask(NATS) +// Producer: document.go pattern — Create(IngestionTask) → PublishTask(NATS) // Consumer: Ingestor.Start() core logic — calls each actual function in sequence func TestRealProducerConsumer(t *testing.T) { // ── 1. NATS ── @@ -89,9 +89,9 @@ func TestRealProducerConsumer(t *testing.T) { DatasetID: "kb1", Status: common.CREATED, } - created, err := dao.NewIngestionTaskDAO().CheckAndCreate(ingestionTask) + created, err := dao.NewIngestionTaskDAO().Create(ingestionTask) if err != nil { - t.Fatalf("CheckAndCreate: %v", err) + t.Fatalf("Create: %v", err) } t.Logf("Producer: IngestionTask created id=%s status=%s", created.ID, created.Status) diff --git a/internal/service/document.go b/internal/service/document.go index e65433b2dc..69a0908d22 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -62,6 +62,7 @@ type DocumentService struct { kbDAO *dao.KnowledgebaseDAO ingestionTaskDAO *dao.IngestionTaskDAO ingestionTaskLogDAO *dao.IngestionTaskLogDAO + ingestionTaskSvc *IngestionTaskService docEngine engine.DocEngine engineType server.EngineType metadataSvc *MetadataService @@ -70,16 +71,19 @@ type DocumentService struct { fileDAO *dao.FileDAO canvasDAO *dao.UserCanvasDAO api4ConvDAO *dao.API4ConversationDAO - publishTask func(subject string, payload []byte) error } // NewDocumentService create document service func NewDocumentService() *DocumentService { cfg := server.GetConfig() + publisher := NewMessageQueueTaskPublisher() + ingestionTaskSvc := NewIngestionTaskService() + ingestionTaskSvc.SetTaskPublisher(publisher) return &DocumentService{ documentDAO: dao.NewDocumentDAO(), ingestionTaskDAO: dao.NewIngestionTaskDAO(), ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), + ingestionTaskSvc: ingestionTaskSvc, kbDAO: dao.NewKnowledgebaseDAO(), docEngine: engine.Get(), engineType: cfg.DocEngine.Type, @@ -89,13 +93,6 @@ func NewDocumentService() *DocumentService { fileDAO: dao.NewFileDAO(), canvasDAO: dao.NewUserCanvasDAO(), api4ConvDAO: dao.NewAPI4ConversationDAO(), - publishTask: func(subject string, payload []byte) error { - msgQueueEngine := engine.GetMessageQueueEngine() - if msgQueueEngine == nil { - return fmt.Errorf("message queue engine is not initialized") - } - return msgQueueEngine.PublishTask(subject, payload) - }, } } @@ -720,7 +717,7 @@ func (s *DocumentService) deleteDocumentFull(docID string) error { return err } if ingestionTask != nil { - taskInfo, err := s.ingestionTaskDAO.Delete(ingestionTask.ID, &ingestionTask.UserID) + taskInfo, err := s.ingestionTaskSvc.Remove(ingestionTask.ID, &ingestionTask.UserID) if err != nil { return err } @@ -1163,21 +1160,7 @@ func (s *DocumentService) GetDocumentsByAuthorID(authorID, page, pageSize int) ( } func (s *DocumentService) ListIngestionTasks(userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) { - offset := (page - 1) * pageSize - - var tasks []*entity.IngestionTask - var err error - if datasetID == nil { - tasks, err = s.ingestionTaskDAO.ListByUserID(userID, offset, pageSize) - } else { - tasks, err = s.ingestionTaskDAO.ListByUserIDAndDatasetID(userID, *datasetID, offset, pageSize) - } - - if err != nil { - return nil, err - } - - return tasks, nil + return s.ingestionTaskSvc.ListByUser(userID, datasetID, page, pageSize) } type ParseDocumentResponse struct { @@ -1186,112 +1169,20 @@ type ParseDocumentResponse struct { } func (s *DocumentService) IngestDocuments(datasetID, userID string, docIDs []string) ([]*ParseDocumentResponse, error) { - // deduplicate the document id - uniqueDocIDs := common.Deduplicate(docIDs) - if uniqueDocIDs == nil || len(uniqueDocIDs) == 0 { - return nil, fmt.Errorf("no documents to parse") + responses, err := s.ingestionTaskSvc.CreateForDocuments(datasetID, userID, docIDs) + if err != nil { + return nil, err } - - var responses []*ParseDocumentResponse - - // query database, if the document ids are valid - for _, docID := range uniqueDocIDs { - doc, err := s.documentDAO.GetByID(docID) - - if err != nil { - errorMessage := err.Error() - responses = append(responses, &ParseDocumentResponse{ - DocumentID: docID, - Result: errorMessage, - }) - continue - } - - if doc == nil { - errorMessage := "no such document" - responses = append(responses, &ParseDocumentResponse{ - DocumentID: docID, - Result: errorMessage, - }) - continue - } - - task := &entity.IngestionTask{ - DocumentID: docID, - UserID: userID, - DatasetID: datasetID, - Schema: nil, - Status: common.CREATED, - } - - // save the task to database - task, err = s.ingestionTaskDAO.CheckAndCreate(task) - if err != nil { - errorMessage := err.Error() - responses = append(responses, &ParseDocumentResponse{ - DocumentID: docID, - Result: errorMessage, - }) - continue - } - - msgQueueEngine := engine.GetMessageQueueEngine() - - taskMessage := common.TaskMessage{ - TaskID: task.ID, - TaskType: common.TaskTypeIngestionTask, - } - - // convert task - taskMessageStr, err := json.Marshal(taskMessage) - if err != nil { - return nil, err - } - - err = msgQueueEngine.PublishTask("tasks.RAGFLOW", taskMessageStr) - if err != nil { - return nil, err - } - - responses = append(responses, &ParseDocumentResponse{ - DocumentID: docID, - Result: fmt.Sprintf("task_id: %s", task.ID), - }) - } - common.Info(fmt.Sprintf("parse documents, dataset: %s, documents: %v", datasetID, docIDs)) return responses, nil } func (s *DocumentService) StopIngestionTasks(tasks []string, userID string) ([]*entity.IngestionTask, error) { - - var taskResponses []*entity.IngestionTask - for _, taskID := range tasks { - task, err := s.ingestionTaskDAO.SetStoppingByAPIServer(taskID) - if err != nil { - return nil, err - } - taskResponses = append(taskResponses, task) - } - return taskResponses, nil + return s.ingestionTaskSvc.RequestStopMany(tasks, &userID) } func (s *DocumentService) RemoveIngestionTasks(tasks []string, userID string) ([]map[string]string, error) { - - var deletedTasks []map[string]string - for _, taskID := range tasks { - taskRecord := map[string]string{ - "task_id": taskID, - } - _, err := s.ingestionTaskDAO.Delete(taskID, &userID) - if err != nil { - taskRecord["remove"] = fmt.Sprintf("fail: %s", err.Error()) - } else { - taskRecord["remove"] = "success" - } - deletedTasks = append(deletedTasks, taskRecord) - } - return deletedTasks, nil + return s.ingestionTaskSvc.RemoveMany(tasks, &userID) } type IngestDocumentRequest struct { @@ -1563,26 +1454,14 @@ func (s *DocumentService) queueDocumentDataflowTask(kb *entity.Knowledgebase, do if err := s.beginDocumentParse(doc.ID); err != nil { return err } - task := &entity.IngestionTask{ + _, err := s.ingestionTaskSvc.CreateAndEnqueue(&entity.IngestionTask{ DocumentID: doc.ID, UserID: userID, DatasetID: kb.ID, Schema: nil, Status: common.CREATED, - } - task, err := s.ingestionTaskDAO.CheckAndCreate(task) - if err != nil { - return err - } - taskMessage := common.TaskMessage{ - TaskID: task.ID, - TaskType: common.TaskTypeIngestionTask, - } - payload, err := json.Marshal(taskMessage) - if err != nil { - return err - } - return s.publishTask("tasks.RAGFLOW", payload) + }) + return err } func (s *DocumentService) newDocumentParseTasks(doc *entity.Document, bucket, objectName string, priority int64) ([]*entity.Task, error) { diff --git a/internal/service/document_test.go b/internal/service/document_test.go index 0ec66ed8fe..d1d7b9670f 100644 --- a/internal/service/document_test.go +++ b/internal/service/document_test.go @@ -422,6 +422,7 @@ func testDocumentService(t *testing.T) *DocumentService { file2DocumentDAO: dao.NewFile2DocumentDAO(), fileDAO: dao.NewFileDAO(), ingestionTaskDAO: dao.NewIngestionTaskDAO(), + ingestionTaskSvc: NewIngestionTaskService(), docEngine: nil, metadataSvc: nil, // nil engine → metadata ops skipped } @@ -1082,14 +1083,9 @@ func TestQueueDocumentDataflowTask_PublishesIngestionTaskMessage(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) insertTestDoc(t, "doc-1", "kb-1", 9, 4) - var publishedSubject string - var publishedPayload []byte + publisher := &recordingTaskPublisher{} svc := testDocumentService(t) - svc.publishTask = func(subject string, payload []byte) error { - publishedSubject = subject - publishedPayload = append([]byte(nil), payload...) - return nil - } + svc.ingestionTaskSvc.SetTaskPublisher(publisher) kb, err := svc.kbDAO.GetByID("kb-1") if err != nil { @@ -1104,17 +1100,14 @@ func TestQueueDocumentDataflowTask_PublishesIngestionTaskMessage(t *testing.T) { t.Fatalf("queueDocumentDataflowTask: %v", err) } - if publishedSubject != "tasks.RAGFLOW" { - t.Fatalf("subject = %q, want %q", publishedSubject, "tasks.RAGFLOW") + if publisher.subject != "tasks.RAGFLOW" { + t.Fatalf("subject = %q, want %q", publisher.subject, "tasks.RAGFLOW") } - if len(publishedPayload) == 0 { - t.Fatal("expected published payload") + if len(publisher.messages) != 1 { + t.Fatalf("expected 1 published message, got %d", len(publisher.messages)) } - var msg common.TaskMessage - if err := json.Unmarshal(publishedPayload, &msg); err != nil { - t.Fatalf("unmarshal payload: %v", err) - } + msg := publisher.messages[0] if msg.TaskType != common.TaskTypeIngestionTask { t.Fatalf("task type = %q, want %q", msg.TaskType, common.TaskTypeIngestionTask) } diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go new file mode 100644 index 0000000000..586a94dc56 --- /dev/null +++ b/internal/service/ingestion_task_service.go @@ -0,0 +1,357 @@ +package service + +import ( + "errors" + "fmt" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + + "gorm.io/gorm" +) + +type InvalidTaskTransitionError struct { + TaskID string + From string + To string +} + +func (e *InvalidTaskTransitionError) Error() string { + return fmt.Sprintf("task %s status cannot transition from %s to %s", e.TaskID, e.From, e.To) +} + +type TaskStatusConflictError struct { + TaskID string + ExpectedFrom string + AttemptedTo string + ActualCurrent string +} + +func (e *TaskStatusConflictError) Error() string { + return fmt.Sprintf("task %s status conflict: expected %s -> %s, actual current %s", e.TaskID, e.ExpectedFrom, e.AttemptedTo, e.ActualCurrent) +} + +type IngestionTaskService struct { + documentDAO *dao.DocumentDAO + userDAO *dao.UserDAO + ingestionTaskDAO *dao.IngestionTaskDAO + ingestionTaskLogDAO *dao.IngestionTaskLogDAO + taskPublisher TaskPublisher +} + +func NewIngestionTaskService() *IngestionTaskService { + return &IngestionTaskService{ + documentDAO: dao.NewDocumentDAO(), + userDAO: dao.NewUserDAO(), + ingestionTaskDAO: dao.NewIngestionTaskDAO(), + ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), + taskPublisher: NewMessageQueueTaskPublisher(), + } +} + +func (s *IngestionTaskService) SetTaskPublisher(taskPublisher TaskPublisher) { + if taskPublisher == nil { + return + } + s.taskPublisher = taskPublisher +} + +func (s *IngestionTaskService) ListByUser(userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) { + if datasetID == nil { + return s.ingestionTaskDAO.ListByUserID(userID, page, pageSize) + } + return s.ingestionTaskDAO.ListByUserIDAndDatasetID(userID, *datasetID, page, pageSize) +} + +func (s *IngestionTaskService) CreateForDocuments(datasetID, userID string, docIDs []string) ([]*ParseDocumentResponse, error) { + uniqueDocIDs := common.Deduplicate(docIDs) + if len(uniqueDocIDs) == 0 { + return nil, fmt.Errorf("no documents to parse") + } + + responses := make([]*ParseDocumentResponse, 0, len(uniqueDocIDs)) + for _, docID := range uniqueDocIDs { + doc, err := s.documentDAO.GetByID(docID) + if err != nil { + responses = append(responses, &ParseDocumentResponse{ + DocumentID: docID, + Result: err.Error(), + }) + continue + } + if doc == nil { + responses = append(responses, &ParseDocumentResponse{ + DocumentID: docID, + Result: "no such document", + }) + continue + } + + task := &entity.IngestionTask{ + DocumentID: docID, + UserID: userID, + DatasetID: datasetID, + Schema: nil, + Status: common.CREATED, + } + task, err = s.CreateAndEnqueue(task) + if err != nil { + responses = append(responses, &ParseDocumentResponse{ + DocumentID: docID, + Result: err.Error(), + }) + continue + } + + responses = append(responses, &ParseDocumentResponse{ + DocumentID: docID, + Result: fmt.Sprintf("task_id: %s", task.ID), + }) + } + return responses, nil +} + +func (s *IngestionTaskService) RequestStopMany(tasks []string, ownerUserID *string) ([]*entity.IngestionTask, error) { + taskResponses := make([]*entity.IngestionTask, 0, len(tasks)) + for _, taskID := range tasks { + if ownerUserID != nil { + task, err := s.getTask(taskID) + if err != nil { + return nil, err + } + if task.UserID != *ownerUserID { + return nil, errors.New("task does not belong to the user") + } + } + task, err := s.RequestStop(taskID) + if err != nil { + return nil, err + } + taskResponses = append(taskResponses, task) + } + return taskResponses, nil +} + +func (s *IngestionTaskService) RemoveMany(tasks []string, ownerUserID *string) ([]map[string]string, error) { + deletedTasks := make([]map[string]string, 0, len(tasks)) + for _, taskID := range tasks { + taskRecord := map[string]string{"task_id": taskID} + if _, err := s.Remove(taskID, ownerUserID); err != nil { + taskRecord["remove"] = fmt.Sprintf("fail: %s", err.Error()) + } else { + taskRecord["remove"] = "success" + } + deletedTasks = append(deletedTasks, taskRecord) + } + return deletedTasks, nil +} + +func (s *IngestionTaskService) ListAllForAdmin() ([]map[string]interface{}, error) { + ingestionTasks, err := s.ingestionTaskDAO.GetAllTasks(0, 0) + if err != nil { + return nil, err + } + + showTasks := make([]map[string]interface{}, 0, len(ingestionTasks)) + for _, task := range ingestionTasks { + user, err := s.userDAO.GetByTenantID(task.UserID) + if err != nil { + return nil, err + } + + showTask := map[string]interface{}{ + "id": task.ID, + "user_id": task.UserID, + "user": user.Email, + "document_id": task.DocumentID, + "status": task.Status, + } + + latestLog, err := s.ingestionTaskLogDAO.LatestLogByTaskID(task.ID) + if err == nil && latestLog != nil && latestLog.Checkpoint != nil { + if step, ok := latestLog.Checkpoint["current_step"].(float64); ok { + showTask["step"] = int(step) + } + } + + showTasks = append(showTasks, showTask) + } + return showTasks, nil +} + +func (s *IngestionTaskService) StartRunning(taskID string) (*entity.IngestionTask, error) { + task, err := s.getTask(taskID) + if err != nil { + return nil, err + } + switch task.Status { + case common.CREATED: + return s.transition(taskID, common.RUNNING) + case common.STOPPING: + return s.transition(taskID, common.STOPPED) + case common.RUNNING, common.COMPLETED, common.STOPPED, common.FAILED: + return task, nil + default: + return task, fmt.Errorf("task %s has unsupported status %s", taskID, task.Status) + } +} + +func (s *IngestionTaskService) RequestStop(taskID string) (*entity.IngestionTask, error) { + task, err := s.getTask(taskID) + if err != nil { + return nil, err + } + switch task.Status { + case common.CREATED: + return s.transition(taskID, common.STOPPED) + case common.RUNNING: + return s.transition(taskID, common.STOPPING) + default: + return task, nil + } +} + +func (s *IngestionTaskService) MarkCompleted(taskID string) error { + _, err := s.transition(taskID, common.COMPLETED) + return err +} + +func (s *IngestionTaskService) MarkFailed(taskID string) error { + _, err := s.transition(taskID, common.FAILED) + return err +} + +func (s *IngestionTaskService) Remove(taskID string, userID *string) (*dao.TaskInfo, error) { + return s.ingestionTaskDAO.Delete(taskID, userID) +} + +func (s *IngestionTaskService) getTask(taskID string) (*entity.IngestionTask, error) { + task, err := s.ingestionTaskDAO.GetByID(taskID) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, common.ErrTaskNotFound + } + return nil, err + } + return task, nil +} + +func validateTransition(from, to string) error { + switch from { + case common.CREATED: + if to == common.RUNNING || to == common.STOPPED { + return nil + } + case common.RUNNING: + if to == common.STOPPING || to == common.COMPLETED || to == common.FAILED { + return nil + } + case common.STOPPING: + if to == common.STOPPED { + return nil + } + case common.FAILED, common.STOPPED: + if to == common.CREATED { + return nil + } + } + return &InvalidTaskTransitionError{From: from, To: to} +} + +func (s *IngestionTaskService) newTaskStatusConflictError(taskID, expectedFrom, attemptedTo string) error { + current, err := s.getTask(taskID) + if err != nil { + return err + } + return &TaskStatusConflictError{ + TaskID: taskID, + ExpectedFrom: expectedFrom, + AttemptedTo: attemptedTo, + ActualCurrent: current.Status, + } +} + +func (s *IngestionTaskService) transition(taskID string, to string) (*entity.IngestionTask, error) { + task, err := s.getTask(taskID) + if err != nil { + return nil, err + } + if err := validateTransition(task.Status, to); err != nil { + var transitionErr *InvalidTaskTransitionError + if errors.As(err, &transitionErr) { + return task, &InvalidTaskTransitionError{TaskID: taskID, From: transitionErr.From, To: transitionErr.To} + } + return task, err + } + updated, err := s.ingestionTaskDAO.UpdateStatusIfCurrent(taskID, task.Status, to) + if err != nil { + return nil, err + } + if !updated { + return nil, s.newTaskStatusConflictError(taskID, task.Status, to) + } + task.Status = to + return task, nil +} + +func (s *IngestionTaskService) CreateAndEnqueue(task *entity.IngestionTask) (*entity.IngestionTask, error) { + existing, err := s.ingestionTaskDAO.GetByDocumentID(task.DocumentID) + if err != nil { + return nil, err + } + if existing != nil { + switch existing.Status { + case common.FAILED, common.STOPPED: + originalStatus := existing.Status + existing, err = s.transition(existing.ID, common.CREATED) + if err != nil { + return nil, err + } + if err := s.enqueueTask(existing.ID); err != nil { + if rollbackErr := s.rollbackRetriedTask(existing.ID, originalStatus); rollbackErr != nil { + return nil, fmt.Errorf("enqueue task %s: %w (rollback failed: %v)", existing.ID, err, rollbackErr) + } + return nil, err + } + return existing, nil + default: + return nil, fmt.Errorf("document id %s already exists, status: %s, task id: %s", task.DocumentID, existing.Status, existing.ID) + } + } + created, err := s.ingestionTaskDAO.Create(task) + if err != nil { + return nil, err + } + if err := s.enqueueTask(created.ID); err != nil { + if rollbackErr := s.rollbackCreatedTask(created.ID); rollbackErr != nil { + return nil, fmt.Errorf("enqueue task %s: %w (rollback failed: %v)", created.ID, err, rollbackErr) + } + return nil, err + } + return created, nil +} + +func (s *IngestionTaskService) rollbackRetriedTask(taskID, status string) error { + updated, err := s.ingestionTaskDAO.UpdateStatusIfCurrent(taskID, common.CREATED, status) + if err != nil { + return err + } + if !updated { + return s.newTaskStatusConflictError(taskID, common.CREATED, status) + } + return nil +} + +func (s *IngestionTaskService) rollbackCreatedTask(taskID string) error { + _, err := s.ingestionTaskDAO.Delete(taskID, nil) + return err +} + +func (s *IngestionTaskService) enqueueTask(taskID string) error { + taskMessage := common.TaskMessage{ + TaskID: taskID, + TaskType: common.TaskTypeIngestionTask, + } + return s.taskPublisher.PublishTaskMessage("tasks.RAGFLOW", taskMessage) +} diff --git a/internal/service/ingestion_task_service_test.go b/internal/service/ingestion_task_service_test.go new file mode 100644 index 0000000000..498e2da3d5 --- /dev/null +++ b/internal/service/ingestion_task_service_test.go @@ -0,0 +1,465 @@ +package service + +import ( + "errors" + "testing" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" +) + +type recordingTaskPublisher struct { + subject string + messages []common.TaskMessage + err error +} + +func (p *recordingTaskPublisher) PublishTaskMessage(subject string, msg common.TaskMessage) error { + p.subject = subject + p.messages = append(p.messages, msg) + return p.err +} + +func TestIngestionTaskServiceCreateForDocumentsPublishesTaskMessages(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0) + insertTestDoc(t, "doc-1", "kb-1", 0, 0) + + publisher := &recordingTaskPublisher{} + svc := NewIngestionTaskService() + svc.taskPublisher = publisher + + resp, err := svc.CreateForDocuments("kb-1", "user-1", []string{"doc-1"}) + if err != nil { + t.Fatalf("CreateForDocuments failed: %v", err) + } + if len(resp) != 1 { + t.Fatalf("expected 1 response, got %d", len(resp)) + } + if publisher.subject != "tasks.RAGFLOW" { + t.Fatalf("subject = %q, want %q", publisher.subject, "tasks.RAGFLOW") + } + if len(publisher.messages) != 1 { + t.Fatalf("expected 1 published message, got %d", len(publisher.messages)) + } + msg := publisher.messages[0] + if msg.TaskType != common.TaskTypeIngestionTask { + t.Fatalf("task type = %q, want %q", msg.TaskType, common.TaskTypeIngestionTask) + } + task, err := dao.NewIngestionTaskDAO().GetByID(msg.TaskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.DocumentID != "doc-1" || task.DatasetID != "kb-1" || task.UserID != "user-1" { + t.Fatalf("unexpected task: %+v", task) + } +} + +func TestIngestionTaskServiceListByUserFiltersDataset(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + insertTestIngestionTask(t, "task-2", "user-1", "doc-2", "kb-2") + insertTestIngestionTask(t, "task-3", "user-2", "doc-3", "kb-1") + + svc := NewIngestionTaskService() + datasetID := "kb-1" + tasks, err := svc.ListByUser("user-1", &datasetID, 0, 0) + if err != nil { + t.Fatalf("ListByUser failed: %v", err) + } + if len(tasks) != 1 { + t.Fatalf("expected 1 task, got %d", len(tasks)) + } + if tasks[0].ID != "task-1" { + t.Fatalf("task ID = %q, want %q", tasks[0].ID, "task-1") + } +} + +func TestIngestionTaskServiceRequestStopManyStopsOwnedTasks(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + userID := "user-1" + svc := NewIngestionTaskService() + tasks, err := svc.RequestStopMany([]string{"task-1"}, &userID) + if err != nil { + t.Fatalf("RequestStopMany failed: %v", err) + } + if len(tasks) != 1 { + t.Fatalf("expected 1 task response, got %d", len(tasks)) + } + if tasks[0].Status != common.STOPPED { + t.Fatalf("status = %q, want %q", tasks[0].Status, common.STOPPED) + } +} + +func TestIngestionTaskServiceRequestStopManyRejectsOtherUsersTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + userID := "user-2" + svc := NewIngestionTaskService() + if _, err := svc.RequestStopMany([]string{"task-1"}, &userID); err == nil { + t.Fatal("expected RequestStopMany to reject non-owner") + } + task, err := dao.NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.CREATED { + t.Fatalf("status = %q, want %q", task.Status, common.CREATED) + } +} + +func TestIngestionTaskServiceRequestStopManyAllowsAdmin(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + tasks, err := svc.RequestStopMany([]string{"task-1"}, nil) + if err != nil { + t.Fatalf("RequestStopMany admin failed: %v", err) + } + if len(tasks) != 1 { + t.Fatalf("expected 1 task response, got %d", len(tasks)) + } + if tasks[0].Status != common.STOPPED { + t.Fatalf("status = %q, want %q", tasks[0].Status, common.STOPPED) + } +} + +func TestIngestionTaskServiceRemoveManyRemovesOwnedTasks(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + userID := "user-1" + svc := NewIngestionTaskService() + result, err := svc.RemoveMany([]string{"task-1"}, &userID) + if err != nil { + t.Fatalf("RemoveMany failed: %v", err) + } + if len(result) != 1 || result[0]["remove"] != "success" { + t.Fatalf("unexpected remove result: %+v", result) + } + if _, err := dao.NewIngestionTaskDAO().GetByID("task-1"); err == nil { + t.Fatal("task should be removed") + } +} + +func TestIngestionTaskServiceListAllForAdminIncludesStepAndUserEmail(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + status := "1" + if err := dao.DB.Create(&entity.User{ + ID: "user-1", + Email: "user-1@test.com", + Nickname: "user-1", + IsAuthenticated: "1", + IsActive: "1", + IsAnonymous: "0", + Status: &status, + }).Error; err != nil { + t.Fatalf("insert user: %v", err) + } + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + if err := dao.DB.Create(&entity.IngestionTaskLog{ + TaskID: "task-1", + Checkpoint: entity.JSONMap{"current_step": 3}, + }).Error; err != nil { + t.Fatalf("insert task log: %v", err) + } + + svc := NewIngestionTaskService() + tasks, err := svc.ListAllForAdmin() + if err != nil { + t.Fatalf("ListAllForAdmin failed: %v", err) + } + if len(tasks) != 1 { + t.Fatalf("expected 1 task, got %d", len(tasks)) + } + if tasks[0]["user"] != "user-1@test.com" { + t.Fatalf("user = %v, want user-1@test.com", tasks[0]["user"]) + } + if tasks[0]["step"] != 3 { + t.Fatalf("step = %v, want 3", tasks[0]["step"]) + } +} + +func TestIngestionTaskServiceStartRunningTransitionsCreatedTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + task, err := svc.StartRunning("task-1") + if err != nil { + t.Fatalf("StartRunning failed: %v", err) + } + if task.Status != common.RUNNING { + t.Fatalf("status = %q, want %q", task.Status, common.RUNNING) + } +} + +func TestIngestionTaskServiceRequestStopTransitionsCreatedTaskToStopped(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + task, err := svc.RequestStop("task-1") + if err != nil { + t.Fatalf("RequestStop failed: %v", err) + } + if task.Status != common.STOPPED { + t.Fatalf("status = %q, want %q", task.Status, common.STOPPED) + } +} + +func TestIngestionTaskServiceMarkCompletedRejectsNonRunningTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + if err := svc.MarkCompleted("task-1"); err == nil { + t.Fatal("expected MarkCompleted to reject non-running task") + } + task, err := dao.NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.CREATED { + t.Fatalf("status = %q, want %q", task.Status, common.CREATED) + } +} + +func TestIngestionTaskServiceMarkCompletedUpdatesTaskStatus(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", common.RUNNING).Error; err != nil { + t.Fatalf("set running status: %v", err) + } + + svc := NewIngestionTaskService() + if err := svc.MarkCompleted("task-1"); err != nil { + t.Fatalf("MarkCompleted failed: %v", err) + } + task, err := dao.NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.COMPLETED { + t.Fatalf("status = %q, want %q", task.Status, common.COMPLETED) + } +} + +func TestIngestionTaskServiceMarkFailedUpdatesTaskStatus(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", common.RUNNING).Error; err != nil { + t.Fatalf("set running status: %v", err) + } + + svc := NewIngestionTaskService() + if err := svc.MarkFailed("task-1"); err != nil { + t.Fatalf("MarkFailed failed: %v", err) + } + task, err := dao.NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.FAILED { + t.Fatalf("status = %q, want %q", task.Status, common.FAILED) + } +} + +func TestIngestionTaskServiceNewTaskStatusConflictErrorLoadsActualStatus(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", common.STOPPING).Error; err != nil { + t.Fatalf("set stopping status: %v", err) + } + + svc := NewIngestionTaskService() + err := svc.newTaskStatusConflictError("task-1", common.CREATED, common.RUNNING) + var conflictErr *TaskStatusConflictError + if !errors.As(err, &conflictErr) { + t.Fatalf("expected TaskStatusConflictError, got %T", err) + } + if conflictErr.TaskID != "task-1" || conflictErr.ExpectedFrom != common.CREATED || conflictErr.AttemptedTo != common.RUNNING || conflictErr.ActualCurrent != common.STOPPING { + t.Fatalf("unexpected conflict error: %+v", conflictErr) + } +} + +func TestIngestionTaskServiceMarkCompletedReturnsTaskIDInTransitionError(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + err := svc.MarkCompleted("task-1") + var transitionErr *InvalidTaskTransitionError + if !errors.As(err, &transitionErr) { + t.Fatalf("expected InvalidTaskTransitionError, got %T", err) + } + if transitionErr.TaskID != "task-1" || transitionErr.From != common.CREATED || transitionErr.To != common.COMPLETED { + t.Fatalf("unexpected transition error: %+v", transitionErr) + } +} + +func TestIngestionTaskServiceCreateAndEnqueueRetriesTerminalTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + publisher := &recordingTaskPublisher{} + svc := NewIngestionTaskService() + svc.taskPublisher = publisher + + testCases := []struct { + name string + status string + }{ + {name: "failed", status: common.FAILED}, + {name: "stopped", status: common.STOPPED}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + publisher.subject = "" + publisher.messages = nil + if err := dao.DB.Where("id = ?", "task-1").Delete(&entity.IngestionTask{}).Error; err != nil { + t.Fatalf("clear task: %v", err) + } + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", tc.status).Error; err != nil { + t.Fatalf("set terminal status: %v", err) + } + + task, err := svc.CreateAndEnqueue(&entity.IngestionTask{ + DocumentID: "doc-1", + UserID: "user-1", + DatasetID: "kb-1", + Status: common.CREATED, + }) + if err != nil { + t.Fatalf("CreateAndEnqueue failed: %v", err) + } + if task.ID != "task-1" { + t.Fatalf("task ID = %q, want task-1", task.ID) + } + if task.Status != common.CREATED { + t.Fatalf("status = %q, want %q", task.Status, common.CREATED) + } + if len(publisher.messages) != 1 || publisher.messages[0].TaskID != "task-1" { + t.Fatalf("unexpected published messages: %+v", publisher.messages) + } + reloaded, err := dao.NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("reload task: %v", err) + } + if reloaded.Status != common.CREATED { + t.Fatalf("reloaded status = %q, want %q", reloaded.Status, common.CREATED) + } + }) + } +} + +func TestIngestionTaskServiceCreateAndEnqueueRejectsActiveExistingTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + publisher := &recordingTaskPublisher{} + svc := NewIngestionTaskService() + svc.taskPublisher = publisher + + _, err := svc.CreateAndEnqueue(&entity.IngestionTask{DocumentID: "doc-1", UserID: "user-1", DatasetID: "kb-1", Status: common.CREATED}) + if err == nil { + t.Fatal("expected CreateAndEnqueue to reject existing created task") + } + if len(publisher.messages) != 0 { + t.Fatalf("expected no published messages, got %+v", publisher.messages) + } +} + +func TestIngestionTaskServiceCreateAndEnqueueRollsBackNewTaskOnPublishFailure(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + publisher := &recordingTaskPublisher{err: errors.New("publish failed")} + svc := NewIngestionTaskService() + svc.taskPublisher = publisher + + _, err := svc.CreateAndEnqueue(&entity.IngestionTask{ + DocumentID: "doc-1", + UserID: "user-1", + DatasetID: "kb-1", + Status: common.CREATED, + }) + if err == nil || err.Error() != "publish failed" { + t.Fatalf("expected publish failure, got %v", err) + } + task, getErr := dao.NewIngestionTaskDAO().GetByDocumentID("doc-1") + if getErr != nil { + t.Fatalf("reload task by document id: %v", getErr) + } + if task != nil { + t.Fatalf("expected created task to be deleted after publish failure, got %+v", task) + } +} + +func TestIngestionTaskServiceCreateAndEnqueueRollsBackRetriedTaskOnPublishFailure(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", common.FAILED).Error; err != nil { + t.Fatalf("set failed status: %v", err) + } + + publisher := &recordingTaskPublisher{err: errors.New("publish failed")} + svc := NewIngestionTaskService() + svc.taskPublisher = publisher + + _, err := svc.CreateAndEnqueue(&entity.IngestionTask{ + DocumentID: "doc-1", + UserID: "user-1", + DatasetID: "kb-1", + Status: common.CREATED, + }) + if err == nil || err.Error() != "publish failed" { + t.Fatalf("expected publish failure, got %v", err) + } + reloaded, getErr := dao.NewIngestionTaskDAO().GetByID("task-1") + if getErr != nil { + t.Fatalf("reload task: %v", getErr) + } + if reloaded.Status != common.FAILED { + t.Fatalf("status = %q, want %q", reloaded.Status, common.FAILED) + } +} + +func TestIngestionTaskServiceRemoveDeletesOwnedTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + userID := "user-1" + svc := NewIngestionTaskService() + info, err := svc.Remove("task-1", &userID) + if err != nil { + t.Fatalf("Remove failed: %v", err) + } + if info == nil || info.TaskID != "task-1" { + t.Fatalf("unexpected task info: %+v", info) + } + if _, err := dao.NewIngestionTaskDAO().GetByID("task-1"); err == nil { + t.Fatal("task should be removed") + } +} diff --git a/internal/service/ingestion_task_transition_test.go b/internal/service/ingestion_task_transition_test.go new file mode 100644 index 0000000000..2b05d0fdeb --- /dev/null +++ b/internal/service/ingestion_task_transition_test.go @@ -0,0 +1,55 @@ +package service + +import ( + "errors" + "testing" + + "ragflow/internal/common" +) + +func TestValidateTransitionAllowsExpectedEdges(t *testing.T) { + testCases := []struct { + from string + to string + }{ + {from: common.CREATED, to: common.RUNNING}, + {from: common.CREATED, to: common.STOPPED}, + {from: common.RUNNING, to: common.STOPPING}, + {from: common.RUNNING, to: common.COMPLETED}, + {from: common.RUNNING, to: common.FAILED}, + {from: common.STOPPING, to: common.STOPPED}, + {from: common.FAILED, to: common.CREATED}, + {from: common.STOPPED, to: common.CREATED}, + } + + for _, tc := range testCases { + if err := validateTransition(tc.from, tc.to); err != nil { + t.Fatalf("validateTransition(%q, %q) failed: %v", tc.from, tc.to, err) + } + } +} + +func TestValidateTransitionRejectsInvalidEdge(t *testing.T) { + err := validateTransition(common.CREATED, common.COMPLETED) + if err == nil { + t.Fatal("expected invalid transition to be rejected") + } + var transitionErr *InvalidTaskTransitionError + if !errors.As(err, &transitionErr) { + t.Fatalf("expected InvalidTaskTransitionError, got %T", err) + } + if transitionErr.TaskID != "" || transitionErr.From != common.CREATED || transitionErr.To != common.COMPLETED { + t.Fatalf("unexpected transition error: %+v", transitionErr) + } +} + +func TestTaskStatusConflictErrorSupportsErrorsAs(t *testing.T) { + err := &TaskStatusConflictError{TaskID: "task-1", ExpectedFrom: common.CREATED, AttemptedTo: common.RUNNING, ActualCurrent: common.STOPPING} + var conflictErr *TaskStatusConflictError + if !errors.As(err, &conflictErr) { + t.Fatalf("expected TaskStatusConflictError, got %T", err) + } + if conflictErr.TaskID != "task-1" || conflictErr.ExpectedFrom != common.CREATED || conflictErr.AttemptedTo != common.RUNNING || conflictErr.ActualCurrent != common.STOPPING { + t.Fatalf("unexpected conflict error: %+v", conflictErr) + } +} diff --git a/internal/service/task_publisher.go b/internal/service/task_publisher.go new file mode 100644 index 0000000000..b8d14753f3 --- /dev/null +++ b/internal/service/task_publisher.go @@ -0,0 +1,31 @@ +package service + +import ( + "encoding/json" + "fmt" + + "ragflow/internal/common" + "ragflow/internal/engine" +) + +type TaskPublisher interface { + PublishTaskMessage(subject string, msg common.TaskMessage) error +} + +type MessageQueueTaskPublisher struct{} + +func NewMessageQueueTaskPublisher() *MessageQueueTaskPublisher { + return &MessageQueueTaskPublisher{} +} + +func (p *MessageQueueTaskPublisher) PublishTaskMessage(subject string, msg common.TaskMessage) error { + msgQueueEngine := engine.GetMessageQueueEngine() + if msgQueueEngine == nil { + return fmt.Errorf("message queue engine is not initialized") + } + payload, err := json.Marshal(msg) + if err != nil { + return err + } + return msgQueueEngine.PublishTask(subject, payload) +}