zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-10 13:45:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(go-file): KB counter drift when deleting files with linked documents (#16584)

Browse Source 
## 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.
This commit is contained in:
Hz_
2026-07-03 14:07:54 +08:00
committed by GitHub
parent ee942711c4
commit ac5d0c4615
3 changed files with 107 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/service/document.go
View File

@@ -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 {

49
internal/service/file.go
View File

@@ -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_<tenant_id>_<kb_id>
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 {

105
internal/service/file_test.go
View File

@@ -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) {