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

Go: add context, part3 (#17369)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-24 22:00:09 +08:00
committed by GitHub
parent 08332501a8
commit cc1eb6fb58
13 changed files with 65 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
internal/dao/file2document.go
View File

@@ -17,7 +17,10 @@
package dao
import (
"context"
"ragflow/internal/entity"
"gorm.io/gorm"
)
// File2DocumentDAO file to document mapping data access object
@@ -29,10 +32,10 @@ func NewFile2DocumentDAO() *File2DocumentDAO {
}
// GetKBInfoByFileID gets knowledge base info by file ID
func (dao *File2DocumentDAO) GetKBInfoByFileID(fileID string) ([]map[string]interface{}, error) {
func (dao *File2DocumentDAO) GetKBInfoByFileID(ctx context.Context, db *gorm.DB, fileID string) ([]map[string]interface{}, error) {
var results []map[string]interface{}
rows, err := DB.Model(&entity.File{}).
rows, err := db.WithContext(ctx).Model(&entity.File{}).
Select("knowledgebase.id, knowledgebase.name, file2document.document_id").
Joins("JOIN file2document ON file2document.file_id = ?", fileID).
Joins("JOIN document ON document.id = file2document.document_id").
@@ -46,7 +49,7 @@ func (dao *File2DocumentDAO) GetKBInfoByFileID(fileID string) ([]map[string]inte
for rows.Next() {
var kbID, kbName, docID string
if err := rows.Scan(&kbID, &kbName, &docID); err != nil {
if err = rows.Scan(&kbID, &kbName, &docID); err != nil {
continue
}
results = append(results, map[string]interface{}{
@@ -60,30 +63,30 @@ func (dao *File2DocumentDAO) GetKBInfoByFileID(fileID string) ([]map[string]inte
}
// GetByFileID gets file2document mappings by file ID
func (dao *File2DocumentDAO) GetByFileID(fileID string) ([]*entity.File2Document, error) {
func (dao *File2DocumentDAO) GetByFileID(ctx context.Context, db *gorm.DB, fileID string) ([]*entity.File2Document, error) {
var mappings []*entity.File2Document
err := DB.Where("file_id = ?", fileID).Find(&mappings).Error
err := db.WithContext(ctx).Where("file_id = ?", fileID).Find(&mappings).Error
return mappings, err
}
// DeleteByFileID deletes file2document mappings by file ID
func (dao *File2DocumentDAO) DeleteByFileID(fileID string) error {
return DB.Unscoped().Where("file_id = ?", fileID).Delete(&entity.File2Document{}).Error
func (dao *File2DocumentDAO) DeleteByFileID(ctx context.Context, db *gorm.DB, fileID string) error {
return db.WithContext(ctx).Unscoped().Where("file_id = ?", fileID).Delete(&entity.File2Document{}).Error
}
// GetByDocumentID gets file2document mappings by document ID
func (dao *File2DocumentDAO) GetByDocumentID(docID string) ([]*entity.File2Document, error) {
func (dao *File2DocumentDAO) GetByDocumentID(ctx context.Context, db *gorm.DB, docID string) ([]*entity.File2Document, error) {
var mappings []*entity.File2Document
err := DB.Where("document_id = ?", docID).Find(&mappings).Error
err := db.WithContext(ctx).Where("document_id = ?", docID).Find(&mappings).Error
return mappings, err
}
// DeleteByDocumentID deletes file2document mappings by document ID
func (dao *File2DocumentDAO) DeleteByDocumentID(docID string) error {
return DB.Unscoped().Where("document_id = ?", docID).Delete(&entity.File2Document{}).Error
func (dao *File2DocumentDAO) DeleteByDocumentID(ctx context.Context, db *gorm.DB, docID string) error {
return db.WithContext(ctx).Unscoped().Where("document_id = ?", docID).Delete(&entity.File2Document{}).Error
}
// Create inserts a new file2document mapping record.
func (dao *File2DocumentDAO) Create(mapping *entity.File2Document) error {
return DB.Create(mapping).Error
func (dao *File2DocumentDAO) Create(ctx context.Context, db *gorm.DB, mapping *entity.File2Document) error {
return db.WithContext(ctx).Create(mapping).Error
}

20
internal/dao/file2document_test.go
View File

@@ -17,6 +17,7 @@
package dao
import (
"context"
"testing"
"github.com/glebarez/sqlite"
@@ -74,7 +75,8 @@ func TestFile2DocumentDAO_GetByDocumentID(t *testing.T) {
f2d := testFile2Document(t, "file-1", "doc-1")
results, err := dao.GetByDocumentID("doc-1")
ctx := context.Background()
results, err := dao.GetByDocumentID(ctx, db, "doc-1")
if err != nil {
t.Fatalf("GetByDocumentID failed: %v", err)
}
@@ -94,7 +96,8 @@ func TestFile2DocumentDAO_GetByDocumentID_NotFound(t *testing.T) {
pushDB(t, db)
dao := NewFile2DocumentDAO()
results, err := dao.GetByDocumentID("nonexistent")
ctx := context.Background()
results, err := dao.GetByDocumentID(ctx, db, "nonexistent")
if err != nil {
t.Fatalf("GetByDocumentID failed: %v", err)
}
@@ -111,7 +114,8 @@ func TestFile2DocumentDAO_GetByDocumentID_MultipleResults(t *testing.T) {
testFile2Document(t, "file-1", "doc-shared")
testFile2Document(t, "file-2", "doc-shared")
results, err := dao.GetByDocumentID("doc-shared")
ctx := context.Background()
results, err := dao.GetByDocumentID(ctx, db, "doc-shared")
if err != nil {
t.Fatalf("GetByDocumentID failed: %v", err)
}
@@ -129,13 +133,14 @@ func TestFile2DocumentDAO_DeleteByDocumentID(t *testing.T) {
testFile2Document(t, "file-2", "doc-del")
testFile2Document(t, "file-3", "doc-keep")
err := dao.DeleteByDocumentID("doc-del")
ctx := context.Background()
err := dao.DeleteByDocumentID(ctx, db, "doc-del")
if err != nil {
t.Fatalf("DeleteByDocumentID failed: %v", err)
}
// Verify deleted records are gone
results, err := dao.GetByDocumentID("doc-del")
results, err := dao.GetByDocumentID(ctx, db, "doc-del")
if err != nil {
t.Fatalf("GetByDocumentID failed: %v", err)
}
@@ -144,7 +149,7 @@ func TestFile2DocumentDAO_DeleteByDocumentID(t *testing.T) {
}
// Verify other document's records are untouched
results, err = dao.GetByDocumentID("doc-keep")
results, err = dao.GetByDocumentID(ctx, db, "doc-keep")
if err != nil {
t.Fatalf("GetByDocumentID failed: %v", err)
}
@@ -158,7 +163,8 @@ func TestFile2DocumentDAO_DeleteByDocumentID_Noop(t *testing.T) {
pushDB(t, db)
dao := NewFile2DocumentDAO()
err := dao.DeleteByDocumentID("nonexistent")
ctx := context.Background()
err := dao.DeleteByDocumentID(ctx, db, "nonexistent")
if err != nil {
t.Fatalf("DeleteByDocumentID should not error on missing: %v", err)
}

4
internal/entity/base.go
View File

@@ -157,6 +157,10 @@ func (j *JSONMap) Scan(value interface{}) error {
return json.Unmarshal(b, j)
}
func (j *JSONMap) GormDataType() string {
return "longtext"
}
// JSONSlice is a slice type that can store JSON array data
type JSONSlice []interface{}

2
internal/ingestion/component/document_storage.go
View File

@@ -92,7 +92,7 @@ func ResolveDocumentStorage(ctx context.Context, docID string) (*DocumentStorage
}
ref := &DocumentStorageRef{Name: documentNameOrID(doc)}
mappings, err := dao.NewFile2DocumentDAO().GetByDocumentID(doc.ID)
mappings, err := dao.NewFile2DocumentDAO().GetByDocumentID(ctx, dao.DB, doc.ID)
if err != nil {
return nil, err
}

12
internal/ingestion/task/pipeline_real_integration_test.go
View File

@@ -595,10 +595,16 @@ func mustSeedTaskRealPipelineDocumentBytes(
content []byte,
) {
t.Helper()
ocrID := "ocr-default"
if err := db.Create(&entity.Tenant{
ID: tenantID,
LLMID: "gpt-4",
Status: taskStrPtr("1"),
ID: tenantID,
LLMID: "gpt-4",
ASRID: "asr-default",
Img2TxtID: "img2txt-default",
RerankID: "rerank-default",
ParserIDs: "parser-default",
OCRID: &ocrID,
Status: taskStrPtr("1"),
}).Error; err != nil {
t.Fatalf("create tenant: %v", err)
}

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

@@ -37,7 +37,7 @@ func (s *DocumentService) GetDocumentStorageAddress(ctx context.Context, doc *en
file2DocumentDAO := dao.NewFile2DocumentDAO()
fileDAO := dao.NewFileDAO()
mappings, err := file2DocumentDAO.GetByDocumentID(doc.ID)
mappings, err := file2DocumentDAO.GetByDocumentID(ctx, dao.DB, doc.ID)
if err != nil {
return "", "", err
}
@@ -400,7 +400,7 @@ func (s *DocumentService) rollbackAddFileFromKBError(ctx context.Context, doc *e
// 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(ctx context.Context, docID string) error {
mappings, mapErr := s.file2DocumentDAO.GetByDocumentID(docID)
mappings, mapErr := s.file2DocumentDAO.GetByDocumentID(ctx, dao.DB, docID)
if mapErr != nil {
common.Logger.Warn(fmt.Sprintf("cleanupFileReferences: failed to get f2d mappings for %s: %v", docID, mapErr))
return mapErr
@@ -421,7 +421,7 @@ func (s *DocumentService) cleanupFileReferences(ctx context.Context, docID strin
}
// Delete all file2document rows for this document
if delErr := s.file2DocumentDAO.DeleteByDocumentID(docID); delErr != nil {
if delErr := s.file2DocumentDAO.DeleteByDocumentID(ctx, dao.DB, docID); delErr != nil {
common.Logger.Warn(fmt.Sprintf("cleanupFileReferences: failed to delete f2d for %s: %v", docID, delErr))
return delErr
}
@@ -429,7 +429,7 @@ func (s *DocumentService) cleanupFileReferences(ctx context.Context, docID strin
// For each file, only delete the record and blob when it is a
// knowledgebase-owned upload and no other doc references it
for _, fileID := range fileIDs {
remaining, remErr := s.file2DocumentDAO.GetByFileID(fileID)
remaining, remErr := s.file2DocumentDAO.GetByFileID(ctx, dao.DB, fileID)
if remErr != nil {
common.Logger.Warn(fmt.Sprintf("cleanupFileReferences: failed to check remaining f2d for %s: %v", fileID, remErr))
continue

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

@@ -347,7 +347,7 @@ func (s *DocumentService) updateDocumentNameOnly(ctx context.Context, doc *entit
return errors.New("database error (Document rename)")
}
mappings, err := s.file2DocumentDAO.GetByDocumentID(doc.ID)
mappings, err := s.file2DocumentDAO.GetByDocumentID(ctx, dao.DB, doc.ID)
if err == nil && len(mappings) > 0 && mappings[0].FileID != nil && s.fileDAO != nil {
err = s.fileDAO.UpdateByID(ctx, dao.DB, *mappings[0].FileID, map[string]interface{}{"name": newName})
if err != nil {

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

@@ -658,7 +658,7 @@ func TestDeleteDocumentFull_CleansUpFile2Document(t *testing.T) {
// Verify f2d mapping deleted
f2dDAO := dao.NewFile2DocumentDAO()
mappings, _ := f2dDAO.GetByDocumentID("doc-1")
mappings, _ := f2dDAO.GetByDocumentID(ctx, db, "doc-1")
if len(mappings) != 0 {
t.Fatalf("expected 0 f2d mappings, got %d", len(mappings))
}
@@ -696,7 +696,7 @@ func TestDeleteDocumentFull_SharedFilePreserved(t *testing.T) {
// f2d mapping for doc-1 should be gone
f2dDAO := dao.NewFile2DocumentDAO()
mappings, _ := f2dDAO.GetByDocumentID("doc-1")
mappings, _ := f2dDAO.GetByDocumentID(ctx, db, "doc-1")
if len(mappings) != 0 {
t.Fatalf("expected 0 f2d mappings for doc-1, got %d", len(mappings))
}
@@ -708,7 +708,7 @@ func TestDeleteDocumentFull_SharedFilePreserved(t *testing.T) {
}
// f2d mapping for doc-2 should still exist
mappings, _ = f2dDAO.GetByDocumentID("doc-2")
mappings, _ = f2dDAO.GetByDocumentID(ctx, db, "doc-2")
if len(mappings) != 1 {
t.Fatalf("expected 1 f2d mapping for doc-2, got %d", len(mappings))
}
@@ -1495,7 +1495,7 @@ func TestCleanupFileReferences_SingleFileDeleted(t *testing.T) {
svc.cleanupFileReferences(ctx, "doc-1")
// f2d gone
mappings, _ := dao.NewFile2DocumentDAO().GetByDocumentID("doc-1")
mappings, _ := dao.NewFile2DocumentDAO().GetByDocumentID(ctx, db, "doc-1")
if len(mappings) != 0 {
t.Fatalf("expected 0 f2d after cleanup, got %d", len(mappings))
}
@@ -1520,7 +1520,7 @@ func TestCleanupFileReferences_SharedFileSurvives(t *testing.T) {
svc.cleanupFileReferences(ctx, "doc-1")
// f2d for doc-1 gone
mappings, _ := dao.NewFile2DocumentDAO().GetByDocumentID("doc-1")
mappings, _ := dao.NewFile2DocumentDAO().GetByDocumentID(ctx, db, "doc-1")
if len(mappings) != 0 {
t.Fatalf("expected 0 f2d for doc-1, got %d", len(mappings))
}
@@ -1530,7 +1530,7 @@ func TestCleanupFileReferences_SharedFileSurvives(t *testing.T) {
t.Fatalf("expected 1 file record, got %d", len(files))
}
// f2d for doc-2 survives
mappings, _ = dao.NewFile2DocumentDAO().GetByDocumentID("doc-2")
mappings, _ = dao.NewFile2DocumentDAO().GetByDocumentID(ctx, db, "doc-2")
if len(mappings) != 1 {
t.Fatalf("expected 1 f2d for doc-2, got %d", len(mappings))
}

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

@@ -195,7 +195,7 @@ func (s *DocumentService) newAFileFromKB(ctx context.Context, tenantID, name, pa
// dataset folder plus a file2document mapping. Mirrors Python
// FileService.add_file_from_kb (idempotent on the document mapping).
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 {
if existing, err := s.file2DocumentDAO.GetByDocumentID(ctx, dao.DB, doc.ID); err == nil && len(existing) > 0 {
return nil
}
name := ""
@@ -222,7 +222,7 @@ func (s *DocumentService) addFileFromKB(ctx context.Context, doc *entity.Documen
return err
}
docID := doc.ID
if err := s.file2DocumentDAO.Create(&entity.File2Document{
if err := s.file2DocumentDAO.Create(ctx, dao.DB, &entity.File2Document{
ID: utility.GenerateToken(),
FileID: &fileID,
DocumentID: &docID,

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

@@ -169,7 +169,7 @@ func (s *File2DocumentService) LinkToDatasets(ctx context.Context, userID string
func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs []string, userID, mode string) error {
replaceExisting := mode != "add"
for _, fileID := range fileIDs {
mappings, err := s.file2DocumentDAO.GetByFileID(fileID)
mappings, err := s.file2DocumentDAO.GetByFileID(ctx, dao.DB, fileID)
if err != nil {
common.Warn("convertFiles: GetByFileID failed", zap.String("fileID", fileID), zap.Error(err))
}
@@ -191,7 +191,7 @@ func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs
}
// Drop the file2document mappings for this file (mirrors Python
// File2DocumentService.delete_by_file_id, done once per file).
if err := s.file2DocumentDAO.DeleteByFileID(fileID); err != nil {
if err = s.file2DocumentDAO.DeleteByFileID(ctx, dao.DB, fileID); err != nil {
common.Warn("convertFiles: DeleteByFileID failed", zap.String("fileID", fileID), zap.Error(err))
}
} else {
@@ -271,7 +271,7 @@ func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs
FileID: &fileID,
DocumentID: &doc.ID,
}
if err := s.file2DocumentDAO.Create(mapping); err != nil {
if err = s.file2DocumentDAO.Create(ctx, dao.DB, mapping); err != nil {
common.Warn("convertFiles: Create file2document mapping failed",
zap.String("fileID", fileID), zap.String("docID", doc.ID), zap.Error(err))
}

2
internal/service/file/file_content.go
View File

@@ -30,7 +30,7 @@ func (s *FileService) GetFileContent(ctx context.Context, uid, fileID string) (*
// Matches Python's File2DocumentService.get_storage_address function
func (s *FileService) GetStorageAddress(ctx context.Context, fileID string) (*StorageAddress, error) {
// Get file2document mapping
f2d, err := s.file2DocumentDAO.GetByFileID(fileID)
f2d, err := s.file2DocumentDAO.GetByFileID(ctx, dao.DB, fileID)
if err != nil || len(f2d) == 0 {
return nil, fmt.Errorf("file2document mapping not found")
}

4
internal/service/file/file_delete.go
View File

@@ -69,7 +69,7 @@ func (s *FileService) deleteSingleFile(ctx context.Context, file *entity.File) e
}
// 2. Handle associated documents
informs, err := s.file2DocumentDAO.GetByFileID(file.ID)
informs, err := s.file2DocumentDAO.GetByFileID(ctx, dao.DB, file.ID)
if err != nil {
return fmt.Errorf("failed to get file2document mappings: %w", err)
}
@@ -90,7 +90,7 @@ func (s *FileService) deleteSingleFile(ctx context.Context, file *entity.File) e
}
// Delete file2document mapping (outside the loop, called once - matching Python behavior)
if err = s.file2DocumentDAO.DeleteByFileID(file.ID); err != nil {
if err = s.file2DocumentDAO.DeleteByFileID(ctx, dao.DB, file.ID); err != nil {
return fmt.Errorf("failed to delete file2document mapping: %w", err)
}
}

6
internal/service/file/file_folder.go
View File

@@ -84,7 +84,7 @@ func (s *FileService) ListFiles(ctx context.Context, tenantID, pfID string, page
fileInfo.KbsInfo = []map[string]interface{}{}
} else {
// Get KB info for non-folder files
kbsInfo, err := s.file2DocumentDAO.GetKBInfoByFileID(file.ID)
kbsInfo, err := s.file2DocumentDAO.GetKBInfoByFileID(ctx, dao.DB, file.ID)
if err != nil {
kbsInfo = []map[string]interface{}{}
}
@@ -490,7 +490,7 @@ func (s *FileService) MoveFiles(ctx context.Context, uid string, srcFileIDs []st
}
// Update associated document name if exists
informs, err := s.file2DocumentDAO.GetByFileID(file.ID)
informs, err := s.file2DocumentDAO.GetByFileID(ctx, dao.DB, file.ID)
if err == nil && len(informs) > 0 && informs[0].DocumentID != nil {
docID := *informs[0].DocumentID
documentDAO := dao.NewDocumentDAO()
@@ -589,7 +589,7 @@ func (s *FileService) moveEntryRecursive(ctx context.Context, sourceFile *entity
// Update associated document name if renamed
if overrideName != "" {
informs, err := s.file2DocumentDAO.GetByFileID(sourceFile.ID)
informs, err := s.file2DocumentDAO.GetByFileID(ctx, dao.DB, sourceFile.ID)
if err == nil && len(informs) > 0 && informs[0].DocumentID != nil {
docID := *informs[0].DocumentID
documentDAO := dao.NewDocumentDAO()