From ac5d0c4615040b6d5e3c507bd9ceb96c2d273d1b Mon Sep 17 00:00:00 2001 From: Hz_ Date: Fri, 3 Jul 2026 14:07:54 +0800 Subject: [PATCH] fix(go-file): KB counter drift when deleting files with linked documents (#16584) ## Summary Use `DocumentService.RemoveDocumentKeepFile` when deleting files that are linked to documents. ## Change - inject `DocumentService` into `FileService` - replace direct document deletion in `deleteSingleFile` - remove the obsolete file-local engine deletion helper ## Result Deleting a file now cleans up linked documents through the same service path used elsewhere, keeping KB counters and document engine cleanup consistent. --- internal/service/document.go | 2 + internal/service/file.go | 49 ++-------------- internal/service/file_test.go | 105 ++++++++++++++++++++++++++++++++-- 3 files changed, 107 insertions(+), 49 deletions(-) diff --git a/internal/service/document.go b/internal/service/document.go index 5281cacc3..cb6b3a3f9 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -3718,6 +3718,7 @@ func (s *DocumentService) UploadDocumentInfos(userID string, files []*multipart. fileSvc := &FileService{ fileDAO: s.fileDAO, file2DocumentDAO: s.file2DocumentDAO, + documentService: s, } data, err := fileSvc.UploadInfos(userID, files) if err != nil { @@ -3730,6 +3731,7 @@ func (s *DocumentService) UploadDocumentInfoByURL(userID, rawURL string) (map[st fileSvc := &FileService{ fileDAO: s.fileDAO, file2DocumentDAO: s.file2DocumentDAO, + documentService: s, } data, err := fileSvc.UploadFromURL(userID, rawURL) if err != nil { diff --git a/internal/service/file.go b/internal/service/file.go index f9ad80817..4ecd3d3d5 100644 --- a/internal/service/file.go +++ b/internal/service/file.go @@ -30,7 +30,6 @@ import ( "path/filepath" "ragflow/internal/common" "ragflow/internal/dao" - "ragflow/internal/engine" "ragflow/internal/entity" "ragflow/internal/ingestion/parser" "ragflow/internal/storage" @@ -46,6 +45,7 @@ import ( type FileService struct { fileDAO *dao.FileDAO file2DocumentDAO *dao.File2DocumentDAO + documentService *DocumentService } // NewFileService create file service @@ -53,6 +53,7 @@ func NewFileService() *FileService { return &FileService{ fileDAO: dao.NewFileDAO(), file2DocumentDAO: dao.NewFile2DocumentDAO(), + documentService: NewDocumentService(), } } @@ -692,35 +693,14 @@ func (s *FileService) deleteSingleFile(ctx context.Context, file *entity.File) e return fmt.Errorf("failed to get file2document mappings: %w", err) } if len(informs) > 0 { - documentDAO := dao.NewDocumentDAO() - datasetDAO := dao.NewKnowledgebaseDAO() - for _, inform := range informs { if inform.DocumentID == nil { continue } docID := *inform.DocumentID - - doc, err := documentDAO.GetByID(docID) - if err == nil && doc != nil { - // Get tenant ID from KB - ds, err := datasetDAO.GetByID(doc.KbID) - if err == nil && ds != nil { - tenantID := ds.TenantID - if tenantID != "" { - // Delete from document engine - if err := s.deleteDocumentFromEngine(ctx, doc, tenantID); err != nil { - common.Logger.Error(fmt.Sprintf("Fail to delete document from engine: %s, error: %v", doc.ID, err)) - } - } - } - - // Delete document record - if _, err := documentDAO.Delete(docID); err != nil { - common.Logger.Error(fmt.Sprintf("Fail to delete document: %s, error: %v", docID, err)) - } + if err := s.documentService.RemoveDocumentKeepFile(docID); err != nil { + common.Logger.Error(fmt.Sprintf("Fail to remove document: %s, error: %v", docID, err)) } - } // Delete file2document mapping (outside the loop, called once - matching Python behavior) @@ -737,27 +717,6 @@ func (s *FileService) deleteSingleFile(ctx context.Context, file *entity.File) e return nil } -// deleteDocumentFromEngine deletes a document from the document engine -func (s *FileService) deleteDocumentFromEngine(ctx context.Context, doc *entity.Document, tenantID string) error { - // Get document engine - docEngine := engine.Get() - if docEngine == nil { - return nil - } - - // Build index name: ragflow__ - indexName := fmt.Sprintf("ragflow_%s_%s", tenantID, doc.KbID) - - // Delete document from engine with timeout - reqCtx, cancel := context.WithTimeout(ctx, 300*time.Second) - defer cancel() - condition := map[string]interface{}{"doc_id": doc.ID} - if _, err := docEngine.DeleteChunks(reqCtx, condition, indexName, doc.KbID); err != nil { - return fmt.Errorf("delete document from engine: %w", err) - } - return nil -} - // deleteFolderRecursive recursively deletes a folder and its contents // Matches Python's _delete_folder_recursive function func (s *FileService) deleteFolderRecursive(ctx context.Context, folder *entity.File, uid string) error { diff --git a/internal/service/file_test.go b/internal/service/file_test.go index db1ec167c..cacfd126c 100644 --- a/internal/service/file_test.go +++ b/internal/service/file_test.go @@ -2,6 +2,7 @@ package service import ( "bytes" + "context" "errors" "io" "net/http" @@ -10,6 +11,8 @@ import ( "testing" "time" + "ragflow/internal/dao" + "ragflow/internal/entity" "ragflow/internal/storage" ) @@ -22,6 +25,14 @@ type fakeStorage struct { exists bool } +func testFileService() *FileService { + return &FileService{ + fileDAO: dao.NewFileDAO(), + file2DocumentDAO: dao.NewFile2DocumentDAO(), + documentService: &DocumentService{}, + } +} + func (f *fakeStorage) Health() bool { return true } @@ -83,7 +94,7 @@ func TestFileService_DownloadAgentFile_Success(t *testing.T) { factory.SetStorage(originalStorage) }) - svc := NewFileService() + svc := testFileService() tenantID := "tenant123" location := "file-abc.txt" @@ -118,7 +129,7 @@ func TestFileService_DownloadAgentFile_Error(t *testing.T) { factory.SetStorage(originalStorage) }) - svc := NewFileService() + svc := testFileService() tenantID := "tenant123" location := "file-abc.txt" @@ -160,7 +171,7 @@ func TestFileService_UploadFromURL_PDFAddsExtensionAndStoresToDownloads(t *testi factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) - svc := NewFileService() + svc := testFileService() resp, err := svc.UploadFromURL("tenant123", server.URL+"/report") if err != nil { t.Fatalf("UploadFromURL failed: %v", err) @@ -206,7 +217,7 @@ func TestFileService_UploadFromURL_HTMLNormalizesReadableContent(t *testing.T) { factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) - svc := NewFileService() + svc := testFileService() resp, err := svc.UploadFromURL("tenant123", server.URL+"/page") if err != nil { t.Fatalf("UploadFromURL failed: %v", err) @@ -252,6 +263,92 @@ func TestReadUploadInfoData_RejectsOversizedInput(t *testing.T) { } } +func TestFileService_DeleteSingleFile_RemovesLinkedDocumentThroughDocumentService(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + + mockStorage := newFakeUploadStorage() + factory := storage.GetStorageFactory() + originalStorage := factory.GetStorage() + factory.SetStorage(mockStorage) + t.Cleanup(func() { factory.SetStorage(originalStorage) }) + + insertTestKB(t, "kb-file-delete", "tenant-1", 1, 30, 10) + insertTestDoc(t, "doc-file-delete", "kb-file-delete", 30, 10) + insertTestTask(t, "task-file-delete", "doc-file-delete") + + location := "doc.pdf" + insertTestFile(t, "file-delete", "folder-delete", "doc.pdf", &location) + insertTestFile2Document(t, "f2d-file-delete", "file-delete", "doc-file-delete") + if err := mockStorage.Put("folder-delete", location, []byte("blob")); err != nil { + t.Fatalf("put test blob: %v", err) + } + + file, err := dao.NewFileDAO().GetByID("file-delete") + if err != nil { + t.Fatalf("get file: %v", err) + } + docSvc := testDocumentService(t) + docEngine := &rerunDeleteDocEngine{} + docSvc.docEngine = docEngine + svc := &FileService{ + fileDAO: dao.NewFileDAO(), + file2DocumentDAO: dao.NewFile2DocumentDAO(), + documentService: docSvc, + } + + if err := svc.deleteSingleFile(context.Background(), file); err != nil { + t.Fatalf("deleteSingleFile failed: %v", err) + } + + if mockStorage.ObjExist("folder-delete", location) { + t.Fatal("expected storage object to be removed") + } + if _, err := dao.NewDocumentDAO().GetByID("doc-file-delete"); err == nil { + t.Fatal("expected linked document to be deleted") + } + var taskCount int64 + if err := dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-file-delete").Count(&taskCount).Error; err != nil { + t.Fatalf("count tasks: %v", err) + } + if taskCount != 0 { + t.Fatalf("expected linked tasks to be deleted, got %d", taskCount) + } + mappings, err := dao.NewFile2DocumentDAO().GetByFileID("file-delete") + if err != nil { + t.Fatalf("get file2document mappings: %v", err) + } + if len(mappings) != 0 { + t.Fatalf("expected file2document mappings to be deleted, got %d", len(mappings)) + } + files, err := dao.NewFileDAO().GetByIDs([]string{"file-delete"}) + if err != nil { + t.Fatalf("get deleted file: %v", err) + } + if len(files) != 0 { + t.Fatalf("expected file record to be deleted, got %d", len(files)) + } + if docEngine.deleteCalls != 1 { + t.Fatalf("expected document engine cleanup to be called once, got %d", docEngine.deleteCalls) + } + if docEngine.indexName != "ragflow_tenant-1" { + t.Fatalf("expected document engine index ragflow_tenant-1, got %q", docEngine.indexName) + } + if docEngine.datasetID != "kb-file-delete" { + t.Fatalf("expected document engine dataset kb-file-delete, got %q", docEngine.datasetID) + } + if docEngine.condition["doc_id"] != "doc-file-delete" { + t.Fatalf("expected document engine doc_id condition doc-file-delete, got %v", docEngine.condition["doc_id"]) + } + kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-file-delete") + if err != nil { + t.Fatalf("get kb: %v", err) + } + if kb.DocNum != 0 || kb.TokenNum != 0 || kb.ChunkNum != 0 { + t.Fatalf("expected KB counters to be decremented to zero, got doc=%d token=%d chunk=%d", kb.DocNum, kb.TokenNum, kb.ChunkNum) + } +} + type zeroReader struct{} func (zeroReader) Read(p []byte) (int, error) {