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

Go: fix context (#18060)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-08-10 19:11:43 +08:00
committed by GitHub
parent cb649e8d56
commit dbd062f65e
10 changed files with 41 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
internal/handler/mcp.go
View File

@@ -370,7 +370,8 @@ func (h *MCPHandler) TestMCPServer(c *gin.Context) {
return
}
tools, err := h.mcpService.TestServer(mcpID, &req)
ctx := c.Request.Context()
tools, err := h.mcpService.TestServer(ctx, mcpID, &req)
if mcpErrorResponse(c, err) {
return
}

24
internal/service/dataset/tags.go
View File

@@ -51,7 +51,7 @@ func (d *DatasetService) AggregateTags(ctx context.Context, datasetIDs []string,
merged := make(map[string]int)
for tenantID, kbIDs := range datasetIDsByTenant {
for offset := 0; ; offset += pageSize {
searchResp, err := d.docEngine.Search(context.Background(), &enginetypes.SearchRequest{
searchResp, err := d.docEngine.Search(ctx, &enginetypes.SearchRequest{
IndexNames: []string{fmt.Sprintf("ragflow_%s", tenantID)},
KbIDs: kbIDs,
Offset: offset,
@@ -99,7 +99,7 @@ func (d *DatasetService) AggregateTags(ctx context.Context, datasetIDs []string,
func (d *DatasetService) ListTags(ctx context.Context, datasetID, userID string) ([]map[string]interface{}, common.ErrorCode, error) {
datasetID = strings.TrimSpace(datasetID)
if datasetID == "" {
return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"")
return nil, common.CodeDataError, errors.New("lack of \"Dataset ID\"")
}
normalizedID, err := normalizeDatasetID(datasetID)
if err != nil {
@@ -110,16 +110,16 @@ func (d *DatasetService) ListTags(ctx context.Context, datasetID, userID string)
return nil, common.CodeDataError, errors.New("no authorization")
}
if d.docEngine == nil {
return nil, common.CodeServerError, errors.New("Document engine is not initialized")
return nil, common.CodeServerError, errors.New("document engine is not initialized")
}
kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID)
if err != nil || kb == nil {
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
return nil, common.CodeDataError, errors.New("invalid Dataset ID")
}
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
newCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
exists, err := d.docEngine.ChunkStoreExists(ctx, indexName, datasetID)
exists, err := d.docEngine.ChunkStoreExists(newCtx, indexName, datasetID)
if err != nil {
return nil, common.CodeServerError, fmt.Errorf("failed to inspect chunk store: %w", err)
}
@@ -129,10 +129,10 @@ func (d *DatasetService) ListTags(ctx context.Context, datasetID, userID string)
const pageSize = 10000
counts := make(map[string]int)
for offset := 0; ; offset += pageSize {
if err = ctx.Err(); err != nil {
if err = newCtx.Err(); err != nil {
return nil, common.CodeServerError, fmt.Errorf("list tags timeout or canceled: %w", err)
}
searchResp, err := d.docEngine.Search(ctx, &enginetypes.SearchRequest{
searchResp, err := d.docEngine.Search(newCtx, &enginetypes.SearchRequest{
IndexNames: []string{indexName},
KbIDs: []string{datasetID},
Offset: offset,
@@ -197,17 +197,17 @@ func (d *DatasetService) RenameTag(ctx context.Context, datasetID, userID, fromT
return nil, common.CodeDataError, err
}
if strings.TrimSpace(datasetID) == "" {
return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"")
return nil, common.CodeDataError, errors.New("lack of \"Dataset ID\"")
}
if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) {
return nil, common.CodeDataError, errors.New("no authorization")
}
if d.docEngine == nil {
return nil, common.CodeServerError, errors.New("Document engine is not initialized")
return nil, common.CodeServerError, errors.New("document engine is not initialized")
}
kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID)
if err != nil || kb == nil {
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
return nil, common.CodeDataError, errors.New("invalid Dataset ID")
}
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
condition := map[string]interface{}{
@@ -222,7 +222,7 @@ func (d *DatasetService) RenameTag(ctx context.Context, datasetID, userID, fromT
"tag_kwd": toTag,
},
}
err = d.docEngine.UpdateChunks(context.Background(), condition, newValue, indexName, datasetID)
err = d.docEngine.UpdateChunks(ctx, condition, newValue, indexName, datasetID)
if err != nil {
return nil, common.CodeServerError, fmt.Errorf("failed to rename tag: %w", err)
}

10
internal/service/dataset/update.go
View File

@@ -352,7 +352,7 @@ func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, tenantID
}
if pagerankUpdate != nil {
if err = d.updateDatasetPagerankChunks(*pagerankUpdate); err != nil {
if err = d.updateDatasetPagerankChunks(ctx, *pagerankUpdate); err != nil {
return nil, common.CodeServerError, err
}
}
@@ -362,14 +362,14 @@ func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, tenantID
return data, common.CodeSuccess, nil
}
func (d *DatasetService) updateDatasetPagerankChunks(update datasetPagerankUpdate) error {
ctx, cancel := context.WithTimeout(context.Background(), datasetPagerankUpdateTimeout)
func (d *DatasetService) updateDatasetPagerankChunks(ctx context.Context, update datasetPagerankUpdate) error {
newCtx, cancel := context.WithTimeout(ctx, datasetPagerankUpdateTimeout)
defer cancel()
var err error
if update.value > 0 {
err = d.docEngine.UpdateChunks(ctx, map[string]interface{}{"kb_id": update.datasetID}, map[string]interface{}{common.PAGERANK_FLD: update.value}, update.index, update.datasetID)
err = d.docEngine.UpdateChunks(newCtx, map[string]interface{}{"kb_id": update.datasetID}, map[string]interface{}{common.PAGERANK_FLD: update.value}, update.index, update.datasetID)
} else {
err = d.docEngine.UpdateChunks(ctx, map[string]interface{}{"exists": common.PAGERANK_FLD}, map[string]interface{}{"remove": common.PAGERANK_FLD}, update.index, update.datasetID)
err = d.docEngine.UpdateChunks(newCtx, map[string]interface{}{"exists": common.PAGERANK_FLD}, map[string]interface{}{"remove": common.PAGERANK_FLD}, update.index, update.datasetID)
}
if errors.Is(err, types.ErrIndexNotFound) {
// Python's docStoreConn.update logs and returns False on a missing

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

@@ -283,13 +283,13 @@ type BatchUpdateMetadatasResponse struct {
// 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)
func removeObjectBestEffort(ctx context.Context, storageImpl storage.Storage, bucket, object string) error {
newCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
defer cancel()
var lastErr error
for attempt := 0; attempt < 3; attempt++ {
if err := storageImpl.Remove(ctx, bucket, object); err != nil {
if err := storageImpl.Remove(newCtx, 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) {

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

@@ -275,7 +275,7 @@ func (s *DocumentService) deleteDocumentFull(ctx context.Context, docID string)
common.Warn(fmt.Sprintf("need to delete files from taskInfo: %v", taskInfo))
}
s.deleteDocEngineData(docID, kb.TenantID, doc.KbID)
s.deleteDocEngineData(ctx, docID, kb.TenantID, doc.KbID)
if err = s.deleteDocRecordWithCounters(ctx, doc, kb.ID); err != nil {
return err
}
@@ -366,11 +366,10 @@ func (s *DocumentService) resolveDocAndKB(ctx context.Context, docID string) (*e
// deleteDocEngineData removes chunks and metadata from the document engine.
// No-op when the engine is nil.
func (s *DocumentService) deleteDocEngineData(docID, tenantID, kbID string) {
func (s *DocumentService) deleteDocEngineData(ctx context.Context, docID, tenantID, kbID string) {
if s.docEngine == nil {
return
}
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.Warn(fmt.Sprintf("deleteDocEngineData: failed to delete chunks for %s: %v", docID, delErr))
@@ -382,7 +381,7 @@ func (s *DocumentService) deleteDocEngineData(docID, tenantID, kbID string) {
// source/per-doc chunks, so the consumer only owns merged-product cleanup.
// Bound the publish with a timeout so a stalled scheduler (MySQL/NATS) can
// never block the document delete, which already succeeded above.
pubCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
pubCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
defer cancel()
if err := knowledge_compile.PublishDeleted(pubCtx, tenantID, kbID, docID); err != nil {
common.Warn(fmt.Sprintf("deleteDocEngineData: publish doc_deleted for %s failed: %v", docID, err))
@@ -490,7 +489,7 @@ func (s *DocumentService) cleanupFileReferences(ctx context.Context, docID strin
if file.Location != nil && *file.Location != "" {
storageImpl := storage.GetStorageFactory().GetStorage()
if storageImpl != nil {
rmErr := removeObjectBestEffort(storageImpl, file.ParentID, *file.Location)
rmErr := removeObjectBestEffort(ctx, 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))
}

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

@@ -80,8 +80,8 @@ func (s *DocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID,
hasError = true
continue
}
err := s.docEngine.UpdateChunks(
context.Background(),
err = s.docEngine.UpdateChunks(
ctx,
map[string]interface{}{"doc_id": docID},
map[string]interface{}{"available_int": statusInt},
fmt.Sprintf("ragflow_%s", kb.TenantID),
@@ -397,7 +397,7 @@ func (s *DocumentService) updateDocumentNameOnly(ctx context.Context, doc *entit
titleSmTks, _ := tokenizer.FineGrainedTokenize(titleTks)
indexName := fmt.Sprintf("ragflow_%s", tenantID)
return s.docEngine.UpdateChunks(
context.Background(),
ctx,
map[string]interface{}{"doc_id": doc.ID},
map[string]interface{}{
"docnm_kwd": newName,

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

@@ -136,13 +136,14 @@ func (s *DocumentService) Ingest(ctx context.Context, userID string, req *Ingest
}
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
if s.docEngine != nil {
exists, err := s.docEngine.ChunkStoreExists(context.Background(), indexName, doc.KbID)
var exists bool
exists, err = s.docEngine.ChunkStoreExists(ctx, indexName, doc.KbID)
if err != nil {
common.Error(fmt.Sprintf("go side, doc %s, ChunkStoreExists failed", doc.ID), err)
return common.CodeExceptionError, err
}
if exists {
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 {
common.Error(fmt.Sprintf("go side, doc %s, DeleteChunks failed", doc.ID), err)
return common.CodeExceptionError, err
}

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

@@ -96,7 +96,7 @@ 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.
rmErr := removeObjectBestEffort(storageImpl, kb.ID, location)
rmErr := removeObjectBestEffort(ctx, 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))
}
@@ -107,7 +107,7 @@ 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)
rmErr := removeObjectBestEffort(storageImpl, kb.ID, location)
rmErr := removeObjectBestEffort(ctx, storageImpl, kb.ID, location)
if rmErr != nil {
common.Warn(fmt.Sprintf("UploadLocalDocuments: failed to remove blob %s/%s: %v", kb.ID, location, rmErr))
}
@@ -285,7 +285,7 @@ func (s *DocumentService) UploadWebDocument(ctx context.Context, kb *entity.Know
doc := s.newDatasetDocument(kb, tenantID, filename, location, string(filetype), kb.ParserConfig, "web", int64(len(blob)), blob)
if err = s.InsertDocument(doc); err != nil {
rmErr := removeObjectBestEffort(storageImpl, kb.ID, location)
rmErr := removeObjectBestEffort(ctx, 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))
}
@@ -293,7 +293,7 @@ func (s *DocumentService) UploadWebDocument(ctx context.Context, kb *entity.Know
}
if err = s.addFileFromKB(ctx, doc, kbFolder.ID, kb.TenantID); err != nil {
err = s.rollbackAddFileFromKBError(ctx, doc, kb.ID, err)
rmErr := removeObjectBestEffort(storageImpl, kb.ID, location)
rmErr := removeObjectBestEffort(ctx, 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))
}

3
internal/service/mcp.go
View File

@@ -703,7 +703,7 @@ type TestServerRequest struct {
}
// TestServer opens a live MCP session and returns the tools the server advertises.
func (s *MCPService) TestServer(mcpID string, req *TestServerRequest) ([]map[string]interface{}, error) {
func (s *MCPService) TestServer(ctx context.Context, mcpID string, req *TestServerRequest) ([]map[string]interface{}, error) {
if req == nil || req.URL == "" {
return nil, fmt.Errorf("%w: Invalid MCP url", ErrMCPInvalidURL)
}
@@ -738,7 +738,6 @@ func (s *MCPService) TestServer(mcpID string, req *TestServerRequest) ([]map[str
vars[k] = sv
}
}
ctx := context.Background()
mcpCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()

7
internal/service/mcp_test.go
View File

@@ -90,19 +90,20 @@ func TestUpdateMCPServerRejectsDuplicatedName(t *testing.T) {
func TestServerInputValidation(t *testing.T) {
s := &MCPService{}
ctx := t.Context()
// Empty URL is rejected before any connection attempt.
if _, err := s.TestServer("id-1", &TestServerRequest{ServerType: mcpServerTypeSSE}); !errors.Is(err, ErrMCPInvalidURL) {
if _, err := s.TestServer(ctx, "id-1", &TestServerRequest{ServerType: mcpServerTypeSSE}); !errors.Is(err, ErrMCPInvalidURL) {
t.Errorf("expected ErrMCPInvalidURL for empty url, got %v", err)
}
// nil body is treated as empty URL.
if _, err := s.TestServer("id-1", nil); !errors.Is(err, ErrMCPInvalidURL) {
if _, err := s.TestServer(ctx, "id-1", nil); !errors.Is(err, ErrMCPInvalidURL) {
t.Errorf("expected ErrMCPInvalidURL for nil body, got %v", err)
}
// Invalid server type is rejected before connecting.
if _, err := s.TestServer("id-1", &TestServerRequest{URL: "http://example.com/sse", ServerType: "stdio"}); !errors.Is(err, ErrMCPInvalidType) {
if _, err := s.TestServer(ctx, "id-1", &TestServerRequest{URL: "http://example.com/sse", ServerType: "stdio"}); !errors.Is(err, ErrMCPInvalidType) {
t.Errorf("expected ErrMCPInvalidType for bad type, got %v", err)
}
}