zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-03 06:17:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: add context (#17354)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-24 20:19:41 +08:00
committed by GitHub
parent 64dbc518aa
commit d9e359d481
25 changed files with 438 additions and 351 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/service/document/document_artifact.go
View File

@@ -199,7 +199,7 @@ func (s *DocumentService) GetDocumentPreview(ctx context.Context, docID string)
return nil, err
}
bucket, name, err := s.GetDocumentStorageAddress(doc)
bucket, name, err := s.GetDocumentStorageAddress(ctx, doc)
if err != nil {
return nil, err
}

23
internal/service/document/document_crud.go
View File

@@ -29,7 +29,7 @@ func (s *DocumentService) Accessible(ctx context.Context, docID, userID string)
return s.kbDAO.Accessible(doc.KbID, userID)
}
func (s *DocumentService) GetDocumentStorageAddress(doc *entity.Document) (string, string, error) {
func (s *DocumentService) GetDocumentStorageAddress(ctx context.Context, doc *entity.Document) (string, string, error) {
if doc == nil {
return "", "", fmt.Errorf("document is nil")
}
@@ -43,7 +43,7 @@ func (s *DocumentService) GetDocumentStorageAddress(doc *entity.Document) (strin
}
if len(mappings) > 0 && mappings[0].FileID != nil {
file, err := fileDAO.GetByID(*mappings[0].FileID)
file, err := fileDAO.GetByID(ctx, dao.DB, *mappings[0].FileID)
if err != nil {
return "", "", err
}
@@ -70,7 +70,7 @@ func (s *DocumentService) DownloadDocument(ctx context.Context, datasetID, docID
if err != nil || doc.KbID != datasetID {
return nil, fmt.Errorf("Document not found!")
}
bucket, name, err := s.GetDocumentStorageAddress(doc)
bucket, name, err := s.GetDocumentStorageAddress(ctx, doc)
if err != nil {
return nil, err
}
@@ -257,7 +257,11 @@ func (s *DocumentService) deleteDocumentFull(ctx context.Context, docID string)
if err = s.deleteDocRecordWithCounters(ctx, doc, kb.ID); err != nil {
return err
}
s.cleanupFileReferences(docID)
cleanupCtx := context.WithoutCancel(ctx)
if err = s.cleanupFileReferences(cleanupCtx, docID); err != nil {
return fmt.Errorf("document deleted but file cleanup failed: %w", err)
}
return nil
}
@@ -395,13 +399,14 @@ func (s *DocumentService) rollbackAddFileFromKBError(ctx context.Context, doc *e
// the file is a knowledgebase-owned upload (source_type == knowledgebase) and
// no other document still references the same file_id. Files linked from file
// management are only unlinked — the file record and blob stay intact.
func (s *DocumentService) cleanupFileReferences(docID string) {
func (s *DocumentService) cleanupFileReferences(ctx context.Context, docID string) error {
mappings, mapErr := s.file2DocumentDAO.GetByDocumentID(docID)
if mapErr != nil {
common.Logger.Warn(fmt.Sprintf("cleanupFileReferences: failed to get f2d mappings for %s: %v", docID, mapErr))
return mapErr
}
if len(mappings) == 0 {
return
return nil
}
// Collect unique file_ids
@@ -418,6 +423,7 @@ func (s *DocumentService) cleanupFileReferences(docID string) {
// Delete all file2document rows for this document
if delErr := s.file2DocumentDAO.DeleteByDocumentID(docID); delErr != nil {
common.Logger.Warn(fmt.Sprintf("cleanupFileReferences: failed to delete f2d for %s: %v", docID, delErr))
return delErr
}
// For each file, only delete the record and blob when it is a
@@ -433,7 +439,7 @@ func (s *DocumentService) cleanupFileReferences(docID string) {
}
fileDAO := dao.NewFileDAO()
file, fErr := fileDAO.GetByID(fileID)
file, fErr := fileDAO.GetByID(ctx, dao.DB, fileID)
if fErr != nil || file == nil {
common.Logger.Warn(fmt.Sprintf("cleanupFileReferences: file not found %s: %v", fileID, fErr))
continue
@@ -441,7 +447,7 @@ func (s *DocumentService) cleanupFileReferences(docID string) {
if entity.FileSource(file.SourceType) != entity.FileSourceKnowledgebase {
continue // linked from file management — unlink only, keep the file
}
if _, delErr := fileDAO.DeleteByIDs([]string{fileID}); delErr != nil {
if _, delErr := fileDAO.DeleteByIDs(ctx, dao.DB, []string{fileID}); delErr != nil {
common.Logger.Warn(fmt.Sprintf("cleanupFileReferences: failed to delete file %s: %v", fileID, delErr))
continue // keep the blob so the live file row still has its object
}
@@ -454,4 +460,5 @@ func (s *DocumentService) cleanupFileReferences(docID string) {
}
}
}
return nil
}

5
internal/service/document/document_dataset_update.go
View File

@@ -349,7 +349,10 @@ func (s *DocumentService) updateDocumentNameOnly(ctx context.Context, doc *entit
mappings, err := s.file2DocumentDAO.GetByDocumentID(doc.ID)
if err == nil && len(mappings) > 0 && mappings[0].FileID != nil && s.fileDAO != nil {
_ = s.fileDAO.UpdateByID(*mappings[0].FileID, map[string]interface{}{"name": newName})
err = s.fileDAO.UpdateByID(ctx, dao.DB, *mappings[0].FileID, map[string]interface{}{"name": newName})
if err != nil {
return fmt.Errorf("file rename failed after document rename: %w", err)
}
}
if s.docEngine == nil {

2
internal/service/document/document_parse.go
View File

@@ -29,7 +29,7 @@ func (s *DocumentService) StartParseDocuments(ctx context.Context, doc *entity.D
// Validate storage first so we don't clear prior results and then fail
// because the document can't be read, leaving the document with neither
// old nor new parse results.
if _, _, err := s.GetDocumentStorageAddress(doc); err != nil {
if _, _, err := s.GetDocumentStorageAddress(ctx, doc); err != nil {
return err
}

21
internal/service/document/document_test.go
View File

@@ -664,7 +664,7 @@ func TestDeleteDocumentFull_CleansUpFile2Document(t *testing.T) {
}
// Verify file record deleted (hard delete)
files, _ := dao.NewFileDAO().GetByIDs([]string{"file-1"})
files, _ := dao.NewFileDAO().GetByIDs(ctx, db, []string{"file-1"})
if len(files) != 0 {
t.Fatalf("expected 0 files, got %d", len(files))
}
@@ -702,7 +702,7 @@ func TestDeleteDocumentFull_SharedFilePreserved(t *testing.T) {
}
// file record should still exist (doc-2 still references it)
files, _ := dao.NewFileDAO().GetByIDs([]string{"file-shared"})
files, _ := dao.NewFileDAO().GetByIDs(ctx, db, []string{"file-shared"})
if len(files) != 1 {
t.Fatalf("expected 1 file record to survive, got %d", len(files))
}
@@ -1478,7 +1478,8 @@ func TestCleanupFileReferences_NoMappings(t *testing.T) {
svc := testDocumentService(t)
// Should not panic with no f2d mappings
svc.cleanupFileReferences("no-mappings")
ctx := t.Context()
svc.cleanupFileReferences(ctx, "no-mappings")
}
func TestCleanupFileReferences_SingleFileDeleted(t *testing.T) {
@@ -1490,7 +1491,8 @@ func TestCleanupFileReferences_SingleFileDeleted(t *testing.T) {
insertTestFile2Document(t, "f2d-1", "file-1", "doc-1")
svc := testDocumentService(t)
svc.cleanupFileReferences("doc-1")
ctx := t.Context()
svc.cleanupFileReferences(ctx, "doc-1")
// f2d gone
mappings, _ := dao.NewFile2DocumentDAO().GetByDocumentID("doc-1")
@@ -1498,7 +1500,7 @@ func TestCleanupFileReferences_SingleFileDeleted(t *testing.T) {
t.Fatalf("expected 0 f2d after cleanup, got %d", len(mappings))
}
// file record gone
files, _ := dao.NewFileDAO().GetByIDs([]string{"file-1"})
files, _ := dao.NewFileDAO().GetByIDs(ctx, db, []string{"file-1"})
if len(files) != 0 {
t.Fatalf("expected 0 files after cleanup, got %d", len(files))
}
@@ -1514,7 +1516,8 @@ func TestCleanupFileReferences_SharedFileSurvives(t *testing.T) {
insertTestFile2Document(t, "f2d-2", "file-shared", "doc-2")
svc := testDocumentService(t)
svc.cleanupFileReferences("doc-1")
ctx := t.Context()
svc.cleanupFileReferences(ctx, "doc-1")
// f2d for doc-1 gone
mappings, _ := dao.NewFile2DocumentDAO().GetByDocumentID("doc-1")
@@ -1522,7 +1525,7 @@ func TestCleanupFileReferences_SharedFileSurvives(t *testing.T) {
t.Fatalf("expected 0 f2d for doc-1, got %d", len(mappings))
}
// file record survives
files, _ := dao.NewFileDAO().GetByIDs([]string{"file-shared"})
files, _ := dao.NewFileDAO().GetByIDs(ctx, db, []string{"file-shared"})
if len(files) != 1 {
t.Fatalf("expected 1 file record, got %d", len(files))
}
@@ -1715,7 +1718,7 @@ func TestUpdateDatasetDocumentRenameUpdatesDocumentAndFile(t *testing.T) {
if doc.Name == nil || *doc.Name != newName {
t.Fatalf("document name = %v, want %q", doc.Name, newName)
}
file, _ := dao.NewFileDAO().GetByID("file-1")
file, _ := dao.NewFileDAO().GetByID(ctx, db, "file-1")
if file.Name != newName {
t.Fatalf("file name = %q, want %q", file.Name, newName)
}
@@ -2887,7 +2890,7 @@ func TestFileDeleteRemovesLinkedDocument(t *testing.T) {
docSvc := testDocumentService(t)
fileSvc := file.NewFileService(
func(_ *dao.FileDAO, _ *entity.File, _ string) bool { return true },
func(_ context.Context, _ *dao.FileDAO, _ *entity.File, _ string) bool { return true },
docSvc,
)

36
internal/service/document/document_upload.go
View File

@@ -36,7 +36,7 @@ func (s *DocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.K
// Resolve (and create if needed) the dataset's file-manager folder up front.
// Without the File / file2document linkage the document list (which inner-joins
// file2document + file) would never surface the uploaded files.
kbFolder, err := s.ensureKBFolder(kb, tenantID)
kbFolder, err := s.ensureKBFolder(ctx, kb, tenantID)
if err != nil {
return nil, []string{err.Error()}
}
@@ -101,7 +101,7 @@ func (s *DocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.K
errMsgs = append(errMsgs, fh.Filename+": "+err.Error())
continue
}
if err = s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil {
if err = s.addFileFromKB(ctx, 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.
err = s.rollbackAddFileFromKBError(ctx, doc, kb.ID, err)
@@ -131,7 +131,7 @@ func (s *DocumentService) UploadEmptyDocument(ctx context.Context, kb *entity.Kn
}
}
kbFolder, err := s.ensureKBFolder(kb, tenantID)
kbFolder, err := s.ensureKBFolder(ctx, kb, tenantID)
if err != nil {
return nil, common.CodeServerError, err
}
@@ -140,7 +140,7 @@ func (s *DocumentService) UploadEmptyDocument(ctx context.Context, kb *entity.Kn
if err = s.InsertDocument(doc); err != nil {
return nil, common.CodeServerError, err
}
if err = s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil {
if err = s.addFileFromKB(ctx, doc, kbFolder.ID, kb.TenantID); err != nil {
return nil, common.CodeServerError, s.rollbackAddFileFromKBError(ctx, doc, kb.ID, err)
}
return docToRawMap(doc), common.CodeSuccess, nil
@@ -149,22 +149,26 @@ func (s *DocumentService) UploadEmptyDocument(ctx context.Context, kb *entity.Kn
// ensureKBFolder resolves (creating as needed) the per-dataset file-manager
// folder: root -> .knowledgebase -> <dataset name>. Mirrors Python
// get_root_folder + get_kb_folder + new_a_file_from_kb.
func (s *DocumentService) ensureKBFolder(kb *entity.Knowledgebase, tenantID string) (*entity.File, error) {
root, err := s.fileDAO.GetRootFolder(tenantID)
func (s *DocumentService) ensureKBFolder(ctx context.Context, kb *entity.Knowledgebase, tenantID string) (*entity.File, error) {
root, err := s.fileDAO.GetRootFolder(ctx, dao.DB, tenantID)
if err != nil {
return nil, err
}
kbRoot, err := s.newAFileFromKB(tenantID, knowledgebaseFolderName, root.ID)
kbRoot, err := s.newAFileFromKB(ctx, tenantID, knowledgebaseFolderName, root.ID)
if err != nil {
return nil, err
}
return s.newAFileFromKB(kb.TenantID, kb.Name, kbRoot.ID)
return s.newAFileFromKB(ctx, kb.TenantID, kb.Name, kbRoot.ID)
}
// newAFileFromKB returns the existing folder named name under parentID, or
// creates it. Mirrors Python FileService.new_a_file_from_kb.
func (s *DocumentService) newAFileFromKB(tenantID, name, parentID string) (*entity.File, error) {
for _, f := range s.fileDAO.Query(name, parentID, tenantID) {
func (s *DocumentService) newAFileFromKB(ctx context.Context, tenantID, name, parentID string) (*entity.File, error) {
existingFolders, err := s.fileDAO.Query(ctx, dao.DB, name, parentID, tenantID)
if err != nil {
return nil, err
}
for _, f := range existingFolders {
if f.TenantID == tenantID {
return f, nil
}
@@ -181,7 +185,7 @@ func (s *DocumentService) newAFileFromKB(tenantID, name, parentID string) (*enti
Location: &loc,
SourceType: string(entity.FileSourceKnowledgebase),
}
if err := s.fileDAO.Create(folder); err != nil {
if err := s.fileDAO.Create(ctx, dao.DB, folder); err != nil {
return nil, err
}
return folder, nil
@@ -190,7 +194,7 @@ func (s *DocumentService) newAFileFromKB(tenantID, name, parentID string) (*enti
// addFileFromKB links a document into the file manager: a File row under the
// dataset folder plus a file2document mapping. Mirrors Python
// FileService.add_file_from_kb (idempotent on the document mapping).
func (s *DocumentService) addFileFromKB(doc *entity.Document, kbFolderID, tenantID string) error {
func (s *DocumentService) addFileFromKB(ctx context.Context, doc *entity.Document, kbFolderID, tenantID string) error {
if existing, err := s.file2DocumentDAO.GetByDocumentID(doc.ID); err == nil && len(existing) > 0 {
return nil
}
@@ -214,7 +218,7 @@ func (s *DocumentService) addFileFromKB(doc *entity.Document, kbFolderID, tenant
Location: &loc,
SourceType: string(entity.FileSourceKnowledgebase),
}
if err := s.fileDAO.Create(file); err != nil {
if err := s.fileDAO.Create(ctx, dao.DB, file); err != nil {
return err
}
docID := doc.ID
@@ -223,7 +227,7 @@ func (s *DocumentService) addFileFromKB(doc *entity.Document, kbFolderID, tenant
FileID: &fileID,
DocumentID: &docID,
}); err != nil {
_ = s.fileDAO.Delete(fileID)
_ = s.fileDAO.Delete(ctx, dao.DB, fileID)
return err
}
return nil
@@ -235,7 +239,7 @@ func (s *DocumentService) UploadWebDocument(ctx context.Context, kb *entity.Know
return nil, common.CodeServerError, fmt.Errorf("storage not initialized")
}
kbFolder, err := s.ensureKBFolder(kb, tenantID)
kbFolder, err := s.ensureKBFolder(ctx, kb, tenantID)
if err != nil {
return nil, common.CodeServerError, err
}
@@ -279,7 +283,7 @@ func (s *DocumentService) UploadWebDocument(ctx context.Context, kb *entity.Know
_ = storageImpl.Remove(kb.ID, location)
return nil, common.CodeServerError, err
}
if err = s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil {
if err = s.addFileFromKB(ctx, doc, kbFolder.ID, kb.TenantID); err != nil {
err = s.rollbackAddFileFromKBError(ctx, doc, kb.ID, err)
_ = storageImpl.Remove(kb.ID, location)
return nil, common.CodeServerError, err

21
internal/service/document/file2document.go
View File

@@ -84,7 +84,7 @@ type LinkToDatasetsRequest struct {
// handler can map it to a Python-compatible response without leaking internals.
func (s *File2DocumentService) LinkToDatasets(ctx context.Context, userID string, req *LinkToDatasetsRequest, mode string) error {
// ── 1. Validate files exist ───────────────────────────────────────────────
files, err := s.fileDAO.GetByIDs(req.FileIDs)
files, err := s.fileDAO.GetByIDs(ctx, dao.DB, req.FileIDs)
if err != nil {
common.Warn("LinkToDatasets: GetByIDs failed", zap.Error(err))
return ErrLinkInternal
@@ -117,7 +117,7 @@ func (s *File2DocumentService) LinkToDatasets(ctx context.Context, userID string
for _, id := range req.FileIDs {
file := filesSet[id]
if file.Type == "folder" {
inner, err := s.getAllInnermostFileIDs(id)
inner, err := s.getAllInnermostFileIDs(ctx, id)
if err != nil {
common.Warn("LinkToDatasets: folder expansion failed", zap.String("fileID", id), zap.Error(err))
return ErrLinkInternal
@@ -131,11 +131,11 @@ func (s *File2DocumentService) LinkToDatasets(ctx context.Context, userID string
// ── 4. Validate expanded file permissions ─────────────────────────────────
for _, id := range allFileIDs {
file, err := s.fileDAO.GetByID(id)
file, err := s.fileDAO.GetByID(ctx, dao.DB, id)
if err != nil || file == nil {
return ErrLinkFileNotFound
}
if !service.CheckFileTeamPermission(s.fileDAO, file, userID) {
if !service.CheckFileTeamPermission(ctx, s.fileDAO, file, userID) {
return ErrLinkNoAuthorization
}
}
@@ -150,7 +150,8 @@ func (s *File2DocumentService) LinkToDatasets(ctx context.Context, userID string
// ── 6. Run conversion in background (fire-and-forget) ────────────────────
kbIDs := req.KbIDs
go func() {
if err = s.convertFiles(ctx, allFileIDs, kbIDs, userID, mode); err != nil {
newCtx := context.Background()
if err = s.convertFiles(newCtx, allFileIDs, kbIDs, userID, mode); err != nil {
common.Warn("file2document.convertFiles failed",
zap.Strings("file_ids", allFileIDs),
zap.Strings("kb_ids", kbIDs),
@@ -211,7 +212,7 @@ func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs
}
// Reload the source file.
file, err := s.fileDAO.GetByID(fileID)
file, err := s.fileDAO.GetByID(ctx, dao.DB, fileID)
if err != nil || file == nil {
continue
}
@@ -259,7 +260,7 @@ func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs
// InsertDocument creates the row and increments KB doc_num in one
// transaction, so a failed insert never leaves a stale counter.
if err := s.documentSvc.InsertDocument(doc); err != nil {
if err = s.documentSvc.InsertDocument(doc); err != nil {
common.Warn("convertFiles: InsertDocument failed",
zap.String("kbID", kbID), zap.String("fileID", fileID), zap.Error(err))
continue
@@ -281,15 +282,15 @@ func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs
// getAllInnermostFileIDs recursively collects all non-folder file IDs under a folder.
// Mirrors Python FileService.get_all_innermost_file_ids.
func (s *File2DocumentService) getAllInnermostFileIDs(folderID string) ([]string, error) {
children, err := s.fileDAO.ListByParentID(folderID)
func (s *File2DocumentService) getAllInnermostFileIDs(ctx context.Context, folderID string) ([]string, error) {
children, err := s.fileDAO.ListByParentID(ctx, dao.DB, folderID)
if err != nil {
return nil, err
}
var ids []string
for _, child := range children {
if child.Type == "folder" {
sub, err := s.getAllInnermostFileIDs(child.ID)
sub, err := s.getAllInnermostFileIDs(ctx, child.ID)
if err != nil {
return nil, err
}