From 9b83d0f154e2dbbc747121fe911922b9f27b3065 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Thu, 2 Jul 2026 10:34:14 +0800 Subject: [PATCH] fix(go): document count in kb (#16490) ### 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. --- internal/service/document.go | 51 +++++++++++++++++----------- internal/service/document_test.go | 55 +++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 20 deletions(-) diff --git a/internal/service/document.go b/internal/service/document.go index b62a94f53..50c0b56e3 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -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 } diff --git a/internal/service/document_test.go b/internal/service/document_test.go index 4b3d1bb78..e93510fb6 100644 --- a/internal/service/document_test.go +++ b/internal/service/document_test.go @@ -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)