mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-12 11:43:39 +08:00
Feat: ingestion cleanup (#16953)
### Summary 1. Remove dead code (replaced by builtin ingestion pipeline) 2. Refactor (move document parsing progress from http api into ingestion executor)
This commit is contained in:
@@ -194,6 +194,21 @@ func (dao *IngestionTaskDAO) GetByDocumentID(documentId string) (*entity.Ingesti
|
||||
return tasks[0], nil
|
||||
}
|
||||
|
||||
// DeleteIfTerminal deletes ingestion tasks for a document that are in a
|
||||
// terminal state (COMPLETED, STOPPED, FAILED) or still queued (CREATED).
|
||||
// RUNNING and STOPPING tasks are NOT deleted because an in-flight worker
|
||||
// would keep writing chunks and corrupt a new run's results.
|
||||
// Returns the number of rows deleted.
|
||||
func (dao *IngestionTaskDAO) DeleteIfTerminal(documentID string) (int64, error) {
|
||||
result := DB.Where("document_id = ? AND status NOT IN (?, ?)",
|
||||
documentID, common.RUNNING, common.STOPPING).
|
||||
Delete(&entity.IngestionTask{})
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
type IngestionTaskLogDAO struct{}
|
||||
|
||||
func NewIngestionTaskLogDAO() *IngestionTaskLogDAO {
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"fmt"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/entity"
|
||||
|
||||
@@ -133,3 +134,61 @@ func TestIngestionTaskDAOUpdateStatusIfCurrentRejectsMismatchedStatus(t *testing
|
||||
t.Fatalf("status = %q, want %q", reloaded.Status, common.STOPPING)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIngestionTaskDAODeleteIfTerminal_RemovesOnlyTerminal(t *testing.T) {
|
||||
db := setupTaskTestDB(t)
|
||||
orig := DB
|
||||
DB = db
|
||||
t.Cleanup(func() { DB = orig })
|
||||
|
||||
// Create tasks in different statuses, each with a unique docID.
|
||||
statuses := []string{common.CREATED, common.RUNNING, common.STOPPING, common.COMPLETED, common.STOPPED, common.FAILED}
|
||||
for i, status := range statuses {
|
||||
docID := fmt.Sprintf("doc-%d", i)
|
||||
task := &entity.IngestionTask{
|
||||
ID: fmt.Sprintf("task-%d", i),
|
||||
UserID: "user-1",
|
||||
DocumentID: docID,
|
||||
DatasetID: "kb-1",
|
||||
Status: status,
|
||||
}
|
||||
if err := db.Create(task).Error; err != nil {
|
||||
t.Fatalf("create task %s: %v", status, err)
|
||||
}
|
||||
}
|
||||
|
||||
// DeleteIfTerminal deletes everything except RUNNING and STOPPING.
|
||||
// CREATED is safe to delete (no worker has claimed it yet);
|
||||
// COMPLETED/STOPPED/FAILED are terminal.
|
||||
// Call it for every doc and verify the negative cases survived.
|
||||
for i := 0; i < len(statuses); i++ {
|
||||
docID := fmt.Sprintf("doc-%d", i)
|
||||
_, err := NewIngestionTaskDAO().DeleteIfTerminal(docID)
|
||||
if err != nil {
|
||||
t.Fatalf("DeleteIfTerminal(doc-%d): %v", i, err)
|
||||
}
|
||||
}
|
||||
|
||||
// RUNNING and STOPPING must survive.
|
||||
for _, i := range []int{1, 2} {
|
||||
docID := fmt.Sprintf("doc-%d", i)
|
||||
task, err := NewIngestionTaskDAO().GetByDocumentID(docID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetByDocumentID %s: %v", docID, err)
|
||||
}
|
||||
if task == nil {
|
||||
t.Fatalf("%s task (doc=%d) must not be deleted", statuses[i], i)
|
||||
}
|
||||
}
|
||||
// CREATED, COMPLETED, STOPPED, FAILED must be gone.
|
||||
for _, i := range []int{0, 3, 4, 5} {
|
||||
docID := fmt.Sprintf("doc-%d", i)
|
||||
task, err := NewIngestionTaskDAO().GetByDocumentID(docID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetByDocumentID %s: %v", docID, err)
|
||||
}
|
||||
if task != nil {
|
||||
t.Fatalf("%s task (doc=%d) should be deleted, still present", statuses[i], i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user