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

fix(go): document count in kb (#16490)

Browse Source 
### Summary
This PR fixes incorrect dataset document counters in the Go service.

Several document creation paths inserted document records directly
through documentDAO.Create, bypassing the shared InsertDocument logic
that increments knowledgebase.doc_num. As a result, datasets could
contain documents while doc_num remained 0.
This commit is contained in:
Hz_
2026-07-02 10:34:14 +08:00
committed by GitHub
parent 0de69e5bba
commit 9b83d0f154
2 changed files with 86 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
internal/service/document.go
View File

@@ -556,17 +556,19 @@ func (s *DocumentService) DownloadDocument(datasetID, docID string) (*DownloadDo
// CreateDocument create document
func (s *DocumentService) CreateDocument(req *CreateDocumentRequest) (*entity.Document, error) {
document := &entity.Document{
Name: &req.Name,
KbID: req.KbID,
ParserID: req.ParserID,
CreatedBy: req.CreatedBy,
Type: req.Type,
SourceType: req.Source,
Suffix: ".doc",
Status: func() *string { s := "0"; return &s }(),
ID: common.GenerateUUID(),
Name: &req.Name,
KbID: req.KbID,
ParserID: req.ParserID,
ParserConfig: entity.JSONMap{},
CreatedBy: req.CreatedBy,
Type: req.Type,
SourceType: req.Source,
Suffix: ".doc",
Status: func() *string { s := "0"; return &s }(),
}
if err := s.documentDAO.Create(document); err != nil {
if err := s.InsertDocument(document); err != nil {
return nil, fmt.Errorf("failed to create document: %w", err)
}
@@ -772,20 +774,30 @@ func (s *DocumentService) deleteDocRecordWithCounters(doc *entity.Document, kbID
return nil // already deleted by a concurrent request — skip counters
}
decErr := tx.Model(&entity.Knowledgebase{}).
result = tx.Model(&entity.Knowledgebase{}).
Where("id = ?", kbID).
Updates(map[string]interface{}{
"doc_num": gorm.Expr("doc_num - 1"),
"chunk_num": gorm.Expr("chunk_num - ?", doc.ChunkNum),
"token_num": gorm.Expr("token_num - ?", doc.TokenNum),
}).Error
if decErr != nil {
common.Logger.Warn(fmt.Sprintf("deleteDocRecordWithCounters: failed to decrement KB %s: %v", kbID, decErr))
})
if result.Error != nil {
return fmt.Errorf("failed to decrement counters for KB %s: %w", kbID, result.Error)
}
if result.RowsAffected == 0 {
return fmt.Errorf("knowledgebase %s not found", kbID)
}
return nil
})
}
func (s *DocumentService) rollbackAddFileFromKBError(doc *entity.Document, kbID string, err error) error {
if cleanupErr := s.deleteDocRecordWithCounters(doc, kbID); cleanupErr != nil {
return fmt.Errorf("%w; rollback cleanup failed: %w", err, cleanupErr)
}
return err
}
// cleanupFileReferences deletes file2document mappings for docID, and for each
// referenced file, only hard-deletes the file record and its storage blob when
// no other document still references the same file_id.
@@ -2993,7 +3005,7 @@ func (s *DocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantI
}
doc := s.newDatasetDocument(kb, tenantID, filename, location, string(filetype), merged, "local", int64(len(blob)), blob)
if err := s.documentDAO.Create(doc); err != nil {
if err := s.InsertDocument(doc); err != nil {
// Roll back the orphaned blob so a failed insert doesn't leak storage.
_ = storageImpl.Remove(kb.ID, location)
errMsgs = append(errMsgs, fh.Filename+": "+err.Error())
@@ -3002,7 +3014,7 @@ func (s *DocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantI
if err := s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil {
// Linkage failed: roll back the document row and blob so the partial
// state doesn't leave an invisible (unlisted) document behind.
_, _ = s.documentDAO.Delete(doc.ID)
err = s.rollbackAddFileFromKBError(doc, kb.ID, err)
_ = storageImpl.Remove(kb.ID, location)
errMsgs = append(errMsgs, fh.Filename+": "+err.Error())
continue
@@ -3035,12 +3047,11 @@ func (s *DocumentService) UploadEmptyDocument(kb *entity.Knowledgebase, tenantID
}
doc := s.newDatasetDocument(kb, tenantID, name, "", "virtual", kb.ParserConfig, "local", 0, nil)
if err := s.documentDAO.Create(doc); err != nil {
if err := s.InsertDocument(doc); err != nil {
return nil, common.CodeServerError, err
}
if err := s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil {
_, _ = s.documentDAO.Delete(doc.ID)
return nil, common.CodeServerError, err
return nil, common.CodeServerError, s.rollbackAddFileFromKBError(doc, kb.ID, err)
}
return docToRawMap(doc), common.CodeSuccess, nil
}
@@ -3178,12 +3189,12 @@ func (s *DocumentService) UploadWebDocument(kb *entity.Knowledgebase, tenantID,
}
doc := s.newDatasetDocument(kb, tenantID, filename, location, string(filetype), kb.ParserConfig, "web", int64(len(blob)), blob)
if err := s.documentDAO.Create(doc); err != nil {
if err := s.InsertDocument(doc); err != nil {
_ = storageImpl.Remove(kb.ID, location)
return nil, common.CodeServerError, err
}
if err := s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil {
_, _ = s.documentDAO.Delete(doc.ID)
err = s.rollbackAddFileFromKBError(doc, kb.ID, err)
_ = storageImpl.Remove(kb.ID, location)
return nil, common.CodeServerError, err
}

55
internal/service/document_test.go
View File

@@ -400,6 +400,17 @@ func insertTestKB(t *testing.T, id, tenantID string, docNum, tokenNum, chunkNum
}
}
func assertKBDocNum(t *testing.T, kbID string, want int64) {
t.Helper()
var got int64
if err := dao.DB.Model(&entity.Knowledgebase{}).Select("doc_num").Where("id = ?", kbID).Scan(&got).Error; err != nil {
t.Fatalf("get kb doc_num %s: %v", kbID, err)
}
if got != want {
t.Fatalf("kb %s doc_num=%d, want %d", kbID, got, want)
}
}
func insertTestDoc(t *testing.T, id, kbID string, tokenNum, chunkNum int64) {
t.Helper()
doc := &entity.Document{
@@ -458,6 +469,29 @@ func insertTestFile(t *testing.T, id, parentID, name string, location *string) {
}
}
func TestCreateDocumentIncrementsKBDocNum(t *testing.T) {
db := setupServiceTestDB(t)
pushServiceDB(t, db)
insertTestKB(t, "kb-create", "tenant-1", 0, 0, 0)
svc := testDocumentService(t)
doc, err := svc.CreateDocument(&CreateDocumentRequest{
Name: "created.txt",
KbID: "kb-create",
ParserID: "naive",
CreatedBy: "tenant-1",
Type: "doc",
Source: "local",
})
if err != nil {
t.Fatalf("CreateDocument failed: %v", err)
}
if doc == nil || doc.KbID != "kb-create" {
t.Fatalf("unexpected doc: %+v", doc)
}
assertKBDocNum(t, "kb-create", 1)
}
func TestDeleteDocumentFull_Basic(t *testing.T) {
db := setupServiceTestDB(t)
pushServiceDB(t, db)
@@ -644,6 +678,7 @@ func TestUploadLocalDocuments_MirrorsPythonCoreFields(t *testing.T) {
ParserConfig: entity.JSONMap{
"existing": "value",
},
DocNum: 1,
}
if err := dao.DB.Create(kb).Error; err != nil {
t.Fatalf("insert kb: %v", err)
@@ -695,6 +730,7 @@ func TestUploadLocalDocuments_MirrorsPythonCoreFields(t *testing.T) {
if string(storedBlob) != "abc" {
t.Fatalf("stored blob=%q, want abc", storedBlob)
}
assertKBDocNum(t, kb.ID, 2)
}
func TestUploadEmptyDocument_CreatesVirtualDocumentAndFileLink(t *testing.T) {
@@ -742,6 +778,7 @@ func TestUploadEmptyDocument_CreatesVirtualDocumentAndFileLink(t *testing.T) {
if linkCount != 1 {
t.Fatalf("link count=%d, want 1", linkCount)
}
assertKBDocNum(t, kb.ID, 1)
}
func insertUserTenantForAccessCheck(t *testing.T, userID, tenantID string) {
@@ -1254,6 +1291,24 @@ func TestDeleteDocRecordWithCounters_DocAlreadyDeleted(t *testing.T) {
}
}
func TestDeleteDocRecordWithCounters_KBUpdateFailureRollsBackDocumentDelete(t *testing.T) {
db := setupServiceTestDB(t)
pushServiceDB(t, db)
insertTestDoc(t, "doc-1", "missing-kb", 10, 5)
doc, _ := dao.NewDocumentDAO().GetByID("doc-1")
svc := testDocumentService(t)
if err := svc.deleteDocRecordWithCounters(doc, "missing-kb"); err == nil {
t.Fatal("expected missing KB counter update to return an error")
}
if _, err := dao.NewDocumentDAO().GetByID("doc-1"); err != nil {
t.Fatalf("expected document delete to roll back, got: %v", err)
}
}
func TestCleanupFileReferences_NoMappings(t *testing.T) {
db := setupServiceTestDB(t)
pushServiceDB(t, db)