zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-05 15:20:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: add context to storage (#17690)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-08-03 15:08:21 +08:00
committed by GitHub
parent 2dbdaf951d
commit bddc941814
42 changed files with 634 additions and 527 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
internal/service/document/document.go
View File

@@ -17,8 +17,10 @@
package document
import (
"context"
"errors"
"ragflow/internal/service"
"ragflow/internal/storage"
"regexp"
"time"
@@ -219,10 +221,10 @@ type IngestDocumentRequest struct {
// StartParseOptions controls StartParseDocuments behavior.
type StartParseOptions struct {
// ApplyKB merges the knowledgebase's parser_config (llm_id, metadata)
// ApplyKB merges the knowledge base's parser_config (llm_id, metadata)
// into the document before parsing.
ApplyKB bool
// RerunWithDelete clears prior chunks/tasks/counters before re-parsing.
// RerunWithDelete clears prior chunks/tasks/counters before reparsing.
RerunWithDelete bool
}
@@ -252,7 +254,7 @@ const knowledgebaseFolderName = ".knowledgebase"
const maxUploadDocSize = 128 * 1024 * 1024
// MetadataUpdate is one update item: set key to value.
type DocumentMetadataUpdate struct {
type MetadataUpdate struct {
Key string `json:"key"`
Value interface{} `json:"value"`
Match interface{} `json:"match,omitempty"`
@@ -260,19 +262,44 @@ type DocumentMetadataUpdate struct {
}
// MetadataDelete removes a whole key, or a specific value from a list field.
type DocumentMetadataDelete struct {
type MetadataDelete struct {
Key string `json:"key"`
Value interface{} `json:"value,omitempty"`
}
// MetadataSelector selects which documents to target.
type DocumentMetadataSelector struct {
type MetadataSelector struct {
DocumentIDs []string `json:"document_ids"`
MetadataCondition map[string]interface{} `json:"metadata_condition"`
}
// BatchUpdateDocumentMetadatasResponse summarises the operation.
type BatchUpdateDocumentMetadatasResponse struct {
// BatchUpdateMetadatasResponse summarises the operation.
type BatchUpdateMetadatasResponse struct {
Updated int `json:"updated"`
MatchedDocs int `json:"matched_docs"`
}
// removeObjectBestEffort retries blob deletion on a context that survives the
// originating request. Bounded by a timeout so a wedged storage SDK cannot
// block the caller forever. It always uses the parent request's storage impl,
// but deliberately NOT the request context, because a cancelled request must
// not leak the blob it already wrote (or orphan a blob whose row was deleted).
func removeObjectBestEffort(storageImpl storage.Storage, bucket, object string) error {
ctx, cancel := context.WithTimeout(context.WithoutCancel(context.Background()), 30*time.Second)
defer cancel()
var lastErr error
for attempt := 0; attempt < 3; attempt++ {
if err := storageImpl.Remove(ctx, bucket, object); err != nil {
lastErr = err
// Treat cancellation of the *new* cleanup ctx as terminal.
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
break
}
time.Sleep(time.Duration(attempt+1) * 100 * time.Millisecond)
continue
}
return nil
}
return lastErr
}

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

@@ -25,7 +25,7 @@ func (s *DocumentService) GetDocumentImage(ctx context.Context, imageID string)
return nil, fmt.Errorf("storage not initialized")
}
return storageImpl.Get(parts[0], parts[1])
return storageImpl.Get(ctx, parts[0], parts[1])
}
// GetDocumentArtifact retrieves a sandbox artifact from object storage.
@@ -61,11 +61,11 @@ func (s *DocumentService) GetDocumentArtifact(ctx context.Context, filename, use
}
bucket := sandboxArtifactBucket()
if !storageImpl.ObjExist(bucket, basename) {
if !storageImpl.ObjExist(ctx, bucket, basename) {
return nil, ErrArtifactNotFound
}
data, err := storageImpl.Get(bucket, basename)
data, err := storageImpl.Get(ctx, bucket, basename)
if err != nil {
return nil, err
}
@@ -209,7 +209,7 @@ func (s *DocumentService) GetDocumentPreview(ctx context.Context, docID string)
return nil, fmt.Errorf("storage not initialized")
}
data, err := storageImpl.Get(bucket, name)
data, err := storageImpl.Get(ctx, bucket, name)
if err != nil {
return nil, err
}

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

@@ -83,7 +83,7 @@ func (s *DocumentService) DownloadDocument(ctx context.Context, datasetID, docID
return nil, fmt.Errorf("storage not initialized")
}
data, err := storageImpl.Get(bucket, name)
data, err := storageImpl.Get(ctx, bucket, name)
if err != nil {
return nil, err
}
@@ -279,14 +279,11 @@ func (s *DocumentService) RemoveDocumentKeepFile(ctx context.Context, docID stri
if err != nil {
return err
}
if _, delErr := s.taskDAO.DeleteByDocIDs(ctx, dao.DB, []string{docID}); delErr != nil {
common.Logger.Warn(fmt.Sprintf("RemoveDocumentKeepFile: failed to delete tasks for %s: %v", docID, delErr))
}
if _, delErr := s.taskDAO.DeleteByDocIDs(ctx, dao.DB, []string{docID}); delErr != nil {
if errors.Is(delErr, context.Canceled) || errors.Is(delErr, context.DeadlineExceeded) {
return fmt.Errorf("RemoveDocumentKeepFile: failed to delete tasks for %s: %w", docID, delErr)
}
common.Logger.Warn(fmt.Sprintf("RemoveDocumentKeepFile: failed to delete tasks for %s: %v", docID, delErr))
common.Warn(fmt.Sprintf("RemoveDocumentKeepFile: failed to delete tasks for %s: %v", docID, delErr))
}
return s.deleteDocRecordWithCounters(ctx, doc, kb.ID)
}
@@ -357,7 +354,7 @@ func (s *DocumentService) deleteDocEngineData(docID, tenantID, kbID string) {
ctx := context.Background()
indexName := fmt.Sprintf("ragflow_%s", tenantID)
if _, delErr := s.docEngine.DeleteChunks(ctx, map[string]interface{}{"doc_id": docID}, indexName, kbID); delErr != nil {
common.Logger.Warn(fmt.Sprintf("deleteDocEngineData: failed to delete chunks for %s: %v", docID, delErr))
common.Warn(fmt.Sprintf("deleteDocEngineData: failed to delete chunks for %s: %v", docID, delErr))
}
// Notify the dataset-level post-processing consumer (§11) that this document's
// source + per-doc compiled chunks are gone. The consumer removes the
@@ -369,7 +366,7 @@ func (s *DocumentService) deleteDocEngineData(docID, tenantID, kbID string) {
pubCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
if err := knowledge_compile.PublishDeleted(pubCtx, tenantID, kbID, docID, 0); err != nil {
common.Logger.Warn(fmt.Sprintf("deleteDocEngineData: publish doc_deleted for %s failed: %v", docID, err))
common.Warn(fmt.Sprintf("deleteDocEngineData: publish doc_deleted for %s failed: %v", docID, err))
}
if s.metadataSvc != nil {
_ = s.DeleteDocumentAllMetadata(ctx, docID) // logs internally
@@ -422,7 +419,7 @@ func (s *DocumentService) rollbackAddFileFromKBError(ctx context.Context, doc *e
func (s *DocumentService) cleanupFileReferences(ctx context.Context, docID string) error {
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))
common.Warn(fmt.Sprintf("cleanupFileReferences: failed to get f2d mappings for %s: %v", docID, mapErr))
return mapErr
}
if len(mappings) == 0 {
@@ -442,7 +439,7 @@ func (s *DocumentService) cleanupFileReferences(ctx context.Context, docID strin
// Delete all file2document rows for this document
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))
common.Warn(fmt.Sprintf("cleanupFileReferences: failed to delete f2d for %s: %v", docID, delErr))
return delErr
}
@@ -451,7 +448,7 @@ func (s *DocumentService) cleanupFileReferences(ctx context.Context, docID strin
for _, fileID := range fileIDs {
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))
common.Warn(fmt.Sprintf("cleanupFileReferences: failed to check remaining f2d for %s: %v", fileID, remErr))
continue
}
if len(remaining) > 0 {
@@ -461,21 +458,22 @@ func (s *DocumentService) cleanupFileReferences(ctx context.Context, docID strin
fileDAO := dao.NewFileDAO()
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))
common.Warn(fmt.Sprintf("cleanupFileReferences: file not found %s: %v", fileID, fErr))
continue
}
if entity.FileSource(file.SourceType) != entity.FileSourceKnowledgebase {
continue // linked from file management — unlink only, keep the file
}
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))
common.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
}
if file.Location != nil && *file.Location != "" {
storageImpl := storage.GetStorageFactory().GetStorage()
if storageImpl != nil {
if rmErr := storageImpl.Remove(file.ParentID, *file.Location); rmErr != nil {
common.Logger.Warn(fmt.Sprintf("cleanupFileReferences: failed to remove blob %s/%s: %v", file.ParentID, *file.Location, rmErr))
rmErr := removeObjectBestEffort(storageImpl, file.ParentID, *file.Location)
if rmErr != nil {
common.Warn(fmt.Sprintf("cleanupFileReferences: failed to remove blob %s/%s: %v", file.ParentID, *file.Location, rmErr))
}
}
}

24
internal/service/document/document_metadata.go
View File

@@ -563,12 +563,12 @@ func (s *DocumentService) patchDocumentMetadata(ctx context.Context, docID strin
func (s *DocumentService) BatchUpdateDocumentMetadatas(
ctx context.Context,
datasetID string,
selector *DocumentMetadataSelector,
updates []DocumentMetadataUpdate,
deletes []DocumentMetadataDelete,
) (*BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error) {
selector *MetadataSelector,
updates []MetadataUpdate,
deletes []MetadataDelete,
) (*BatchUpdateMetadatasResponse, common.ErrorCode, error) {
if selector == nil {
selector = &DocumentMetadataSelector{}
selector = &MetadataSelector{}
}
if code, err := validateBatchUpdateDocumentMetadatasRequest(selector, updates, deletes); err != nil {
return nil, code, err
@@ -635,7 +635,7 @@ func (s *DocumentService) BatchUpdateDocumentMetadatas(
// Early-exit when conditions given but nothing matched.
rawConds, _ := selector.MetadataCondition["conditions"]
if rawConds != nil && len(targetDocIDs) == 0 {
return &BatchUpdateDocumentMetadatasResponse{Updated: 0, MatchedDocs: 0}, common.CodeSuccess, nil
return &BatchUpdateMetadatasResponse{Updated: 0, MatchedDocs: 0}, common.CodeSuccess, nil
}
}
@@ -675,13 +675,13 @@ func (s *DocumentService) BatchUpdateDocumentMetadatas(
updated++
}
return &BatchUpdateDocumentMetadatasResponse{Updated: updated, MatchedDocs: len(ids)}, common.CodeSuccess, nil
return &BatchUpdateMetadatasResponse{Updated: updated, MatchedDocs: len(ids)}, common.CodeSuccess, nil
}
func validateBatchUpdateDocumentMetadatasRequest(
selector *DocumentMetadataSelector,
updates []DocumentMetadataUpdate,
deletes []DocumentMetadataDelete,
selector *MetadataSelector,
updates []MetadataUpdate,
deletes []MetadataDelete,
) (common.ErrorCode, error) {
for _, upd := range updates {
if strings.TrimSpace(upd.Key) == "" || upd.Value == nil {
@@ -778,7 +778,7 @@ func cloneDocumentMetadataValue(v interface{}) interface{} {
}
}
func applyDocumentMetadataUpdates(meta map[string]interface{}, updates []DocumentMetadataUpdate) bool {
func applyDocumentMetadataUpdates(meta map[string]interface{}, updates []MetadataUpdate) bool {
changed := false
for _, upd := range updates {
key := strings.TrimSpace(upd.Key)
@@ -854,7 +854,7 @@ func applyDocumentMetadataUpdates(meta map[string]interface{}, updates []Documen
return changed
}
func applyDocumentMetadataDeletes(meta map[string]interface{}, deletes []DocumentMetadataDelete) bool {
func applyDocumentMetadataDeletes(meta map[string]interface{}, deletes []MetadataDelete) bool {
changed := false
for _, del := range deletes {
key := strings.TrimSpace(del.Key)

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

@@ -104,14 +104,14 @@ func (s *DocumentService) clearDocumentParseResults(ctx context.Context, doc *en
}
indexName := fmt.Sprintf("ragflow_%s", tenantID)
exists, err := s.docEngine.ChunkStoreExists(context.Background(), indexName, doc.KbID)
exists, err := s.docEngine.ChunkStoreExists(ctx, indexName, doc.KbID)
if err != nil {
return err
}
if !exists {
return nil
}
if _, err := s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil {
if _, err = s.docEngine.DeleteChunks(ctx, map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil {
return err
}
return nil
@@ -380,8 +380,8 @@ func (s *DocumentService) resetDocumentForReparse(ctx context.Context, doc *enti
}
if s.docEngine != nil {
indexName := fmt.Sprintf("ragflow_%s", tenantID)
s.deleteChunkImages(doc, indexName)
if _, err = s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil {
s.deleteChunkImages(ctx, doc, indexName)
if _, err = s.docEngine.DeleteChunks(ctx, map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil {
return err
}
}
@@ -390,7 +390,7 @@ func (s *DocumentService) resetDocumentForReparse(ctx context.Context, doc *enti
return nil
}
func (s *DocumentService) deleteChunkImages(doc *entity.Document, indexName string) {
func (s *DocumentService) deleteChunkImages(ctx context.Context, doc *entity.Document, indexName string) {
if s.docEngine == nil {
return
}
@@ -401,7 +401,7 @@ func (s *DocumentService) deleteChunkImages(doc *entity.Document, indexName stri
const pageSize = 1000
for offset := 0; ; offset += pageSize {
result, err := s.docEngine.Search(context.Background(), &enginetypes.SearchRequest{
result, err := s.docEngine.Search(ctx, &enginetypes.SearchRequest{
IndexNames: []string{indexName},
KbIDs: []string{doc.KbID},
Offset: offset,
@@ -420,8 +420,8 @@ func (s *DocumentService) deleteChunkImages(doc *entity.Document, indexName stri
if !ok {
continue
}
if storageImpl.ObjExist(doc.KbID, imageKey) {
_ = storageImpl.Remove(doc.KbID, imageKey)
if storageImpl.ObjExist(ctx, doc.KbID, imageKey) {
_ = storageImpl.Remove(ctx, doc.KbID, imageKey)
}
}
}
@@ -499,7 +499,7 @@ func (s *DocumentService) updateDocumentStatusOnly(ctx context.Context, doc *ent
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
return s.docEngine.UpdateChunks(
context.Background(),
ctx,
map[string]interface{}{"doc_id": doc.ID},
map[string]interface{}{"available_int": status},
indexName,

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

@@ -63,36 +63,36 @@ func newFakeUploadStorage() *fakeUploadStorage {
}
func (f *fakeUploadStorage) Type() string { return "fake_upload_storage" }
func (f *fakeUploadStorage) Health() bool { return true }
func (f *fakeUploadStorage) Health(_ context.Context) bool { return true }
func (f *fakeUploadStorage) key(bucket, fnm string) string { return bucket + "/" + fnm }
func (f *fakeUploadStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
func (f *fakeUploadStorage) Put(ctx context.Context, bucket, fnm string, binary []byte, tenantID ...string) error {
f.objects[f.key(bucket, fnm)] = append([]byte(nil), binary...)
return nil
}
func (f *fakeUploadStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
func (f *fakeUploadStorage) Get(ctx context.Context, bucket, fnm string, tenantID ...string) ([]byte, error) {
v, ok := f.objects[f.key(bucket, fnm)]
if !ok {
return nil, errors.New("not found")
}
return append([]byte(nil), v...), nil
}
func (f *fakeUploadStorage) Remove(bucket, fnm string, tenantID ...string) error {
func (f *fakeUploadStorage) Remove(ctx context.Context, bucket, fnm string, tenantID ...string) error {
delete(f.objects, f.key(bucket, fnm))
return nil
}
func (f *fakeUploadStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
func (f *fakeUploadStorage) ObjExist(ctx context.Context, bucket, fnm string, tenantID ...string) bool {
_, ok := f.objects[f.key(bucket, fnm)]
return ok
}
func (f *fakeUploadStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
func (f *fakeUploadStorage) ListObjects(ctx context.Context, bucket string, tenantID ...string) ([]string, error) {
return []string{}, nil
}
func (f *fakeUploadStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
func (f *fakeUploadStorage) GetPresignedURL(ctx context.Context, bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
return "", nil
}
func (f *fakeUploadStorage) BucketExists(bucket string) bool { return true }
func (f *fakeUploadStorage) RemoveBucket(bucket string) error { return nil }
func (f *fakeUploadStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
func (f *fakeUploadStorage) BucketExists(ctx context.Context, bucket string) bool { return true }
func (f *fakeUploadStorage) RemoveBucket(ctx context.Context, bucket string) error { return nil }
func (f *fakeUploadStorage) Copy(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
v, ok := f.objects[f.key(srcBucket, srcPath)]
if !ok {
return false
@@ -100,8 +100,8 @@ func (f *fakeUploadStorage) Copy(srcBucket, srcPath, destBucket, destPath string
f.objects[f.key(destBucket, destPath)] = append([]byte(nil), v...)
return true
}
func (f *fakeUploadStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
if !f.Copy(srcBucket, srcPath, destBucket, destPath) {
func (f *fakeUploadStorage) Move(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
if !f.Copy(ctx, srcBucket, srcPath, destBucket, destPath) {
return false
}
delete(f.objects, f.key(srcBucket, srcPath))
@@ -796,7 +796,7 @@ func TestUploadLocalDocuments_MirrorsPythonCoreFields(t *testing.T) {
t.Fatalf("parser_config=%v", cfg)
}
storedBlob, err := mockStorage.Get(kb.ID, "nested/path/deck(1).pptx")
storedBlob, err := mockStorage.Get(ctx, kb.ID, "nested/path/deck(1).pptx")
if err != nil {
t.Fatalf("blob not stored: %v", err)
}
@@ -2119,12 +2119,12 @@ func TestBatchUpdateDocumentMetadatasMatchesPythonSemantics(t *testing.T) {
svc.docEngine = engine
svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine)
ctx := t.Context()
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &MetadataSelector{
DocumentIDs: []string{"doc-1", "doc-2", "doc-3"},
}, []DocumentMetadataUpdate{
}, []MetadataUpdate{
{Key: "tags", Value: "new", Match: "old"},
{Key: "category", Value: "paper"},
}, []DocumentMetadataDelete{
}, []MetadataDelete{
{Key: "author", Value: "alice"},
})
if err != nil {
@@ -2180,9 +2180,9 @@ func TestBatchUpdateDocumentMetadatasDoesNotReplaceWhenCurrentSearchIsStale(t *t
svc.docEngine = engine
svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine)
ctx := t.Context()
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &MetadataSelector{
DocumentIDs: []string{"doc-1"},
}, []DocumentMetadataUpdate{
}, []MetadataUpdate{
{Key: "category", Value: "paper"},
}, nil)
if err != nil || code != common.CodeSuccess {
@@ -2217,9 +2217,9 @@ func TestBatchUpdateDocumentMetadatasDeletesEmptyMetadataAndNoOps(t *testing.T)
svc.docEngine = engine
svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine)
ctx := t.Context()
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &MetadataSelector{
DocumentIDs: []string{"doc-1", "doc-2"},
}, nil, []DocumentMetadataDelete{{Key: "status", Value: "draft"}})
}, nil, []MetadataDelete{{Key: "status", Value: "draft"}})
if err != nil || code != common.CodeSuccess {
t.Fatalf("delete batch failed: code=%v err=%v", code, err)
}
@@ -2246,9 +2246,9 @@ func TestBatchUpdateDocumentMetadatasNormalizesNumberValues(t *testing.T) {
svc.docEngine = engine
svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine)
ctx := t.Context()
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &MetadataSelector{
DocumentIDs: []string{"doc-1"},
}, []DocumentMetadataUpdate{
}, []MetadataUpdate{
{Key: "score", Value: "42", ValueType: "number"},
}, nil)
if err != nil || code != common.CodeSuccess {
@@ -2276,7 +2276,7 @@ func TestBatchUpdateDocumentMetadatasNormalizesNumberValues(t *testing.T) {
func TestBatchUpdateDocumentMetadatasRejectsMissingValue(t *testing.T) {
svc := testDocumentService(t)
ctx := t.Context()
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{}, []DocumentMetadataUpdate{
resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &MetadataSelector{}, []MetadataUpdate{
{Key: "status"},
}, nil)
if err == nil {

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

@@ -7,13 +7,12 @@ import (
"mime/multipart"
"net/http"
"path/filepath"
"ragflow/internal/dao"
"strings"
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/entity"
"ragflow/internal/storage"
"ragflow/internal/utility"
"strings"
)
// UploadLocalDocuments stores each uploaded file in object storage and inserts a
@@ -86,10 +85,10 @@ func (s *DocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.K
if safeParent != "" {
location = safeParent + "/" + filename
}
for storageImpl.ObjExist(kb.ID, location) {
for storageImpl.ObjExist(ctx, kb.ID, location) {
location += "_"
}
if err = storageImpl.Put(kb.ID, location, blob); err != nil {
if err = storageImpl.Put(ctx, kb.ID, location, blob); err != nil {
errMsgs = append(errMsgs, fh.Filename+": "+err.Error())
continue
}
@@ -97,7 +96,10 @@ func (s *DocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.K
doc := s.newDatasetDocument(kb, tenantID, filename, location, string(filetype), merged, "local", int64(len(blob)), blob)
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)
rmErr := removeObjectBestEffort(storageImpl, kb.ID, location)
if rmErr != nil {
common.Warn(fmt.Sprintf("upload rollback: failed to remove orphaned blob %s/%s: %v", kb.ID, location, rmErr))
}
errMsgs = append(errMsgs, fh.Filename+": "+err.Error())
continue
}
@@ -105,11 +107,14 @@ func (s *DocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.K
// 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)
_ = storageImpl.Remove(kb.ID, location)
rmErr := removeObjectBestEffort(storageImpl, kb.ID, location)
if rmErr != nil {
common.Warn(fmt.Sprintf("UploadLocalDocuments: failed to remove blob %s/%s: %v", kb.ID, location, rmErr))
}
errMsgs = append(errMsgs, fh.Filename+": "+err.Error())
continue
}
// Only reserve the name once the write fully succeeds.
// Only reserve the name once write fully succeeds.
taken[filename] = true
results = append(results, docToRawMap(doc))
}
@@ -271,21 +276,27 @@ func (s *DocumentService) UploadWebDocument(ctx context.Context, kb *entity.Know
}
location := filename
for storageImpl.ObjExist(kb.ID, location) {
for storageImpl.ObjExist(ctx, kb.ID, location) {
location += "_"
}
if err = storageImpl.Put(kb.ID, location, blob); err != nil {
if err = storageImpl.Put(ctx, kb.ID, location, blob); err != nil {
return nil, common.CodeServerError, err
}
doc := s.newDatasetDocument(kb, tenantID, filename, location, string(filetype), kb.ParserConfig, "web", int64(len(blob)), blob)
if err = s.InsertDocument(doc); err != nil {
_ = storageImpl.Remove(kb.ID, location)
rmErr := removeObjectBestEffort(storageImpl, kb.ID, location)
if rmErr != nil {
common.Warn(fmt.Sprintf("UploadWebDocument: failed to insert document, remove blob %s/%s: %v", kb.ID, location, rmErr))
}
return nil, common.CodeServerError, err
}
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)
rmErr := removeObjectBestEffort(storageImpl, kb.ID, location)
if rmErr != nil {
common.Warn(fmt.Sprintf("UploadWebDocument: failed to add file from knowledge base, remove blob %s/%s: %v", kb.ID, location, rmErr))
}
return nil, common.CodeServerError, err
}
return docToRawMap(doc), common.CodeSuccess, nil