diff --git a/internal/dao/file2document.go b/internal/dao/file2document.go index f707eb1d3b..ee8a3bb36d 100644 --- a/internal/dao/file2document.go +++ b/internal/dao/file2document.go @@ -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 } diff --git a/internal/dao/file2document_test.go b/internal/dao/file2document_test.go index 4b1cb548d6..2b42ffe9b8 100644 --- a/internal/dao/file2document_test.go +++ b/internal/dao/file2document_test.go @@ -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) } diff --git a/internal/entity/base.go b/internal/entity/base.go index d2aabb5e7d..a92cfb443e 100644 --- a/internal/entity/base.go +++ b/internal/entity/base.go @@ -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{} diff --git a/internal/ingestion/component/document_storage.go b/internal/ingestion/component/document_storage.go index 2ddd3b22f5..1bbd31cd31 100644 --- a/internal/ingestion/component/document_storage.go +++ b/internal/ingestion/component/document_storage.go @@ -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 } diff --git a/internal/ingestion/task/pipeline_real_integration_test.go b/internal/ingestion/task/pipeline_real_integration_test.go index d474b7d245..5917b9a229 100644 --- a/internal/ingestion/task/pipeline_real_integration_test.go +++ b/internal/ingestion/task/pipeline_real_integration_test.go @@ -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) } diff --git a/internal/service/document/document_crud.go b/internal/service/document/document_crud.go index d8062b48cd..8dad325df9 100644 --- a/internal/service/document/document_crud.go +++ b/internal/service/document/document_crud.go @@ -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 diff --git a/internal/service/document/document_dataset_update.go b/internal/service/document/document_dataset_update.go index 9c995094c6..baeec31845 100644 --- a/internal/service/document/document_dataset_update.go +++ b/internal/service/document/document_dataset_update.go @@ -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 { diff --git a/internal/service/document/document_test.go b/internal/service/document/document_test.go index b9716a9e0c..1de5d945ea 100644 --- a/internal/service/document/document_test.go +++ b/internal/service/document/document_test.go @@ -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)) } diff --git a/internal/service/document/document_upload.go b/internal/service/document/document_upload.go index c0579fcb03..035c6205a4 100644 --- a/internal/service/document/document_upload.go +++ b/internal/service/document/document_upload.go @@ -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, diff --git a/internal/service/document/file2document.go b/internal/service/document/file2document.go index 3d67bb2fb2..9e587082de 100644 --- a/internal/service/document/file2document.go +++ b/internal/service/document/file2document.go @@ -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)) } diff --git a/internal/service/file/file_content.go b/internal/service/file/file_content.go index d85751b9f4..20e094b64c 100644 --- a/internal/service/file/file_content.go +++ b/internal/service/file/file_content.go @@ -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") } diff --git a/internal/service/file/file_delete.go b/internal/service/file/file_delete.go index 785bbb44d8..2a14fa6ac5 100644 --- a/internal/service/file/file_delete.go +++ b/internal/service/file/file_delete.go @@ -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) } } diff --git a/internal/service/file/file_folder.go b/internal/service/file/file_folder.go index 2bafc0322a..1811eeb19b 100644 --- a/internal/service/file/file_folder.go +++ b/internal/service/file/file_folder.go @@ -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()