From bdfc3ada41b41fcc3b56204c9513293c3864e2a2 Mon Sep 17 00:00:00 2001 From: Jin Hai Date: Fri, 24 Jul 2026 16:47:12 +0800 Subject: [PATCH] Go: add context (#17314) ### Summary As title. --------- Signed-off-by: Jin Hai --- internal/admin/handler.go | 2 +- internal/admin/service.go | 4 +- internal/dao/database.go | 2 +- internal/dao/document.go | 105 +++---- internal/dao/document_test.go | 39 ++- internal/handler/agent.go | 12 +- internal/handler/agent_attachment.go | 6 +- internal/handler/agent_download.go | 3 +- internal/handler/agent_download_test.go | 7 +- internal/handler/agent_upload.go | 5 +- internal/handler/bot_test.go | 6 +- internal/handler/chat.go | 3 +- internal/handler/chunk.go | 57 ++-- internal/handler/chunk_test.go | 60 ++-- internal/handler/dataset.go | 10 +- internal/handler/dify_retrieval_handler.go | 6 +- .../handler/dify_retrieval_handler_test.go | 11 +- internal/handler/document.go | 201 +++++++------ internal/handler/document_test.go | 69 ++--- internal/handler/file.go | 13 +- internal/handler/mindmap.go | 4 +- internal/handler/searchbot.go | 13 +- internal/handler/searchbot_test.go | 5 +- .../component/chunker/pdfcrop_cgo.go | 2 +- .../ingestion/component/document_storage.go | 8 +- internal/ingestion/component/file.go | 6 +- internal/ingestion/component/parser.go | 2 +- internal/ingestion/pipeline/pipeline.go | 14 +- internal/ingestion/pipeline/pipeline_test.go | 4 +- internal/ingestion/service/doc_state.go | 19 +- internal/ingestion/service/doc_state_test.go | 22 +- .../ingestion/service/execute_task_test.go | 3 +- .../ingestion/service/ingestion_service.go | 16 +- internal/ingestion/service/progress_sink.go | 11 +- .../ingestion/service/progress_sink_test.go | 26 +- internal/ingestion/task/pipeline_e2e_test.go | 3 +- .../ingestion/task/pipeline_executor_test.go | 4 +- internal/ingestion/task/task_context.go | 5 +- internal/ingestion/task/task_context_test.go | 14 +- .../parser/office_parsers_nocgo_test.go | 2 +- .../parser/parser/pdf_parser_nocgo_test.go | 3 +- internal/service/agent.go | 2 +- internal/service/ask_service.go | 4 +- internal/service/ask_service_test.go | 2 +- internal/service/chunk/chunk.go | 71 ++--- internal/service/chunk/chunk_test.go | 93 +++--- internal/service/chunk_types.go | 25 +- internal/service/dataset/crud.go | 50 ++-- internal/service/dataset/index.go | 38 +-- internal/service/dataset/ingestion.go | 5 +- internal/service/dataset/metadata.go | 9 +- .../service/dataset/metadata_config_test.go | 10 +- internal/service/dataset/task_cleanup_test.go | 5 +- internal/service/dataset/update.go | 6 +- .../service/document/document_artifact.go | 9 +- internal/service/document/document_crud.go | 70 ++--- .../document/document_dataset_update.go | 95 +++---- internal/service/document/document_ingest.go | 25 +- internal/service/document/document_list.go | 37 +-- .../service/document/document_metadata.go | 42 +-- internal/service/document/document_parse.go | 48 ++-- internal/service/document/document_test.go | 266 +++++++++++------- internal/service/document/document_upload.go | 45 +-- internal/service/document/file2document.go | 14 +- internal/service/file/file.go | 3 +- internal/service/file/file_commit.go | 2 +- internal/service/file/file_content.go | 6 +- internal/service/file/file_delete.go | 10 +- internal/service/file/file_folder.go | 19 +- internal/service/file/file_test.go | 12 +- internal/service/file/file_upload.go | 21 +- internal/service/file/file_url.go | 7 +- internal/service/ingestion_task_service.go | 26 +- .../service/ingestion_task_service_test.go | 24 +- internal/service/nlp/retrieval.go | 6 +- internal/service/user.go | 12 +- internal/utility/upload.go | 3 +- internal/utility/upload_test.go | 6 +- 78 files changed, 1031 insertions(+), 904 deletions(-) diff --git a/internal/admin/handler.go b/internal/admin/handler.go index 028d600455..c20d5c339b 100644 --- a/internal/admin/handler.go +++ b/internal/admin/handler.go @@ -309,7 +309,7 @@ func (h *Handler) DeleteUser(c *gin.Context) { return } - result, err := h.service.DeleteUser(username) + result, err := h.service.DeleteUser(c.Request.Context(), username) if err != nil { common.ErrorWithCode(c, common.CodeServerError, err.Error()) return diff --git a/internal/admin/service.go b/internal/admin/service.go index de48d3ae18..4d6a5be16c 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -494,7 +494,7 @@ type DeleteUserResult struct { // Returns: // - *DeleteUserResult // - error: error message -func (s *Service) DeleteUser(username string) (*DeleteUserResult, error) { +func (s *Service) DeleteUser(ctx context.Context, username string) (*DeleteUserResult, error) { result := &DeleteUserResult{ Username: username, DeletedDetails: []string{fmt.Sprintf("Drop user: %s", username)}, @@ -559,7 +559,7 @@ func (s *Service) DeleteUser(username string) (*DeleteUserResult, error) { if len(kbIDs) > 0 { // 2. Get document IDs - docIDs, err := s.documentDAO.GetAllDocIDsByKBIDs(kbIDs) + docIDs, err := s.documentDAO.GetAllDocIDsByKBIDs(ctx, tx, kbIDs) if err != nil { common.Warn("failed to get document IDs", zap.Error(err)) } diff --git a/internal/dao/database.go b/internal/dao/database.go index 863af2df52..97b235fd03 100644 --- a/internal/dao/database.go +++ b/internal/dao/database.go @@ -217,7 +217,7 @@ func GetModelProviderManager() *models.ProviderManager { if err != nil { common.Fatal("Failed to locate model providers", zap.Error(err)) } - if err := models.InitProviderManager(modelConfigDir); err != nil { + if err = models.InitProviderManager(modelConfigDir); err != nil { common.Fatal("Failed to load model providers", zap.Error(err)) } modelProviderManager = models.GetProviderManager() diff --git a/internal/dao/document.go b/internal/dao/document.go index 90e62c2652..4c18747ac2 100644 --- a/internal/dao/document.go +++ b/internal/dao/document.go @@ -17,6 +17,7 @@ package dao import ( + "context" "ragflow/internal/entity" "strings" @@ -31,15 +32,15 @@ func NewDocumentDAO() *DocumentDAO { return &DocumentDAO{} } -// Create create document -func (dao *DocumentDAO) Create(document *entity.Document) error { - return DB.Create(document).Error +// Create document +func (dao *DocumentDAO) Create(ctx context.Context, db *gorm.DB, document *entity.Document) error { + return db.WithContext(ctx).Create(document).Error } // GetByID get document by ID -func (dao *DocumentDAO) GetByID(id string) (*entity.Document, error) { +func (dao *DocumentDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.Document, error) { var document entity.Document - err := DB.First(&document, "id = ?", id).Error + err := db.WithContext(ctx).First(&document, "id = ?", id).Error if err != nil { return nil, err } @@ -47,11 +48,11 @@ func (dao *DocumentDAO) GetByID(id string) (*entity.Document, error) { } // GetByAuthorID get documents by author ID -func (dao *DocumentDAO) GetByAuthorID(authorID string, offset, limit int) ([]*entity.Document, int64, error) { +func (dao *DocumentDAO) GetByAuthorID(ctx context.Context, db *gorm.DB, authorID string, offset, limit int) ([]*entity.Document, int64, error) { var documents []*entity.Document var total int64 - query := DB.Model(&entity.Document{}).Where("created_by = ?", authorID) + query := db.WithContext(ctx).Model(&entity.Document{}).Where("created_by = ?", authorID) if err := query.Count(&total).Error; err != nil { return nil, 0, err } @@ -61,18 +62,18 @@ func (dao *DocumentDAO) GetByAuthorID(authorID string, offset, limit int) ([]*en } // Update update document -func (dao *DocumentDAO) Update(document *entity.Document) error { - return DB.Save(document).Error +func (dao *DocumentDAO) Update(ctx context.Context, db *gorm.DB, document *entity.Document) error { + return db.WithContext(ctx).Save(document).Error } // UpdateByID updates document by ID with the given fields -func (dao *DocumentDAO) UpdateByID(id string, updates map[string]interface{}) error { - return DB.Model(&entity.Document{}).Where("id = ?", id).Updates(updates).Error +func (dao *DocumentDAO) UpdateByID(ctx context.Context, db *gorm.DB, id string, updates map[string]interface{}) error { + return db.WithContext(ctx).Model(&entity.Document{}).Where("id = ?", id).Updates(updates).Error } // IncrementCounts atomically increments chunk_num, token_num, and process_duration for a document -func (dao *DocumentDAO) IncrementCounts(id string, kbID string, chunkNum int64, tokenNum int64, duration float64) error { - return DB.Model(&entity.Document{}). +func (dao *DocumentDAO) IncrementCounts(ctx context.Context, db *gorm.DB, id string, kbID string, chunkNum int64, tokenNum int64, duration float64) error { + return db.WithContext(ctx).Model(&entity.Document{}). Where("id = ? AND kb_id = ?", id, kbID). Updates(map[string]interface{}{ "chunk_num": gorm.Expr("chunk_num + ?", chunkNum), @@ -82,21 +83,21 @@ func (dao *DocumentDAO) IncrementCounts(id string, kbID string, chunkNum int64, } // Delete hard-deletes document by ID. Returns rows affected. -func (dao *DocumentDAO) Delete(id string) (int64, error) { - result := DB.Where("id = ?", id).Delete(&entity.Document{}) +func (dao *DocumentDAO) Delete(ctx context.Context, db *gorm.DB, id string) (int64, error) { + result := db.WithContext(ctx).Where("id = ?", id).Delete(&entity.Document{}) return result.RowsAffected, result.Error } -// List list documents -func (dao *DocumentDAO) List(offset, limit int) ([]*entity.Document, int64, error) { +// List documents +func (dao *DocumentDAO) List(ctx context.Context, db *gorm.DB, offset, limit int) ([]*entity.Document, int64, error) { var documents []*entity.Document var total int64 - if err := DB.Model(&entity.Document{}).Count(&total).Error; err != nil { + if err := db.WithContext(ctx).Model(&entity.Document{}).Count(&total).Error; err != nil { return nil, 0, err } - err := DB.Preload("Author").Offset(offset).Limit(limit).Find(&documents).Error + err := db.WithContext(ctx).Preload("Author").Offset(offset).Limit(limit).Find(&documents).Error return documents, total, err } @@ -119,8 +120,8 @@ type DocumentListOptions struct { } // ListByKBID list documents by knowledge base ID -func (dao *DocumentDAO) ListByKBID(kbID, keywords string, offset, limit int) ([]*entity.DocumentListItem, int64, error) { - return dao.ListByKBIDWithOptions(DocumentListOptions{ +func (dao *DocumentDAO) ListByKBID(ctx context.Context, db *gorm.DB, kbID, keywords string, offset, limit int) ([]*entity.DocumentListItem, int64, error) { + return dao.ListByKBIDWithOptions(ctx, db, DocumentListOptions{ KbID: kbID, Keywords: keywords, OrderBy: "create_time", @@ -131,11 +132,11 @@ func (dao *DocumentDAO) ListByKBID(kbID, keywords string, offset, limit int) ([] } // ListByKBIDWithOptions lists documents by knowledge base ID with filters. -func (dao *DocumentDAO) ListByKBIDWithOptions(opts DocumentListOptions) ([]*entity.DocumentListItem, int64, error) { +func (dao *DocumentDAO) ListByKBIDWithOptions(ctx context.Context, db *gorm.DB, opts DocumentListOptions) ([]*entity.DocumentListItem, int64, error) { var documents []*entity.DocumentListItem var total int64 - listQuery := DB.Table("document"). + listQuery := db.WithContext(ctx).Table("document"). Select(`document.*, user_canvas.title as pipeline_name, user.nickname`). Joins("JOIN file2document ON file2document.document_id = document.id"). Joins("JOIN file ON file.id = file2document.file_id"). @@ -143,7 +144,7 @@ func (dao *DocumentDAO) ListByKBIDWithOptions(opts DocumentListOptions) ([]*enti Joins("LEFT JOIN user ON document.created_by = user.id") listQuery = applyDocumentListFilters(listQuery, opts, true) - countQuery := applyDocumentListFilters(DB.Model(&entity.Document{}), opts, false) + countQuery := applyDocumentListFilters(db.WithContext(ctx).Model(&entity.Document{}), opts, false) if err := countQuery.Count(&total).Error; err != nil { return nil, 0, err @@ -165,14 +166,14 @@ func (dao *DocumentDAO) ListByKBIDWithOptions(opts DocumentListOptions) ([]*enti } // GetFilterByKBID returns aggregate filter counts for documents in a dataset. -func (dao *DocumentDAO) GetFilterByKBID(opts DocumentListOptions) (map[string]interface{}, int64, error) { +func (dao *DocumentDAO) GetFilterByKBID(ctx context.Context, db *gorm.DB, opts DocumentListOptions) (map[string]interface{}, int64, error) { var rows []struct { ID string `gorm:"column:id"` Run *string `gorm:"column:run"` Suffix string `gorm:"column:suffix"` } - query := DB.Table("document"). + query := db.WithContext(ctx).Table("document"). Select("document.id, document.run, document.suffix"). Joins("JOIN file2document ON file2document.document_id = document.id"). Joins("JOIN file ON file.id = file2document.file_id") @@ -201,9 +202,9 @@ func (dao *DocumentDAO) GetFilterByKBID(opts DocumentListOptions) (map[string]in } // ListIDsByKBIDWithOptions lists matching document IDs without pagination. -func (dao *DocumentDAO) ListIDsByKBIDWithOptions(opts DocumentListOptions) ([]string, error) { +func (dao *DocumentDAO) ListIDsByKBIDWithOptions(ctx context.Context, db *gorm.DB, opts DocumentListOptions) ([]string, error) { var ids []string - query := DB.Table("document"). + query := db.WithContext(ctx).Table("document"). Select("document.id"). Joins("JOIN file2document ON file2document.document_id = document.id"). Joins("JOIN file ON file.id = file2document.file_id") @@ -269,11 +270,11 @@ func documentListOrderColumn(orderBy string) string { } // GetByKBID retrieves all documents in a knowledge base ordered by create time. -func (dao *DocumentDAO) GetByKBID(kbID string) ([]*entity.Document, int64, error) { +func (dao *DocumentDAO) GetByKBID(ctx context.Context, db *gorm.DB, kbID string) ([]*entity.Document, int64, error) { var documents []*entity.Document var total int64 - query := DB.Model(&entity.Document{}).Where("kb_id = ?", kbID) + query := db.WithContext(ctx).Model(&entity.Document{}).Where("kb_id = ?", kbID) if err := query.Count(&total).Error; err != nil { return nil, 0, err } @@ -284,7 +285,7 @@ func (dao *DocumentDAO) GetByKBID(kbID string) ([]*entity.Document, int64, error // GetChunkingConfig returns the document, dataset, and tenant fields used to // build a parsing task digest, mirroring DocumentService.get_chunking_config. -func (dao *DocumentDAO) GetChunkingConfig(docID string) (map[string]interface{}, error) { +func (dao *DocumentDAO) GetChunkingConfig(ctx context.Context, db *gorm.DB, docID string) (map[string]interface{}, error) { var row struct { ID string `gorm:"column:id"` KbID string `gorm:"column:kb_id"` @@ -300,7 +301,7 @@ func (dao *DocumentDAO) GetChunkingConfig(docID string) (map[string]interface{}, LLMID string `gorm:"column:llm_id"` } - err := DB.Table("document"). + err := db.WithContext(ctx).Table("document"). Select(` document.id, document.kb_id, @@ -349,18 +350,18 @@ func (dao *DocumentDAO) GetChunkingConfig(docID string) (map[string]interface{}, } // DeleteByTenantID deletes all documents by tenant ID (hard delete) -func (dao *DocumentDAO) DeleteByTenantID(tenantID string) (int64, error) { - result := DB.Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Document{}) +func (dao *DocumentDAO) DeleteByTenantID(ctx context.Context, db *gorm.DB, tenantID string) (int64, error) { + result := db.WithContext(ctx).Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Document{}) return result.RowsAffected, result.Error } // GetAllDocIDsByKBIDs gets all document IDs by knowledge base IDs -func (dao *DocumentDAO) GetAllDocIDsByKBIDs(kbIDs []string) ([]map[string]string, error) { +func (dao *DocumentDAO) GetAllDocIDsByKBIDs(ctx context.Context, db *gorm.DB, kbIDs []string) ([]map[string]string, error) { var docs []struct { ID string `gorm:"column:id"` KbID string `gorm:"column:kb_id"` } - err := DB.Model(&entity.Document{}).Select("id, kb_id").Where("kb_id IN ?", kbIDs).Find(&docs).Error + err := db.WithContext(ctx).Model(&entity.Document{}).Select("id, kb_id").Where("kb_id IN ?", kbIDs).Find(&docs).Error if err != nil { return nil, err } @@ -373,12 +374,12 @@ func (dao *DocumentDAO) GetAllDocIDsByKBIDs(kbIDs []string) ([]map[string]string } // GetByIDs retrieves documents by multiple IDs -func (dao *DocumentDAO) GetByIDs(ids []string) ([]*entity.Document, error) { +func (dao *DocumentDAO) GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error) { if len(ids) == 0 { return nil, nil } var documents []*entity.Document - err := DB.Where("id IN ?", ids).Find(&documents).Error + err := db.WithContext(ctx).Model(&entity.Document{}).Where("id IN ?", ids).Find(&documents).Error if err != nil { return nil, err } @@ -386,12 +387,12 @@ func (dao *DocumentDAO) GetByIDs(ids []string) ([]*entity.Document, error) { } // GetByIDsAndTenantIDs retrieves documents by IDs scoped to knowledgebase owners. -func (dao *DocumentDAO) GetByIDsAndTenantIDs(ids, tenantIDs []string) ([]*entity.Document, error) { +func (dao *DocumentDAO) GetByIDsAndTenantIDs(ctx context.Context, db *gorm.DB, ids, tenantIDs []string) ([]*entity.Document, error) { if len(ids) == 0 || len(tenantIDs) == 0 { return nil, nil } var documents []*entity.Document - err := DB.Model(&entity.Document{}). + err := db.WithContext(ctx).Model(&entity.Document{}). Joins("JOIN knowledgebase ON document.kb_id = knowledgebase.id"). Where("document.id IN ? AND knowledgebase.tenant_id IN ? AND knowledgebase.status = ?", ids, tenantIDs, string(entity.StatusValid)). Find(&documents).Error @@ -402,23 +403,23 @@ func (dao *DocumentDAO) GetByIDsAndTenantIDs(ids, tenantIDs []string) ([]*entity } // GetByDocumentIDAndDatasetID retrieves a document by document ID and dataset/KB ID. -func (dao *DocumentDAO) GetByDocumentIDAndDatasetID(documentID, datasetID string) (*entity.Document, error) { +func (dao *DocumentDAO) GetByDocumentIDAndDatasetID(ctx context.Context, db *gorm.DB, documentID, datasetID string) (*entity.Document, error) { var document entity.Document - err := DB.Where("id = ? AND kb_id = ?", documentID, datasetID).First(&document).Error + err := db.WithContext(ctx).Where("id = ? AND kb_id = ?", documentID, datasetID).First(&document).Error return &document, err } // CountByTenantID counts documents by tenant ID -func (dao *DocumentDAO) CountByTenantID(tenantID string) (int64, error) { +func (dao *DocumentDAO) CountByTenantID(ctx context.Context, db *gorm.DB, tenantID string) (int64, error) { var count int64 - err := DB.Model(&entity.Document{}).Where("created_by = ?", tenantID).Count(&count).Error + err := db.WithContext(ctx).Model(&entity.Document{}).Where("created_by = ?", tenantID).Count(&count).Error return count, err } // SumSizeByDatasetID returns the total document size for a dataset. -func (dao *DocumentDAO) SumSizeByDatasetID(datasetID string) (int64, error) { +func (dao *DocumentDAO) SumSizeByDatasetID(ctx context.Context, db *gorm.DB, datasetID string) (int64, error) { var total int64 - err := DB.Model(&entity.Document{}). + err := db.WithContext(ctx).Model(&entity.Document{}). Select("COALESCE(SUM(size), 0)"). Where("kb_id = ?", datasetID). Scan(&total).Error @@ -427,7 +428,7 @@ func (dao *DocumentDAO) SumSizeByDatasetID(datasetID string) (int64, error) { // GetParsingStatusByKBID aggregates document parsing status counts for a // dataset, mirroring DocumentService.get_parsing_status_by_kb_ids in Python. -func (dao *DocumentDAO) GetParsingStatusByKBID(kbID string) (map[string]int64, error) { +func (dao *DocumentDAO) GetParsingStatusByKBID(ctx context.Context, db *gorm.DB, kbID string) (map[string]int64, error) { result := map[string]int64{ "unstart_count": 0, "running_count": 0, @@ -440,7 +441,7 @@ func (dao *DocumentDAO) GetParsingStatusByKBID(kbID string) (map[string]int64, e Run *string `gorm:"column:run"` Cnt int64 `gorm:"column:cnt"` } - err := DB.Model(&entity.Document{}). + err := db.WithContext(ctx).Model(&entity.Document{}). Select("run, COUNT(id) as cnt"). Where("kb_id = ?", kbID). Group("run"). @@ -467,17 +468,17 @@ func (dao *DocumentDAO) GetParsingStatusByKBID(kbID string) (map[string]int64, e return result, nil } -func (dao *DocumentDAO) GetByNameAndKBID(name, kbID string) ([]*entity.Document, error) { +func (dao *DocumentDAO) GetByNameAndKBID(ctx context.Context, db *gorm.DB, name, kbID string) ([]*entity.Document, error) { var docs []*entity.Document - err := DB.Where("name = ? AND kb_id = ?", name, kbID).Find(&docs).Error + err := db.WithContext(ctx).Where("name = ? AND kb_id = ?", name, kbID).Find(&docs).Error return docs, err } // ListNamesByKbID returns every document name in a dataset, used to compute a // non-colliding upload filename (mirrors Python duplicate_name). -func (dao *DocumentDAO) ListNamesByKbID(kbID string) ([]string, error) { +func (dao *DocumentDAO) ListNamesByKbID(ctx context.Context, db *gorm.DB, kbID string) ([]string, error) { var names []string - err := DB.Model(&entity.Document{}).Where("kb_id = ?", kbID).Pluck("name", &names).Error + err := db.WithContext(ctx).Model(&entity.Document{}).Where("kb_id = ?", kbID).Pluck("name", &names).Error if err != nil { return nil, err } diff --git a/internal/dao/document_test.go b/internal/dao/document_test.go index 484a94f3f0..88d2059fb2 100644 --- a/internal/dao/document_test.go +++ b/internal/dao/document_test.go @@ -52,14 +52,14 @@ func pushDocDB(t *testing.T, testDB *gorm.DB) { func TestDocumentGetByIDs_Success(t *testing.T) { db := setupDocumentTestDB(t) - pushDocDB(t, db) db.Create(&entity.Document{ID: "doc1", KbID: "kb1", Name: sp("Doc 1"), CreatedBy: "user1", ParserConfig: entity.JSONMap{}}) db.Create(&entity.Document{ID: "doc2", KbID: "kb1", Name: sp("Doc 2"), CreatedBy: "user1", ParserConfig: entity.JSONMap{}}) db.Create(&entity.Document{ID: "doc3", KbID: "kb2", Name: sp("Doc 3"), CreatedBy: "user2", ParserConfig: entity.JSONMap{}}) + ctx := t.Context() dao := NewDocumentDAO() - docs, err := dao.GetByIDs([]string{"doc1", "doc3"}) + docs, err := dao.GetByIDs(ctx, db, []string{"doc1", "doc3"}) if err != nil { t.Fatalf("GetByIDs failed: %v", err) } @@ -78,10 +78,9 @@ func TestDocumentGetByIDs_Success(t *testing.T) { func TestDocumentGetByIDs_EmptyIDs(t *testing.T) { db := setupDocumentTestDB(t) - pushDocDB(t, db) - + ctx := t.Context() dao := NewDocumentDAO() - docs, err := dao.GetByIDs([]string{}) + docs, err := dao.GetByIDs(ctx, db, []string{}) if err != nil { t.Fatalf("GetByIDs failed: %v", err) } @@ -92,10 +91,9 @@ func TestDocumentGetByIDs_EmptyIDs(t *testing.T) { func TestDocumentGetByIDs_NilIDs(t *testing.T) { db := setupDocumentTestDB(t) - pushDocDB(t, db) - + ctx := t.Context() dao := NewDocumentDAO() - docs, err := dao.GetByIDs(nil) + docs, err := dao.GetByIDs(ctx, db, nil) if err != nil { t.Fatalf("GetByIDs failed: %v", err) } @@ -106,12 +104,11 @@ func TestDocumentGetByIDs_NilIDs(t *testing.T) { func TestDocumentGetByIDs_NoMatch(t *testing.T) { db := setupDocumentTestDB(t) - pushDocDB(t, db) db.Create(&entity.Document{ID: "doc1", KbID: "kb1", Name: sp("Doc 1"), CreatedBy: "user1", ParserConfig: entity.JSONMap{}}) - + ctx := t.Context() dao := NewDocumentDAO() - docs, err := dao.GetByIDs([]string{"nonexistent"}) + docs, err := dao.GetByIDs(ctx, db, []string{"nonexistent"}) if err != nil { t.Fatalf("GetByIDs failed: %v", err) } @@ -122,7 +119,6 @@ func TestDocumentGetByIDs_NoMatch(t *testing.T) { func TestDocumentGetByKBIDOrdersByCreateTime(t *testing.T) { db := setupDocumentTestDB(t) - pushDocDB(t, db) createTime10 := int64(10) createTime20 := int64(20) @@ -130,8 +126,9 @@ func TestDocumentGetByKBIDOrdersByCreateTime(t *testing.T) { db.Create(&entity.Document{ID: "doc-later", KbID: "kb1", Name: sp("Doc Later"), CreatedBy: "user1", ParserConfig: entity.JSONMap{}, BaseModel: entity.BaseModel{CreateTime: &createTime30}}) db.Create(&entity.Document{ID: "doc-other", KbID: "kb2", Name: sp("Doc Other"), CreatedBy: "user1", ParserConfig: entity.JSONMap{}, BaseModel: entity.BaseModel{CreateTime: &createTime10}}) db.Create(&entity.Document{ID: "doc-earlier", KbID: "kb1", Name: sp("Doc Earlier"), CreatedBy: "user1", ParserConfig: entity.JSONMap{}, BaseModel: entity.BaseModel{CreateTime: &createTime20}}) - - docs, total, err := NewDocumentDAO().GetByKBID("kb1") + ctx := t.Context() + dao := NewDocumentDAO() + docs, total, err := dao.GetByKBID(ctx, db, "kb1") if err != nil { t.Fatalf("GetByKBID failed: %v", err) } @@ -148,12 +145,12 @@ func TestDocumentGetByKBIDOrdersByCreateTime(t *testing.T) { func TestDocumentGetByDocumentIDAndDatasetIDUsesKBID(t *testing.T) { db := setupDocumentTestDB(t) - pushDocDB(t, db) db.Create(&entity.Document{ID: "doc1", KbID: "kb1", Name: sp("Doc 1"), CreatedBy: "user1", ParserConfig: entity.JSONMap{}}) db.Create(&entity.Document{ID: "doc1-other", KbID: "kb2", Name: sp("Doc 2"), CreatedBy: "user1", ParserConfig: entity.JSONMap{}}) - - doc, err := NewDocumentDAO().GetByDocumentIDAndDatasetID("doc1", "kb1") + ctx := t.Context() + dao := NewDocumentDAO() + doc, err := dao.GetByDocumentIDAndDatasetID(ctx, db, "doc1", "kb1") if err != nil { t.Fatalf("GetByDocumentIDAndDatasetID failed: %v", err) } @@ -161,14 +158,13 @@ func TestDocumentGetByDocumentIDAndDatasetIDUsesKBID(t *testing.T) { t.Fatalf("unexpected document: id=%s kb_id=%s", doc.ID, doc.KbID) } - if _, err := NewDocumentDAO().GetByDocumentIDAndDatasetID("doc1", "kb2"); err == nil { + if _, err = dao.GetByDocumentIDAndDatasetID(ctx, db, "doc1", "kb2"); err == nil { t.Fatal("expected no match when document does not belong to dataset") } } func TestDocumentGetChunkingConfigScansParserConfig(t *testing.T) { db := setupDocumentTestDB(t) - pushDocDB(t, db) if err := db.Create(&entity.Tenant{ ID: "tenant1", @@ -207,8 +203,9 @@ func TestDocumentGetChunkingConfigScansParserConfig(t *testing.T) { }).Error; err != nil { t.Fatalf("create document: %v", err) } - - config, err := NewDocumentDAO().GetChunkingConfig("doc1") + ctx := t.Context() + dao := NewDocumentDAO() + config, err := dao.GetChunkingConfig(ctx, db, "doc1") if err != nil { t.Fatalf("GetChunkingConfig failed: %v", err) } diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 5e2c82dd8c..635587d989 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -48,15 +48,13 @@ import ( // the FileHandler (handler/file.go), not by any agent handler, so // the interface deliberately does NOT list it. (Code review CR1.) type agentFileService interface { - DownloadAgentFile(tenantID, location string) ([]byte, error) + DownloadAgentFile(ctx context.Context, tenantID, location string) ([]byte, error) // UploadInfos stores raw bytes in the per-user downloads bucket and - // returns lightweight descriptors. Mirrors python FileService.upload_info - // (multi-file path) used by the agent upload endpoint. - UploadInfos(userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) + // returns lightweight descriptors. + UploadInfos(ctx context.Context, userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) // UploadFromURL downloads a remote file (with SSRF protection) and - // stores it as an info blob. Mirrors python FileService.upload_info - // (single-file path with ?url=) used by the agent upload endpoint. - UploadFromURL(tenantID, rawURL string) (map[string]interface{}, error) + // stores it as an info blob. + UploadFromURL(ctx context.Context, tenantID, rawURL string) (map[string]interface{}, error) } // chatAgentService is the subset of AgentService used by the chat-completion diff --git a/internal/handler/agent_attachment.go b/internal/handler/agent_attachment.go index a06154a63b..70a0f38089 100644 --- a/internal/handler/agent_attachment.go +++ b/internal/handler/agent_attachment.go @@ -29,6 +29,7 @@ package handler import ( + "context" "net/http" "path/filepath" "strings" @@ -42,7 +43,7 @@ import ( // agentAttachmentFileService is the subset of FileService used by // the attachment-download handler. type agentAttachmentFileService interface { - DownloadAgentFile(tenantID, location string) ([]byte, error) + DownloadAgentFile(ctx context.Context, tenantID, location string) ([]byte, error) } // attachmentRequestMetadata holds the parsed query params used when @@ -81,7 +82,8 @@ func (h *AgentHandler) streamAgentAttachment(c *gin.Context, tenantID, attachmen return } - blob, err := h.fileService.DownloadAgentFile(tenantID, attachmentID) + ctx := c.Request.Context() + blob, err := h.fileService.DownloadAgentFile(ctx, tenantID, attachmentID) if err != nil { common.ResponseWithCodeData(c, common.CodeDataError, nil, "Attachment not found!") return diff --git a/internal/handler/agent_download.go b/internal/handler/agent_download.go index e4f155f286..1801a62f2c 100644 --- a/internal/handler/agent_download.go +++ b/internal/handler/agent_download.go @@ -54,6 +54,7 @@ func (h *AgentHandler) DownloadAgentFile(c *gin.Context) { return } + ctx := c.Request.Context() // IDOR note (security review H1): the Go User struct has no // TenantID field — the project collapses user and tenant in a // single-tenant session model. Python's @add_tenant_id_to_kwargs @@ -62,7 +63,7 @@ func (h *AgentHandler) DownloadAgentFile(c *gin.Context) { // per-object ownership check, so this port preserves the python // shape. A future per-object ownership check should be added in // both the python and Go code paths. - blob, err := h.fileService.DownloadAgentFile(user.ID, fileID) + blob, err := h.fileService.DownloadAgentFile(ctx, user.ID, fileID) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error()) return diff --git a/internal/handler/agent_download_test.go b/internal/handler/agent_download_test.go index 3f43e8df55..1b48fb3609 100644 --- a/internal/handler/agent_download_test.go +++ b/internal/handler/agent_download_test.go @@ -17,6 +17,7 @@ package handler import ( + "context" "errors" "mime/multipart" "net/http/httptest" @@ -43,19 +44,19 @@ type fakeAgentFileService struct { urlUploadErr error } -func (f *fakeAgentFileService) UploadInfos(_ string, _ []*multipart.FileHeader) ([]map[string]interface{}, error) { +func (f *fakeAgentFileService) UploadInfos(ctx context.Context, _ string, _ []*multipart.FileHeader) ([]map[string]interface{}, error) { if f.uploadErr != nil { return nil, f.uploadErr } return f.uploadList, nil } -func (f *fakeAgentFileService) UploadFromURL(_ string, _ string) (map[string]interface{}, error) { +func (f *fakeAgentFileService) UploadFromURL(ctx context.Context, _ string, _ string) (map[string]interface{}, error) { if f.urlUploadErr != nil { return nil, f.urlUploadErr } return f.urlUpload, nil } -func (f *fakeAgentFileService) DownloadAgentFile(_ string, _ string) ([]byte, error) { +func (f *fakeAgentFileService) DownloadAgentFile(ctx context.Context, _ string, _ string) ([]byte, error) { if f.err != nil { return nil, f.err } diff --git a/internal/handler/agent_upload.go b/internal/handler/agent_upload.go index 9d4a5d2435..987e164859 100644 --- a/internal/handler/agent_upload.go +++ b/internal/handler/agent_upload.go @@ -116,8 +116,9 @@ func (h *AgentHandler) UploadAgentFile(c *gin.Context) { // into the normal UploadInfos path. We replicate that with a // guard that dispatches to UploadFromURL only when both // conditions are met. + ctx := c.Request.Context() if url := c.Query("url"); url != "" && len(files) == 1 { - uploaded, err := h.fileService.UploadFromURL(user.ID, url) + uploaded, err := h.fileService.UploadFromURL(ctx, user.ID, url) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error()) return @@ -132,7 +133,7 @@ func (h *AgentHandler) UploadAgentFile(c *gin.Context) { return } - results, err := h.fileService.UploadInfos(user.ID, files) + results, err := h.fileService.UploadInfos(ctx, user.ID, files) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error()) return diff --git a/internal/handler/bot_test.go b/internal/handler/bot_test.go index 98fc6eaefa..5fa5aefdd4 100644 --- a/internal/handler/bot_test.go +++ b/internal/handler/bot_test.go @@ -711,15 +711,15 @@ type fakeFileService struct { err error } -func (f *fakeFileService) DownloadAgentFile(tenantID, location string) ([]byte, error) { +func (f *fakeFileService) DownloadAgentFile(ctx context.Context, tenantID, location string) ([]byte, error) { return f.blob, f.err } -func (f *fakeFileService) UploadInfos(userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) { +func (f *fakeFileService) UploadInfos(ctx context.Context, userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) { return nil, nil } -func (f *fakeFileService) UploadFromURL(tenantID, rawURL string) (map[string]interface{}, error) { +func (f *fakeFileService) UploadFromURL(ctx context.Context, tenantID, rawURL string) (map[string]interface{}, error) { return nil, nil } diff --git a/internal/handler/chat.go b/internal/handler/chat.go index 3bf711ad9c..aefe305718 100644 --- a/internal/handler/chat.go +++ b/internal/handler/chat.go @@ -202,7 +202,8 @@ func (h *ChatHandler) MindMap(c *gin.Context) { return } - mindMap, err := runMindMap(mindMapRunConfig{ + ctx := c.Request.Context() + mindMap, err := runMindMap(ctx, mindMapRunConfig{ Question: req.Question, KbIDs: kbIDs, SearchID: req.SearchID, diff --git a/internal/handler/chunk.go b/internal/handler/chunk.go index d21d105c7b..f9584c9bc9 100644 --- a/internal/handler/chunk.go +++ b/internal/handler/chunk.go @@ -33,15 +33,15 @@ import ( // chunkService is the consumer-side interface for ChunkHandler's service dependency. type chunkService interface { - RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) - Get(req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error) - List(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) - SwitchChunks(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error - UpdateChunk(req *service.UpdateChunkRequest, userID string) error - RemoveChunks(req *service.RemoveChunksRequest, userID string) (int64, error) - Parse(userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) + RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) + Get(ctx context.Context, req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error) + List(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) + SwitchChunks(ctx context.Context, userID string, datasetID, documentID string, availableInt int, chunkIDs []string) error + UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest, userID string) error + RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest, userID string) (int64, error) + Parse(ctx context.Context, userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) - StopParsing(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) + StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) } // ChunkHandler chunk handler @@ -62,8 +62,6 @@ func NewChunkHandler(chunkService chunkService, userService *service.UserService // @Summary Retrieval Test // @Description Test retrieval of chunks based on question and knowledge base // @Tags chunks -// @Accept json -// @Produce json // @Param request body service.RetrievalTestRequest true "retrieval test parameters" // @Success 200 {object} map[string]interface{} // @Router /api/v1/datasets/search [post] @@ -125,8 +123,9 @@ func (h *ChunkHandler) RetrievalTest(c *gin.Context) { return } + ctx := c.Request.Context() // Call service with user ID for permission checks - resp, err := h.chunkService.RetrievalTest(&req, user.ID) + resp, err := h.chunkService.RetrievalTest(ctx, &req, user.ID) if err != nil { common.Warn("dataset search failed", zap.String("error", err.Error())) common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, "dataset search failed") @@ -140,8 +139,6 @@ func (h *ChunkHandler) RetrievalTest(c *gin.Context) { // @Summary Get Chunk // @Description Retrieve a single chunk by its ID. // @Tags chunks -// @Accept json -// @Produce json // @Param dataset_id path string true "Dataset ID" // @Param document_id path string true "Document ID" // @Param chunk_id path string true "Chunk ID" @@ -164,7 +161,8 @@ func (h *ChunkHandler) Get(c *gin.Context) { ChunkID: chunkID, } - resp, err := h.chunkService.Get(req, user.ID) + ctx := c.Request.Context() + resp, err := h.chunkService.Get(ctx, req, user.ID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error()) return @@ -198,7 +196,8 @@ func (h *ChunkHandler) Parse(c *gin.Context) { return } - data, code, err := h.chunkService.Parse(userID, datasetId, &req) + ctx := c.Request.Context() + data, code, err := h.chunkService.Parse(ctx, userID, datasetId, &req) if code != common.CodeSuccess { common.ResponseWithCodeData(c, code, nil, err.Error()) return @@ -249,7 +248,8 @@ func (h *ChunkHandler) ListChunks(c *gin.Context) { req.AvailableInt = &available } - resp, err := h.chunkService.List(&req, user.ID) + ctx := c.Request.Context() + resp, err := h.chunkService.List(ctx, &req, user.ID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, err.Error()) return @@ -306,7 +306,8 @@ func (h *ChunkHandler) StopParsing(c *gin.Context) { return } - resp, code, err := h.chunkService.StopParsing(user.ID, datasetID, req) + ctx := c.Request.Context() + resp, code, err := h.chunkService.StopParsing(ctx, user.ID, datasetID, req) if err != nil { var data interface{} if resp != nil { @@ -332,8 +333,6 @@ func (h *ChunkHandler) StopParsing(c *gin.Context) { // @Summary List Chunks // @Description Retrieve paginated chunks for a document with optional filtering. // @Tags chunks -// @Accept json -// @Produce json // @Param request body service.ListChunksRequest true "List chunks parameters" // @Success 200 {object} map[string]interface{} // @Router /api/v1/chunk/list [post] @@ -361,7 +360,8 @@ func (h *ChunkHandler) List(c *gin.Context) { req.Size = &defaultSize } - resp, err := h.chunkService.List(&req, user.ID) + ctx := c.Request.Context() + resp, err := h.chunkService.List(ctx, &req, user.ID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error()) return @@ -420,7 +420,8 @@ func (h *ChunkHandler) SwitchChunks(c *gin.Context) { return } - if err = h.chunkService.SwitchChunks(userID, datasetID, documentID, availableInt, chunkIDs); err != nil { + ctx := c.Request.Context() + if err = h.chunkService.SwitchChunks(ctx, userID, datasetID, documentID, availableInt, chunkIDs); err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, err.Error()) return } @@ -476,18 +477,16 @@ func parseAvailableBody(rawBody map[string]interface{}) (int, error) { return 0, fmt.Errorf("available must be a boolean") } } - return 0, fmt.Errorf("`available_int` or `available` is required.") + return 0, fmt.Errorf("`available_int` or `available` is required") } // UpdateChunk updates a chunk // @Summary Update Chunk // @Description Update chunk fields // @Tags chunks -// @Accept json -// @Produce json // @Param request body service.UpdateChunkRequest true "update chunk" // @Success 200 {object} map[string]interface{} -// @Router /v1/chunk/update [post] +// @Router /api/v1/chunk/update [post] func (h *ChunkHandler) UpdateChunk(c *gin.Context) { user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { @@ -581,7 +580,8 @@ func (h *ChunkHandler) UpdateChunk(c *gin.Context) { req.DocumentID = documentID req.ChunkID = chunkID - err := h.chunkService.UpdateChunk(&req, user.ID) + ctx := c.Request.Context() + err := h.chunkService.UpdateChunk(ctx, &req, user.ID) if err != nil { var coded interface { Code() common.ErrorCode @@ -605,8 +605,6 @@ func (h *ChunkHandler) UpdateChunk(c *gin.Context) { // @Summary Remove Chunks // @Description Remove chunks from a document // @Tags chunks -// @Accept json -// @Produce json // @Param request body service.RemoveChunksRequest true "remove chunks request" // @Success 200 {object} map[string]interface{} // @Router /api/v1/datasets/{dataset_id}/documents/{document_id}/chunks [delete] @@ -637,7 +635,8 @@ func (h *ChunkHandler) RemoveChunks(c *gin.Context) { return } - deletedCount, err := h.chunkService.RemoveChunks(&req, user.ID) + ctx := c.Request.Context() + deletedCount, err := h.chunkService.RemoveChunks(ctx, &req, user.ID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error()) return diff --git a/internal/handler/chunk_test.go b/internal/handler/chunk_test.go index 8a0516475b..15f578879f 100644 --- a/internal/handler/chunk_test.go +++ b/internal/handler/chunk_test.go @@ -20,12 +20,12 @@ import ( // mockChunkSvc implements chunkSvcIface for testing ChunkHandler. // Only the methods actually called by the test are set; others panic. type mockChunkSvc struct { - retrievalTestFn func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) + retrievalTestFn func(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) addChunkFn func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) - listFn func(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) - switchChunksFn func(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error - updateChunkFn func(req *service.UpdateChunkRequest, userID string) error - stopParsingFn func(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) + listFn func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) + switchChunksFn func(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error + updateChunkFn func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error + stopParsingFn func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) } type codedTestError struct { @@ -41,46 +41,46 @@ func (e codedTestError) Code() common.ErrorCode { return e.code } -func (m *mockChunkSvc) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { +func (m *mockChunkSvc) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { if m.retrievalTestFn != nil { - return m.retrievalTestFn(req, userID) + return m.retrievalTestFn(ctx, req, userID) } return &service.RetrievalTestResponse{ Chunks: []map[string]interface{}{{"docnm_kwd": "test", "content_with_weight": "content"}}, Total: 1, }, nil } -func (m *mockChunkSvc) Get(*service.GetChunkRequest, string) (*service.GetChunkResponse, error) { +func (m *mockChunkSvc) Get(context.Context, *service.GetChunkRequest, string) (*service.GetChunkResponse, error) { panic("not implemented") } -func (m *mockChunkSvc) List(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { +func (m *mockChunkSvc) List(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { if m.listFn != nil { - return m.listFn(req, userID) + return m.listFn(ctx, req, userID) } panic("not implemented") } -func (m *mockChunkSvc) SwitchChunks(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error { +func (m *mockChunkSvc) SwitchChunks(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error { if m.switchChunksFn != nil { - return m.switchChunksFn(userID, datasetID, documentID, availableInt, chunkIDs) + return m.switchChunksFn(ctx, userID, datasetID, documentID, availableInt, chunkIDs) } panic("not implemented") } -func (m *mockChunkSvc) UpdateChunk(req *service.UpdateChunkRequest, userID string) error { +func (m *mockChunkSvc) UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest, userID string) error { if m.updateChunkFn != nil { - return m.updateChunkFn(req, userID) + return m.updateChunkFn(ctx, req, userID) } panic("not implemented") } -func (m *mockChunkSvc) RemoveChunks(*service.RemoveChunksRequest, string) (int64, error) { +func (m *mockChunkSvc) RemoveChunks(context.Context, *service.RemoveChunksRequest, string) (int64, error) { panic("not implemented") } -func (m *mockChunkSvc) StopParsing(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { +func (m *mockChunkSvc) StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { if m.stopParsingFn != nil { - return m.stopParsingFn(userID, datasetID, req) + return m.stopParsingFn(ctx, userID, datasetID, req) } panic("not implemented") } -func (m *mockChunkSvc) Parse(string, string, *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) { +func (m *mockChunkSvc) Parse(context.Context, string, string, *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) { panic("not implemented") } func (m *mockChunkSvc) AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) { @@ -139,7 +139,7 @@ func TestChunkHandlerListChunksMapsPathAndQuery(t *testing.T) { r, h := setupChunkHandlerWithUser("user-1", mock) r.GET("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.ListChunks) - mock.listFn = func(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { + mock.listFn = func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { if userID != "user-1" { t.Fatalf("userID = %q, want user-1", userID) } @@ -188,7 +188,7 @@ func TestChunkHandlerListChunksMapsAvailableFalse(t *testing.T) { r, h := setupChunkHandlerWithUser("user-1", mock) r.GET("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.ListChunks) - mock.listFn = func(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { + mock.listFn = func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { if userID != "user-1" { t.Fatalf("userID = %q, want user-1", userID) } @@ -216,7 +216,7 @@ func TestChunkHandlerSwitchChunksCallsService(t *testing.T) { r, h := setupChunkHandlerWithUser("user-1", mock) r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.SwitchChunks) - mock.switchChunksFn = func(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error { + mock.switchChunksFn = func(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error { if userID != "user-1" || datasetID != "kb-1" || documentID != "doc-1" { t.Fatalf("ids = %q/%q/%q, want user-1/kb-1/doc-1", userID, datasetID, documentID) } @@ -267,7 +267,7 @@ func TestChunkHandlerUpdateChunkUsesPathIDs(t *testing.T) { r, h := setupChunkHandlerWithUser("user-1", mock) r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk) - mock.updateChunkFn = func(req *service.UpdateChunkRequest, userID string) error { + mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error { if userID != "user-1" { t.Fatalf("userID = %q, want user-1", userID) } @@ -295,7 +295,7 @@ func TestChunkHandlerUpdateChunkValidationErrorIsBadRequest(t *testing.T) { r, h := setupChunkHandlerWithUser("user-1", mock) r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk) - mock.updateChunkFn = func(req *service.UpdateChunkRequest, userID string) error { + mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error { return codedTestError{code: common.CodeArgumentError, msg: "`tag_feas` values must be finite numbers greater than 0"} } @@ -343,7 +343,7 @@ func TestChunkRetrieval_EmptyQuestion(t *testing.T) { func TestChunkStopParsing_Success(t *testing.T) { r, mock := setupChunkStopParsingTest("user1") - mock.stopParsingFn = func(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { + mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { if userID != "user1" { t.Fatalf("expected user1, got %q", userID) } @@ -378,7 +378,7 @@ func TestChunkStopParsing_Success(t *testing.T) { func TestChunkStopParsingRouteRequiresDocumentIDs(t *testing.T) { r, mock := setupChunkStopParsingTest("user1") - mock.stopParsingFn = func(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { + mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { t.Fatal("service should not be called when document_ids is missing") return nil, common.CodeSuccess, nil } @@ -405,10 +405,10 @@ func TestChunkStopParsingRouteRequiresDocumentIDs(t *testing.T) { func TestChunkStopParsing_InvalidStateIncludesPythonErrorCode(t *testing.T) { r, mock := setupChunkStopParsingTest("user1") - mock.stopParsingFn = func(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { + mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { return &service.StopParsingResponse{ Data: map[string]interface{}{"error_code": "DOC_STOP_PARSING_INVALID_STATE"}, - }, common.CodeDataError, errors.New("Can't stop parsing document that has not started or already completed") + }, common.CodeDataError, errors.New("can't stop parsing document that has not started or already completed") } w := httptest.NewRecorder() @@ -433,7 +433,7 @@ func TestChunkStopParsing_InvalidStateIncludesPythonErrorCode(t *testing.T) { if data["error_code"] != "DOC_STOP_PARSING_INVALID_STATE" { t.Fatalf("unexpected error_code: %v", data["error_code"]) } - if resp["message"] != "Can't stop parsing document that has not started or already completed" { + if resp["message"] != "can't stop parsing document that has not started or already completed" { t.Fatalf("unexpected message: %v", resp["message"]) } } @@ -552,7 +552,7 @@ func TestChunkRetrieval_InvalidJSON(t *testing.T) { func TestChunkRetrieval_Success(t *testing.T) { _, mock := setupChunkRetrievalTest("user1") - mock.retrievalTestFn = func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { + mock.retrievalTestFn = func(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { if userID != "user1" { t.Errorf("expected userID 'user1', got %q", userID) } @@ -598,7 +598,7 @@ func TestChunkRetrieval_Success(t *testing.T) { func TestChunkRetrieval_ServiceError(t *testing.T) { _, mock := setupChunkRetrievalTest("user1") - mock.retrievalTestFn = func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { + mock.retrievalTestFn = func(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { return nil, errors.New("db connection refused") } diff --git a/internal/handler/dataset.go b/internal/handler/dataset.go index 28dbb29670..b21511daeb 100644 --- a/internal/handler/dataset.go +++ b/internal/handler/dataset.go @@ -284,9 +284,9 @@ func (h *DatasetsHandler) GetIngestionSummary(c *gin.Context) { common.ErrorWithCode(c, errorCode, errorMessage) return } - + ctx := c.Request.Context() datasetID := c.Param("dataset_id") - result, code, err := h.datasetsService.GetIngestionSummary(datasetID, user.ID) + result, code, err := h.datasetsService.GetIngestionSummary(ctx, datasetID, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -784,8 +784,9 @@ func (h *DatasetsHandler) RunIndex(c *gin.Context) { return } + ctx := c.Request.Context() indexType := strings.ToLower(strings.TrimSpace(c.Query("type"))) - data, code, err := h.datasetsService.RunIndex(userID, datasetID, indexType) + data, code, err := h.datasetsService.RunIndex(ctx, userID, datasetID, indexType) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -950,7 +951,8 @@ func (h *DatasetsHandler) UpdateDocumentMetadataConfig(c *gin.Context) { return } - data, code, err := h.datasetsService.UpdateDocumentMetadataConfig(userID, datasetID, documentID, req) + ctx := c.Request.Context() + data, code, err := h.datasetsService.UpdateDocumentMetadataConfig(ctx, userID, datasetID, documentID, req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return diff --git a/internal/handler/dify_retrieval_handler.go b/internal/handler/dify_retrieval_handler.go index 2c1794bae7..948fb4e52e 100644 --- a/internal/handler/dify_retrieval_handler.go +++ b/internal/handler/dify_retrieval_handler.go @@ -21,6 +21,7 @@ import ( "errors" "fmt" "net/http" + "ragflow/internal/dao" "strconv" "strings" @@ -66,7 +67,7 @@ type RetrievalServiceIface interface { // DocumentDAOIface abstracts DocumentDAO for the Dify handler. type DocumentDAOIface interface { - GetByIDs(ids []string) ([]*entity.Document, error) + GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error) } // --- Request / Response types --- @@ -321,10 +322,11 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { allDocIDs = append(allDocIDs, id) } + ctx := c.Request.Context() docMap := make(map[string]*entity.Document) if len(allDocIDs) > 0 { var docs []*entity.Document - docs, err = h.docDAO.GetByIDs(allDocIDs) + docs, err = h.docDAO.GetByIDs(ctx, dao.DB, allDocIDs) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, fmt.Sprintf("failed to load documents: %v", err)) return diff --git a/internal/handler/dify_retrieval_handler_test.go b/internal/handler/dify_retrieval_handler_test.go index 7dc46cec78..55c5666db6 100644 --- a/internal/handler/dify_retrieval_handler_test.go +++ b/internal/handler/dify_retrieval_handler_test.go @@ -20,12 +20,13 @@ import ( "context" "encoding/json" "errors" - "gorm.io/gorm" "net/http" "net/http/httptest" "strings" "testing" + "gorm.io/gorm" + "ragflow/internal/common" "ragflow/internal/engine" "ragflow/internal/engine/types" @@ -131,12 +132,12 @@ func (m *mockRetrievalService) Retrieval(ctx context.Context, req *nlp.Retrieval type mockDocDAO struct { DocumentDAOIface - getByIDsFn func(ids []string) ([]*entity.Document, error) + getByIDsFn func(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error) } -func (m *mockDocDAO) GetByIDs(ids []string) ([]*entity.Document, error) { +func (m *mockDocDAO) GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error) { if m.getByIDsFn != nil { - return m.getByIDsFn(ids) + return m.getByIDsFn(ctx, db, ids) } return []*entity.Document{ {ID: "doc1", Name: strPtr("Test Doc")}, @@ -382,7 +383,7 @@ func TestDifyRetrieval_KBDBError(t *testing.T) { func TestDifyRetrieval_DocLoadError(t *testing.T) { h, r := setupDifyTest("user1") h.docDAO = &mockDocDAO{ - getByIDsFn: func(ids []string) ([]*entity.Document, error) { + getByIDsFn: func(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error) { return nil, errors.New("db unavailable") }, } diff --git a/internal/handler/document.go b/internal/handler/document.go index fc73475bfa..4e67d7579d 100644 --- a/internal/handler/document.go +++ b/internal/handler/document.go @@ -18,6 +18,7 @@ package handler import ( "bytes" + "context" "encoding/json" "errors" "fmt" @@ -38,7 +39,7 @@ import ( "ragflow/internal/dao" "ragflow/internal/service" - dataset "ragflow/internal/service/dataset" + "ragflow/internal/service/dataset" "ragflow/internal/service/document" ) @@ -46,46 +47,46 @@ var IMG_BASE64_PREFIX = "data:image/png;base64," // documentServiceIface defines the DocumentService methods used by DocumentHandler. type documentServiceIface interface { - GetDocumentByID(id string) (*document.DocumentResponse, error) - UpdateDocument(id string, req *document.UpdateDocumentRequest) error - DeleteDocument(id string) error - DeleteDocuments(ids []string, deleteAll bool, datasetID, userID string) (int, error) - ParseDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) - StopParseDocuments(datasetID string, docIDs []string) (map[string]interface{}, error) - ListDocuments(page, pageSize int) ([]*document.DocumentResponse, int64, error) - ListDocumentsByDatasetID(kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) - ListDocumentsByDatasetIDWithOptions(opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) - ListDocumentIDsByDatasetIDWithOptions(opts dao.DocumentListOptions) ([]string, error) - GetDocumentFiltersByDatasetID(opts dao.DocumentListOptions) (map[string]interface{}, int64, error) - GetMetadataByKBs(kbIDs []string) (map[string]interface{}, error) - GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error) - GetThumbnails(userID string, docIDs []string) (map[string]string, error) - GetDocumentImage(imageID string) ([]byte, error) - GetMetadataSummary(kbID string, docIDs []string) (map[string]interface{}, error) - SetDocumentMetadata(docID string, meta map[string]interface{}) error - DeleteDocumentMetadata(docID string, keys []string) error - DeleteDocumentAllMetadata(docID string) error - GetDocumentMetadataByID(docID string) (map[string]interface{}, error) - GetDocumentArtifact(filename, userID string) (*document.ArtifactResponse, error) - GetDocumentPreview(docID string) (*document.DocumentPreview, error) - UploadLocalDocuments(kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) - UploadWebDocument(kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) - UploadEmptyDocument(kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) - DownloadDocument(datasetID, docID string) (*document.DownloadDocumentResp, error) - UpdateDatasetDocument(userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) - BatchUpdateDocumentMetadatas(datasetID string, selector *document.DocumentMetadataSelector, updates []document.DocumentMetadataUpdate, deletes []document.DocumentMetadataDelete) (*document.BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error) - ListIngestionTasks(userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) - IngestDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) - StopIngestionTasks(tasks []string, userID string) ([]*entity.IngestionTask, error) - Ingest(userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error) - RemoveIngestionTasks(tasks []string, userID string) ([]map[string]string, error) - BatchUpdateDocumentStatus(userID, datasetID, status string, DocumentIDs []string) (map[string]interface{}, common.ErrorCode, error) + GetDocumentByID(ctx context.Context, id string) (*document.DocumentResponse, error) + UpdateDocument(ctx context.Context, id string, req *document.UpdateDocumentRequest) error + DeleteDocument(ctx context.Context, id string) error + DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error) + ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) + StopParseDocuments(ctx context.Context, datasetID string, docIDs []string) (map[string]interface{}, error) + ListDocuments(ctx context.Context, page, pageSize int) ([]*document.DocumentResponse, int64, error) + ListDocumentsByDatasetID(ctx context.Context, kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) + ListDocumentsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) + ListDocumentIDsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions) ([]string, error) + GetDocumentFiltersByDatasetID(ctx context.Context, opts dao.DocumentListOptions) (map[string]interface{}, int64, error) + GetMetadataByKBs(ctx context.Context, kbIDs []string) (map[string]interface{}, error) + GetDocumentsByAuthorID(ctx context.Context, authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error) + GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error) + GetDocumentImage(ctx context.Context, imageID string) ([]byte, error) + GetMetadataSummary(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error) + SetDocumentMetadata(ctx context.Context, docID string, meta map[string]interface{}) error + DeleteDocumentMetadata(ctx context.Context, docID string, keys []string) error + DeleteDocumentAllMetadata(ctx context.Context, docID string) error + GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]interface{}, error) + GetDocumentArtifact(ctx context.Context, filename, userID string) (*document.ArtifactResponse, error) + GetDocumentPreview(ctx context.Context, docID string) (*document.DocumentPreview, error) + UploadLocalDocuments(ctx context.Context, kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) + UploadWebDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) + UploadEmptyDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) + DownloadDocument(ctx context.Context, datasetID, docID string) (*document.DownloadDocumentResp, error) + UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) + BatchUpdateDocumentMetadatas(ctx context.Context, datasetID string, selector *document.DocumentMetadataSelector, updates []document.DocumentMetadataUpdate, deletes []document.DocumentMetadataDelete) (*document.BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error) + ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) + IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) + StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error) + Ingest(ctx context.Context, userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error) + RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error) + BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, DocumentIDs []string) (map[string]interface{}, common.ErrorCode, error) } // fileUploadIface defines the FileService upload methods used by DocumentHandler. type fileUploadIface interface { - UploadDocumentInfos(userID string, files []*multipart.FileHeader) ([]map[string]interface{}, common.ErrorCode, error) - UploadDocumentInfoByURL(userID, rawURL string) (map[string]interface{}, common.ErrorCode, error) + UploadDocumentInfos(ctx context.Context, userID string, files []*multipart.FileHeader) ([]map[string]interface{}, common.ErrorCode, error) + UploadDocumentInfoByURL(ctx context.Context, userID, rawURL string) (map[string]interface{}, common.ErrorCode, error) } // DocumentHandler document handler @@ -108,8 +109,6 @@ func NewDocumentHandler(documentService documentServiceIface, datasetService *da // @Summary Get Document Info // @Description Get document details by ID // @Tags documents -// @Accept json -// @Produce json // @Param id path int true "document ID" // @Success 200 {object} map[string]interface{} // @Router /api/v1/documents/{id} [get] @@ -128,7 +127,8 @@ func (h *DocumentHandler) GetDocumentByID(c *gin.Context) { return } - document, err := h.documentService.GetDocumentByID(id) + ctx := c.Request.Context() + doc, err := h.documentService.GetDocumentByID(ctx, id) if err != nil { c.JSON(http.StatusNotFound, gin.H{ "error": "document not found", @@ -137,7 +137,7 @@ func (h *DocumentHandler) GetDocumentByID(c *gin.Context) { } c.JSON(http.StatusOK, gin.H{ - "data": document, + "data": doc, }) } @@ -155,7 +155,8 @@ func (h *DocumentHandler) GetThumbnail(c *gin.Context) { return } - result, err := h.documentService.GetThumbnails(user.ID, docIDs) + ctx := c.Request.Context() + result, err := h.documentService.GetThumbnails(ctx, user.ID, docIDs) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error()) return @@ -187,7 +188,8 @@ func parseThumbnailDocIDs(c *gin.Context) []string { // GetDocumentImage returns a document image from object storage. func (h *DocumentHandler) GetDocumentImage(c *gin.Context) { imageID := c.Param("image_id") - data, err := h.documentService.GetDocumentImage(imageID) + ctx := c.Request.Context() + data, err := h.documentService.GetDocumentImage(ctx, imageID) if err != nil { common.ResponseWithCodeData(c, common.CodeDataError, nil, "Image not found.") return @@ -224,7 +226,8 @@ func (h *DocumentHandler) GetDocumentArtifact(c *gin.Context) { return } filename := c.Param("filename") - artifact, err := h.documentService.GetDocumentArtifact(filename, user.ID) + ctx := c.Request.Context() + artifact, err := h.documentService.GetDocumentArtifact(ctx, filename, user.ID) if err != nil { switch { case errors.Is(err, document.ErrArtifactInvalidFilename), @@ -256,7 +259,8 @@ func (h *DocumentHandler) GetDocumentPreview(c *gin.Context) { return } - preview, err := h.documentService.GetDocumentPreview(docID) + ctx := c.Request.Context() + preview, err := h.documentService.GetDocumentPreview(ctx, docID) if err != nil { common.ErrorWithCode(c, common.CodeDataError, "Document not found!") return @@ -298,7 +302,8 @@ func (h *DocumentHandler) UpdateDocument(c *gin.Context) { return } - doc, err := h.documentService.GetDocumentByID(id) + ctx := c.Request.Context() + doc, err := h.documentService.GetDocumentByID(ctx, id) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found!") return @@ -316,7 +321,7 @@ func (h *DocumentHandler) UpdateDocument(c *gin.Context) { return } - if err = h.documentService.UpdateDocument(id, &req); err != nil { + if err = h.documentService.UpdateDocument(ctx, id, &req); err != nil { c.JSON(http.StatusInternalServerError, gin.H{ "error": err.Error(), }) @@ -350,7 +355,8 @@ func (h *DocumentHandler) DeleteDocument(c *gin.Context) { return } - doc, err := h.documentService.GetDocumentByID(id) + ctx := c.Request.Context() + doc, err := h.documentService.GetDocumentByID(ctx, id) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found!") return @@ -360,7 +366,7 @@ func (h *DocumentHandler) DeleteDocument(c *gin.Context) { return } - if err = h.documentService.DeleteDocument(id); err != nil { + if err = h.documentService.DeleteDocument(ctx, id); err != nil { c.JSON(http.StatusInternalServerError, gin.H{ "error": err.Error(), }) @@ -411,7 +417,8 @@ func (h *DocumentHandler) DeleteDocuments(c *gin.Context) { } userID := c.GetString("user_id") - deleted, err := h.documentService.DeleteDocuments(ids, req.DeleteAll, datasetID, userID) + ctx := c.Request.Context() + deleted, err := h.documentService.DeleteDocuments(ctx, ids, req.DeleteAll, datasetID, userID) if err != nil { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) return @@ -473,7 +480,8 @@ func (h *DocumentHandler) BatchUpdateDocumentStatus(c *gin.Context) { return } - result, code, err := h.documentService.BatchUpdateDocumentStatus(userID, datasetID, status, documentIDs) + ctx := c.Request.Context() + result, code, err := h.documentService.BatchUpdateDocumentStatus(ctx, userID, datasetID, status, documentIDs) if err != nil { message := err.Error() if code == common.CodeServerError { @@ -520,8 +528,9 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) { return } + ctx := c.Request.Context() if c.Query("type") == "filter" { - filters, total, err := h.documentService.GetDocumentFiltersByDatasetID(opts) + filters, total, err := h.documentService.GetDocumentFiltersByDatasetID(ctx, opts) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, map[string]interface{}{"total": 0, "filter": map[string]interface{}{}}, "failed to get document filters") return @@ -533,7 +542,7 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) { // Use kbID to filter documents. // Note: create_time_from / create_time_to are applied post-query (not at DB level) // so that total reflects the unfiltered count, matching the Python API contract. - documents, total, err := h.documentService.ListDocumentsByDatasetIDWithOptions(opts, page, pageSize) + documents, total, err := h.documentService.ListDocumentsByDatasetIDWithOptions(ctx, opts, page, pageSize) if err != nil { common.ResponseWithCodeData(c, 1, map[string]interface{}{"total": 0, "docs": []interface{}{}}, "failed to get documents") return @@ -547,7 +556,7 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) { if opts.CreateTimeTo > 0 && doc.CreateTime != nil && *doc.CreateTime > opts.CreateTimeTo { continue } - metaFields, err := h.documentService.GetDocumentMetadataByID(doc.ID) + metaFields, err := h.documentService.GetDocumentMetadataByID(ctx, doc.ID) if err != nil { metaFields = make(map[string]interface{}) } @@ -615,13 +624,15 @@ func (h *DocumentHandler) applyDocumentMetadataFilter(c *gin.Context, opts dao.D return opts, "" } - candidateIDs, err := h.documentService.ListDocumentIDsByDatasetIDWithOptions(opts) + ctx := c.Request.Context() + + candidateIDs, err := h.documentService.ListDocumentIDsByDatasetIDWithOptions(ctx, opts) if err != nil { return opts, "failed to get documents" } candidateSet := stringSet(candidateIDs) - metadataByKey, err := h.documentService.GetMetadataByKBs([]string{opts.KbID}) + metadataByKey, err := h.documentService.GetMetadataByKBs(ctx, []string{opts.KbID}) if err != nil { return opts, err.Error() } @@ -879,8 +890,8 @@ func (h *DocumentHandler) uploadLocalDocuments(c *gin.Context, kb *entity.Knowle } } } - - data, errMsgs := h.documentService.UploadLocalDocuments(kb, tenantID, files, c.PostForm("parent_path"), override) + ctx := c.Request.Context() + data, errMsgs := h.documentService.UploadLocalDocuments(ctx, kb, tenantID, files, c.PostForm("parent_path"), override) if len(data) == 0 && len(errMsgs) > 0 { common.ResponseWithCodeData(c, common.CodeServerError, nil, strings.Join(errMsgs, "\n")) return @@ -929,7 +940,8 @@ func (h *DocumentHandler) uploadEmptyDocument(c *gin.Context, kb *entity.Knowled common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "File name must be 255 bytes or less.") return } - data, code, err := h.documentService.UploadEmptyDocument(kb, tenantID, name) + ctx := c.Request.Context() + data, code, err := h.documentService.UploadEmptyDocument(ctx, kb, tenantID, name) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -956,7 +968,8 @@ func (h *DocumentHandler) uploadWebDocument(c *gin.Context, kb *entity.Knowledge common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "The URL format is invalid") return } - data, code, err := h.documentService.UploadWebDocument(kb, tenantID, name, rawURL) + ctx := c.Request.Context() + data, code, err := h.documentService.UploadWebDocument(ctx, kb, tenantID, name, rawURL) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -1005,8 +1018,8 @@ func (h *DocumentHandler) DownloadDocument(c *gin.Context) { common.ErrorWithCode(c, common.CodeDataError, fmt.Sprintf("The dataset not own the document %s.", docID)) return } - - res, err := h.documentService.DownloadDocument(datasetID, docID) + ctx := c.Request.Context() + res, err := h.documentService.DownloadDocument(ctx, datasetID, docID) if err != nil { common.ErrorWithCode(c, common.CodeDataError, err.Error()) @@ -1140,7 +1153,8 @@ func (h *DocumentHandler) MetadataSummary(c *gin.Context) { return } - summary, err := h.documentService.GetMetadataSummary(kbID, requestBody.DocIDs) + ctx := c.Request.Context() + summary, err := h.documentService.GetMetadataSummary(ctx, kbID, requestBody.DocIDs) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 1, nil, "Failed to get metadata summary: "+err.Error()) return @@ -1216,7 +1230,8 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) { } // Authorization: user must be able to access the document's dataset. - doc, err := h.documentService.GetDocumentByID(req.DocID) + ctx := c.Request.Context() + doc, err := h.documentService.GetDocumentByID(ctx, req.DocID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found") return @@ -1226,7 +1241,7 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) { return } - err = h.documentService.SetDocumentMetadata(req.DocID, meta) + err = h.documentService.SetDocumentMetadata(ctx, req.DocID, meta) if err != nil { errMsg := err.Error() if strings.Contains(errMsg, "no such document") || strings.Contains(errMsg, "document not found") { @@ -1244,8 +1259,6 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) { // @Summary Ingest Document // @Description Ingest a document for processing // @Tags documents -// @Accept json -// @Produce json // @Security ApiKeyAuth // @Param request body document.IngestDocumentRequest true "ingestion info" // @Success 200 {object} map[string]interface{} @@ -1268,8 +1281,8 @@ func (h *DocumentHandler) Ingest(c *gin.Context) { common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) return } - - if code, err := h.documentService.Ingest(userID, &req); err != nil { + ctx := c.Request.Context() + if code, err := h.documentService.Ingest(ctx, userID, &req); err != nil { common.ErrorWithCode(c, code, err.Error()) return } @@ -1293,7 +1306,7 @@ type DeleteMetaRequest struct { // @Security ApiKeyAuth // @Param request body DeleteMetaRequest true "metadata keys to delete or empty to delete all" // @Success 200 {object} map[string]interface{} -// @Router /v1/document/delete_meta [post] +// @Router /api/v1/document/delete_meta [post] func (h *DocumentHandler) DeleteMeta(c *gin.Context) { user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { @@ -1312,8 +1325,8 @@ func (h *DocumentHandler) DeleteMeta(c *gin.Context) { return } - // Authorization: user must be able to access the document's dataset. - doc, err := h.documentService.GetDocumentByID(req.DocID) + ctx := c.Request.Context() + doc, err := h.documentService.GetDocumentByID(ctx, req.DocID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found") return @@ -1337,7 +1350,7 @@ func (h *DocumentHandler) DeleteMeta(c *gin.Context) { return } - err = h.documentService.DeleteDocumentMetadata(req.DocID, keys) + err = h.documentService.DeleteDocumentMetadata(ctx, req.DocID, keys) if err != nil { errMsg := err.Error() if strings.Contains(errMsg, "no such document") || strings.Contains(errMsg, "document not found") { @@ -1349,7 +1362,7 @@ func (h *DocumentHandler) DeleteMeta(c *gin.Context) { } } else { // Delete entire document metadata - err = h.documentService.DeleteDocumentAllMetadata(req.DocID) + err = h.documentService.DeleteDocumentAllMetadata(ctx, req.DocID) if err != nil { errMsg := err.Error() if strings.Contains(errMsg, "no such document") || strings.Contains(errMsg, "document not found") { @@ -1385,8 +1398,8 @@ func (h *DocumentHandler) ListIngestionTasks(c *gin.Context) { return } } - - parseResult, err = h.documentService.ListIngestionTasks(userID, req.DatasetID, 0, 0) + ctx := c.Request.Context() + parseResult, err = h.documentService.ListIngestionTasks(ctx, userID, req.DatasetID, 0, 0) if err != nil { common.ResponseWithCodeData(c, IngestionTaskErrorCode(err), nil, err.Error()) return @@ -1414,8 +1427,8 @@ func (h *DocumentHandler) StartIngestionTask(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeDataError, nil, fmt.Sprintf("You don't own the dataset %s.", datasetID)) return } - - parseResult, err := h.documentService.IngestDocuments(datasetID, userID, req.DocumentIDs) + ctx := c.Request.Context() + parseResult, err := h.documentService.IngestDocuments(ctx, datasetID, userID, req.DocumentIDs) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) return @@ -1442,8 +1455,8 @@ func (h *DocumentHandler) StopIngestionTasks(c *gin.Context) { } userID := c.GetString("user_id") - - parseResult, err := h.documentService.StopIngestionTasks(req.Tasks, userID) + ctx := c.Request.Context() + parseResult, err := h.documentService.StopIngestionTasks(ctx, req.Tasks, userID) if err != nil { common.ResponseWithCodeData(c, IngestionTaskErrorCode(err), nil, err.Error()) return @@ -1468,8 +1481,9 @@ func (h *DocumentHandler) RemoveIngestionTasks(c *gin.Context) { } userID := c.GetString("user_id") + ctx := c.Request.Context() - deletedTasks, err := h.documentService.RemoveIngestionTasks(req.Tasks, userID) + deletedTasks, err := h.documentService.RemoveIngestionTasks(ctx, req.Tasks, userID) if err != nil { common.ResponseWithCodeData(c, IngestionTaskErrorCode(err), nil, err.Error()) return @@ -1496,8 +1510,8 @@ func (h *DocumentHandler) ParseDocuments(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization to access the dataset.") return } - - parseResult, err := h.documentService.ParseDocuments(datasetID, userID, req.Documents) + ctx := c.Request.Context() + parseResult, err := h.documentService.ParseDocuments(ctx, datasetID, userID, req.Documents) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) return @@ -1529,8 +1543,8 @@ func (h *DocumentHandler) StopParseDocuments(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeDataError, nil, fmt.Sprintf("You don't own the dataset %s.", datasetID)) return } - - result, err := h.documentService.StopParseDocuments(datasetID, req.DocumentIDs) + ctx := c.Request.Context() + result, err := h.documentService.StopParseDocuments(ctx, datasetID, req.DocumentIDs) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) return @@ -1559,8 +1573,8 @@ func (h *DocumentHandler) MetadataSummaryByDataset(c *gin.Context) { if docIDsParam := c.Query("doc_ids"); docIDsParam != "" { docIDS = strings.Split(docIDsParam, ",") } - - summary, err := h.documentService.GetMetadataSummary(datasetID, docIDS) + ctx := c.Request.Context() + summary, err := h.documentService.GetMetadataSummary(ctx, datasetID, docIDS) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, "Failed to get metadata summary"+err.Error()) return @@ -1606,8 +1620,8 @@ func (h *DocumentHandler) UpdateDatasetDocument(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) return } - - data, code, err := h.documentService.UpdateDatasetDocument(user.ID, datasetID, documentID, &req, present) + ctx := c.Request.Context() + data, code, err := h.documentService.UpdateDatasetDocument(ctx, user.ID, datasetID, documentID, &req, present) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -1644,8 +1658,9 @@ func (h *DocumentHandler) UploadInfo(c *gin.Context) { return } + ctx := c.Request.Context() if rawURL != "" { - data, code, err := h.fileService.UploadDocumentInfoByURL(user.ID, rawURL) + data, code, err := h.fileService.UploadDocumentInfoByURL(ctx, user.ID, rawURL) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -1654,7 +1669,7 @@ func (h *DocumentHandler) UploadInfo(c *gin.Context) { return } - data, code, err := h.fileService.UploadDocumentInfos(user.ID, fileHeaders) + data, code, err := h.fileService.UploadDocumentInfos(ctx, user.ID, fileHeaders) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -1727,8 +1742,8 @@ func (h *DocumentHandler) handleBatchUpdateDocumentMetadatas(c *gin.Context) { Updates: updates, Deletes: deletes, } - - resp, code, err := h.documentService.BatchUpdateDocumentMetadatas(datasetID, req.Selector, req.Updates, req.Deletes) + ctx := c.Request.Context() + resp, code, err := h.documentService.BatchUpdateDocumentMetadatas(ctx, datasetID, req.Selector, req.Updates, req.Deletes) if err != nil { common.ErrorWithCode(c, code, err.Error()) return diff --git a/internal/handler/document_test.go b/internal/handler/document_test.go index 257a157d94..7a2c9e5204 100644 --- a/internal/handler/document_test.go +++ b/internal/handler/document_test.go @@ -18,6 +18,7 @@ package handler import ( "bytes" + "context" "encoding/json" "fmt" "mime/multipart" @@ -82,7 +83,7 @@ type fakeDocumentService struct { metadataByKBs map[string]interface{} } -func (f *fakeDocumentService) Ingest(userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error) { +func (f *fakeDocumentService) Ingest(ctx context.Context, userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error) { f.ingestUserID = userID f.ingestReq = req if f.ingestCode != 0 || f.ingestErr != nil { @@ -93,14 +94,14 @@ func (f *fakeDocumentService) Ingest(userID string, req *document.IngestDocument const uploadTestDatasetID = "123e4567-e89b-12d3-a456-426614174000" -func (f *fakeDocumentService) UpdateDatasetDocument(userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) { +func (f *fakeDocumentService) UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) { return nil, common.CodeSuccess, nil } -func (f *fakeDocumentService) BatchUpdateDocumentMetadatas(datasetID string, selector *document.DocumentMetadataSelector, updates []document.DocumentMetadataUpdate, deletes []document.DocumentMetadataDelete) (*document.BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error) { +func (f *fakeDocumentService) BatchUpdateDocumentMetadatas(ctx context.Context, datasetID string, selector *document.DocumentMetadataSelector, updates []document.DocumentMetadataUpdate, deletes []document.DocumentMetadataDelete) (*document.BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error) { return nil, common.CodeSuccess, nil } -func (f *fakeDocumentService) GetDocumentArtifact(filename, _ string) (*document.ArtifactResponse, error) { +func (f *fakeDocumentService) GetDocumentArtifact(ctx context.Context, filename, _ string) (*document.ArtifactResponse, error) { if filename == "error.txt" { return nil, document.ErrArtifactNotFound } @@ -114,7 +115,7 @@ func (f *fakeDocumentService) GetDocumentArtifact(filename, _ string) (*document ForceAttachment: false, }, nil } -func (f *fakeDocumentService) GetDocumentPreview(docID string) (*document.DocumentPreview, error) { +func (f *fakeDocumentService) GetDocumentPreview(ctx context.Context, docID string) (*document.DocumentPreview, error) { if docID == "not-found" { return nil, fmt.Errorf("not found") } @@ -124,7 +125,7 @@ func (f *fakeDocumentService) GetDocumentPreview(docID string) (*document.Docume FileName: "preview.txt", }, nil } -func (f *fakeDocumentService) DownloadDocument(datasetID, docID string) (*document.DownloadDocumentResp, error) { +func (f *fakeDocumentService) DownloadDocument(ctx context.Context, datasetID, docID string) (*document.DownloadDocumentResp, error) { if docID == "not-found" { return nil, fmt.Errorf("not found") } @@ -134,7 +135,7 @@ func (f *fakeDocumentService) DownloadDocument(datasetID, docID string) (*docume FileName: "doc.pdf", }, nil } -func (f *fakeDocumentService) GetDocumentByID(id string) (*document.DocumentResponse, error) { +func (f *fakeDocumentService) GetDocumentByID(ctx context.Context, id string) (*document.DocumentResponse, error) { if f.docErr != nil { return nil, f.docErr } @@ -143,109 +144,109 @@ func (f *fakeDocumentService) GetDocumentByID(id string) (*document.DocumentResp } return nil, fmt.Errorf("document not found") } -func (f *fakeDocumentService) UpdateDocument(id string, req *document.UpdateDocumentRequest) error { +func (f *fakeDocumentService) UpdateDocument(ctx context.Context, id string, req *document.UpdateDocumentRequest) error { f.updateCalled = true f.updatedID = id return nil } -func (f *fakeDocumentService) DeleteDocument(id string) error { +func (f *fakeDocumentService) DeleteDocument(ctx context.Context, id string) error { f.deleteCalled = true f.deletedID = id return nil } -func (f *fakeDocumentService) DeleteDocuments(ids []string, deleteAll bool, datasetID, userID string) (int, error) { +func (f *fakeDocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error) { return f.deleted, f.err } -func (f *fakeDocumentService) ParseDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) { +func (f *fakeDocumentService) ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) { return nil, nil } -func (f *fakeDocumentService) StopParseDocuments(datasetID string, docIDs []string) (map[string]interface{}, error) { +func (f *fakeDocumentService) StopParseDocuments(ctx context.Context, datasetID string, docIDs []string) (map[string]interface{}, error) { return f.stopResult, f.stopErr } -func (f *fakeDocumentService) ListDocuments(page, pageSize int) ([]*document.DocumentResponse, int64, error) { +func (f *fakeDocumentService) ListDocuments(ctx context.Context, page, pageSize int) ([]*document.DocumentResponse, int64, error) { return nil, 0, nil } -func (f *fakeDocumentService) ListDocumentsByDatasetID(kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { +func (f *fakeDocumentService) ListDocumentsByDatasetID(ctx context.Context, kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { return nil, 0, nil } -func (f *fakeDocumentService) ListDocumentsByDatasetIDWithOptions(opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { +func (f *fakeDocumentService) ListDocumentsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { f.listOpts = opts return nil, 0, nil } -func (f *fakeDocumentService) ListDocumentIDsByDatasetIDWithOptions(opts dao.DocumentListOptions) ([]string, error) { +func (f *fakeDocumentService) ListDocumentIDsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions) ([]string, error) { f.listOpts = opts return f.listIDs, nil } -func (f *fakeDocumentService) GetDocumentFiltersByDatasetID(opts dao.DocumentListOptions) (map[string]interface{}, int64, error) { +func (f *fakeDocumentService) GetDocumentFiltersByDatasetID(ctx context.Context, opts dao.DocumentListOptions) (map[string]interface{}, int64, error) { f.filterOpts = opts if f.filterResult != nil { return f.filterResult, f.filterTotal, nil } return map[string]interface{}{}, 0, nil } -func (f *fakeDocumentService) GetMetadataByKBs(kbIDs []string) (map[string]interface{}, error) { +func (f *fakeDocumentService) GetMetadataByKBs(ctx context.Context, kbIDs []string) (map[string]interface{}, error) { if f.metadataByKBs != nil { return f.metadataByKBs, nil } return map[string]interface{}{}, nil } -func (f *fakeDocumentService) BatchUpdateDocumentStatus(userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) { +func (f *fakeDocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) { return map[string]interface{}{}, common.CodeSuccess, nil } -func (f *fakeDocumentService) GetThumbnails(userID string, docIDs []string) (map[string]string, error) { +func (f *fakeDocumentService) GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error) { f.thumbnailUserID = userID f.thumbnailDocIDs = append([]string(nil), docIDs...) return f.thumbnails, f.thumbnailErr } -func (f *fakeDocumentService) GetDocumentImage(imageID string) ([]byte, error) { +func (f *fakeDocumentService) GetDocumentImage(ctx context.Context, imageID string) ([]byte, error) { return nil, nil } -func (f *fakeDocumentService) GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error) { +func (f *fakeDocumentService) GetDocumentsByAuthorID(ctx context.Context, authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error) { return nil, 0, nil } -func (f *fakeDocumentService) GetMetadataSummary(kbID string, docIDs []string) (map[string]interface{}, error) { +func (f *fakeDocumentService) GetMetadataSummary(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error) { f.metadataKBID = kbID f.metadataDocIDs = docIDs return f.metadataSummary, f.metadataErr } -func (f *fakeDocumentService) SetDocumentMetadata(docID string, meta map[string]interface{}) error { +func (f *fakeDocumentService) SetDocumentMetadata(ctx context.Context, docID string, meta map[string]interface{}) error { f.setMetaCalled = true f.setMetaDocID = docID f.setMetaValue = meta return nil } -func (f *fakeDocumentService) DeleteDocumentMetadata(docID string, keys []string) error { +func (f *fakeDocumentService) DeleteDocumentMetadata(ctx context.Context, docID string, keys []string) error { return nil } -func (f *fakeDocumentService) DeleteDocumentAllMetadata(docID string) error { +func (f *fakeDocumentService) DeleteDocumentAllMetadata(ctx context.Context, docID string) error { return nil } -func (f *fakeDocumentService) GetDocumentMetadataByID(docID string) (map[string]interface{}, error) { +func (f *fakeDocumentService) GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]interface{}, error) { return nil, nil } -func (f *fakeDocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) { +func (f *fakeDocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) { f.uploadLocalKB = kb f.uploadLocalPath = parentPath f.uploadOverride = parserConfigOverride return f.uploadLocalData, f.uploadLocalErrs } -func (f *fakeDocumentService) UploadWebDocument(kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) { +func (f *fakeDocumentService) UploadWebDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) { return nil, common.CodeServerError, fmt.Errorf("not implemented") } -func (f *fakeDocumentService) UploadEmptyDocument(kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) { +func (f *fakeDocumentService) UploadEmptyDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) { return nil, common.CodeServerError, fmt.Errorf("not implemented") } -func (f *fakeDocumentService) ListIngestionTasks(userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) { +func (f *fakeDocumentService) ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) { return nil, nil } -func (f *fakeDocumentService) IngestDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) { +func (f *fakeDocumentService) IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) { return nil, nil } -func (f *fakeDocumentService) StopIngestionTasks(tasks []string, userID string) ([]*entity.IngestionTask, error) { +func (f *fakeDocumentService) StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error) { return f.stopIngestionTasks, f.stopIngestionTaskErr } -func (f *fakeDocumentService) RemoveIngestionTasks(tasks []string, userID string) ([]map[string]string, error) { +func (f *fakeDocumentService) RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error) { return f.removeIngestionTasks, f.removeIngestionTaskErr } diff --git a/internal/handler/file.go b/internal/handler/file.go index e0358bb097..d830339441 100644 --- a/internal/handler/file.go +++ b/internal/handler/file.go @@ -306,7 +306,8 @@ func (h *FileHandler) UploadFile(c *gin.Context) { } } - result, err := h.fileService.UploadFile(userID, parentID, files) + ctx := c.Request.Context() + result, err := h.fileService.UploadFile(ctx, userID, parentID, files) if err != nil { common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) return @@ -427,7 +428,8 @@ func (h *FileHandler) MoveFiles(c *gin.Context) { return } - success, message := h.fileService.MoveFiles(user.ID, req.SrcFileIDs, req.DestFileID, req.NewName) + ctx := c.Request.Context() + success, message := h.fileService.MoveFiles(ctx, user.ID, req.SrcFileIDs, req.DestFileID, req.NewName) if !success { common.ResponseWithCodeData(c, common.CodeBadRequest, nil, message) return @@ -482,7 +484,8 @@ func (h *FileHandler) Download(c *gin.Context) { // If blob is empty, try fallback via file2document if len(blob) == 0 { - storageAddr, err := h.fileService.GetStorageAddress(fileID) + ctx := c.Request.Context() + storageAddr, err := h.fileService.GetStorageAddress(ctx, fileID) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, "Failed to get file storage address: "+err.Error()) return @@ -564,8 +567,8 @@ func (h *FileHandler) LinkToDatasets(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "mode must be 'add' or 'replace'") return } - - if err := h.file2DocumentService.LinkToDatasets(user.ID, &req, mode); err != nil { + ctx := c.Request.Context() + if err := h.file2DocumentService.LinkToDatasets(ctx, user.ID, &req, mode); err != nil { common.ResponseWithCodeData(c, linkToDatasetsErrorCode(err), nil, err.Error()) return } diff --git a/internal/handler/mindmap.go b/internal/handler/mindmap.go index 439bfc97e2..11c8c99dd4 100644 --- a/internal/handler/mindmap.go +++ b/internal/handler/mindmap.go @@ -42,7 +42,7 @@ type mindMapRunConfig struct { TenantSvc *service.TenantService } -func runMindMap(config mindMapRunConfig) (mindMapNode, error) { +func runMindMap(ctx context.Context, config mindMapRunConfig) (mindMapNode, error) { if config.ChunkSvc == nil { return mindMapNode{}, fmt.Errorf("chunk service not configured") } @@ -54,7 +54,7 @@ func runMindMap(config mindMapRunConfig) (mindMapNode, error) { modelTenantID = config.AuthUserID } retrievalReq := mindMapRetrievalRequest(config.Question, config.KbIDs, config.SearchID, config.SearchConfig) - ranks, err := config.ChunkSvc.RetrievalTest(retrievalReq, config.AuthUserID) + ranks, err := config.ChunkSvc.RetrievalTest(ctx, retrievalReq, config.AuthUserID) if err != nil { return mindMapNode{}, err } diff --git a/internal/handler/searchbot.go b/internal/handler/searchbot.go index 1af77240d9..db814b14e6 100644 --- a/internal/handler/searchbot.go +++ b/internal/handler/searchbot.go @@ -162,9 +162,7 @@ func (h *SearchBotHandler) Handle(c *gin.Context) { // RetrievalTest performs a retrieval test against specified knowledge bases. // @Summary Retrieval Test // @Description Test document retrieval across knowledge bases with optional filters, reranking, and KG search. -// @Tags searchbots -// @Accept json -// @Produce json +// @Tags searchBots // @Param request body SearchBotRetrievalTestRequest true "Retrieval test parameters" // @Success 200 {object} map[string]interface{} // @Router /api/v1/searchbots/retrieval_test [post] @@ -203,11 +201,11 @@ func (h *SearchBotHandler) RetrievalTest(c *gin.Context) { } svcReq := toRetrievalServiceRequest(&req) - - result, err := h.chunkSvc.RetrievalTest(svcReq, user.ID) + ctx := c.Request.Context() + result, err := h.chunkSvc.RetrievalTest(ctx, svcReq, user.ID) if err != nil { common.Warn("search bot retrieval test failed", zap.String("error", err.Error())) - common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeServerError, nil, "retrieval test failed") + common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, "retrieval test failed") return } @@ -360,7 +358,8 @@ func (h *SearchBotHandler) MindMap(c *gin.Context) { searchConfig = searchConfigFromDetail(detail) } - mindMap, err := runMindMap(mindMapRunConfig{ + ctx := c.Request.Context() + mindMap, err := runMindMap(ctx, mindMapRunConfig{ Question: req.Question, KbIDs: filtered, SearchID: req.SearchID, diff --git a/internal/handler/searchbot_test.go b/internal/handler/searchbot_test.go index cca2a40353..fe0002dcce 100644 --- a/internal/handler/searchbot_test.go +++ b/internal/handler/searchbot_test.go @@ -17,6 +17,7 @@ package handler import ( + "context" "encoding/json" "errors" "fmt" @@ -40,7 +41,7 @@ type mockChunkService struct { LastUserID string } -func (m *mockChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { +func (m *mockChunkService) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { m.LastReq = req m.LastUserID = userID if m.retrievalTestFn != nil { @@ -466,7 +467,7 @@ type fakeChunkRetriever struct { err error } -func (f *fakeChunkRetriever) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { +func (f *fakeChunkRetriever) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { if f.err != nil { return nil, f.err } diff --git a/internal/ingestion/component/chunker/pdfcrop_cgo.go b/internal/ingestion/component/chunker/pdfcrop_cgo.go index 9bdf0d1d47..5d964e5484 100644 --- a/internal/ingestion/component/chunker/pdfcrop_cgo.go +++ b/internal/ingestion/component/chunker/pdfcrop_cgo.go @@ -35,7 +35,7 @@ func newPDFEngineFromUpstream(ctx context.Context, up schema.ChunkerFromUpstream data, err = component.FetchBinary(ctx, up.Bucket, up.Path) case up.DocID != "": var ref *component.DocumentStorageRef - ref, err = component.ResolveDocumentStorage(up.DocID) + ref, err = component.ResolveDocumentStorage(ctx, up.DocID) if err == nil && ref != nil { data, err = component.FetchBinary(ctx, ref.Bucket, ref.Path) } diff --git a/internal/ingestion/component/document_storage.go b/internal/ingestion/component/document_storage.go index 73b0aa54eb..db06c9af77 100644 --- a/internal/ingestion/component/document_storage.go +++ b/internal/ingestion/component/document_storage.go @@ -81,12 +81,12 @@ func resolveStorage() storage.Storage { // location. It is exported so downstream components can re-acquire the // source PDF without threading the raw bytes across the component // boundary. -func ResolveDocumentStorage(docID string) (*DocumentStorageRef, error) { +func ResolveDocumentStorage(ctx context.Context, docID string) (*DocumentStorageRef, error) { if ResolveDocumentStorageOverride != nil { return ResolveDocumentStorageOverride(docID) } - doc, err := dao.NewDocumentDAO().GetByID(docID) + doc, err := dao.NewDocumentDAO().GetByID(ctx, dao.DB, docID) if err != nil { return nil, err } @@ -118,7 +118,7 @@ func ResolveDocumentStorage(docID string) (*DocumentStorageRef, error) { return ref, nil } -func resolveDocumentName(docID string) (string, error) { +func resolveDocumentName(ctx context.Context, docID string) (string, error) { if ResolveDocumentStorageOverride != nil { ref, err := ResolveDocumentStorageOverride(docID) if err != nil { @@ -128,7 +128,7 @@ func resolveDocumentName(docID string) (string, error) { return ref.Name, nil } } - doc, err := dao.NewDocumentDAO().GetByID(docID) + doc, err := dao.NewDocumentDAO().GetByID(ctx, dao.DB, docID) if err != nil { return "", err } diff --git a/internal/ingestion/component/file.go b/internal/ingestion/component/file.go index a3073d10d6..8e80027189 100644 --- a/internal/ingestion/component/file.go +++ b/internal/ingestion/component/file.go @@ -131,7 +131,7 @@ func (c *FileComponent) Outputs() map[string]string { func (c *FileComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { // Parse the wire input through the schema type so the // validation errors match the package convention. - in, err := parseFileInputs(inputs) + in, err := parseFileInputs(ctx, inputs) if err != nil { return nil, err } @@ -172,7 +172,7 @@ type fileInputs struct { // parseFileInputs parses and validates the upstream input map. // Mirrors python's branching on `self._canvas._doc_id` vs. // `kwargs.get("file")[0]`. -func parseFileInputs(inputs map[string]any) (fileInputs, error) { +func parseFileInputs(ctx context.Context, inputs map[string]any) (fileInputs, error) { if inputs == nil { return fileInputs{}, fmt.Errorf("file: inputs map is nil") } @@ -229,7 +229,7 @@ func parseFileInputs(inputs map[string]any) (fileInputs, error) { out.path = v } if out.docID != "" { - name, err := resolveDocumentName(out.docID) + name, err := resolveDocumentName(ctx, out.docID) if err != nil { return fileInputs{}, fmt.Errorf("file: resolve doc_id %q: %w", out.docID, err) } diff --git a/internal/ingestion/component/parser.go b/internal/ingestion/component/parser.go index 6288455c3e..3438fe0fa7 100644 --- a/internal/ingestion/component/parser.go +++ b/internal/ingestion/component/parser.go @@ -638,7 +638,7 @@ func readParserBinary(ctx context.Context, inputs map[string]any) ([]byte, error return FetchBinary(ctx, bucket, path) } if docID, ok := getString(inputs, "doc_id"); ok && docID != "" { - ref, err := ResolveDocumentStorage(docID) + ref, err := ResolveDocumentStorage(ctx, docID) if err != nil { return nil, fmt.Errorf("Parser: resolve doc_id %q: %w", docID, err) } diff --git a/internal/ingestion/pipeline/pipeline.go b/internal/ingestion/pipeline/pipeline.go index e1c1d7acbf..54bc0b39da 100644 --- a/internal/ingestion/pipeline/pipeline.go +++ b/internal/ingestion/pipeline/pipeline.go @@ -117,8 +117,8 @@ type ProgressEvent struct { // (internal/ingestion/service). A nil sink is valid: events are dropped and // the pipeline stays DB-independent (unit tests, headless runs). type ProgressSink interface { - OnComponentTotal(taskID string, total int) - OnComponentProgress(ev ProgressEvent) + OnComponentTotal(ctx context.Context, taskID string, total int) + OnComponentProgress(ctx context.Context, ev ProgressEvent) } // WithProgressSink injects a sink that receives component progress events @@ -272,7 +272,7 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, override_para // progress percentage. Best-effort: a DB failure (or headless run // with no DB) must not abort the pipeline — progress is observability. if p.sink != nil { - p.sink.OnComponentTotal(p.taskID, len(p.canvas.Components)) + p.sink.OnComponentTotal(ctx, p.taskID, len(p.canvas.Components)) } runState := canvas.NewCanvasState("", p.taskID) @@ -284,7 +284,7 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, override_para // is nil when the DB is not initialized (unit tests, headless // runs), in which case TrackProgress is a no-op — progress is an // observability concern, not a data dependency. - runCtx = runtime.WithProgressCallback(runCtx, p.componentProgressCallback()) + runCtx = runtime.WithProgressCallback(runCtx, p.componentProgressCallback(ctx)) current := cloneMapOrEmpty(inputs) @@ -303,7 +303,7 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, override_para // Resumable path: record the run, then loop Invoke until the graph // completes or a non-resumable error surfaces. if tracker != nil { - if err := tracker.Start(ctx, p.taskID, "", "", ""); err != nil { + if err = tracker.Start(ctx, p.taskID, "", "", ""); err != nil { common.Error(fmt.Sprintf("pipeline: RunTracker.Start for task %s failed: %v", p.taskID, err), err) } } @@ -459,7 +459,7 @@ func finalizeResult(current, out map[string]any, runState *canvas.CanvasState) m // never touches the DAO layer. Returns nil when no sink is attached, leaving // TrackProgress a no-op and the pipeline DB-independent (unit tests, headless // runs). -func (p *Pipeline) componentProgressCallback() runtime.ProgressCallback { +func (p *Pipeline) componentProgressCallback(ctx context.Context) runtime.ProgressCallback { if p.sink == nil { return nil } @@ -477,7 +477,7 @@ func (p *Pipeline) componentProgressCallback() runtime.ProgressCallback { msg = ev.Component + " Error" } } - p.sink.OnComponentProgress(ProgressEvent{ + p.sink.OnComponentProgress(ctx, ProgressEvent{ TaskID: p.taskID, DocumentID: p.documentID, Component: ev.Component, diff --git a/internal/ingestion/pipeline/pipeline_test.go b/internal/ingestion/pipeline/pipeline_test.go index a6c79f93cc..0ac36966f7 100644 --- a/internal/ingestion/pipeline/pipeline_test.go +++ b/internal/ingestion/pipeline/pipeline_test.go @@ -489,14 +489,14 @@ type recordingSink struct { events []ProgressEvent } -func (r *recordingSink) OnComponentTotal(taskID string, total int) { +func (r *recordingSink) OnComponentTotal(ctx context.Context, taskID string, total int) { r.mu.Lock() defer r.mu.Unlock() r.total = total r.totalSet = true } -func (r *recordingSink) OnComponentProgress(ev ProgressEvent) { +func (r *recordingSink) OnComponentProgress(ctx context.Context, ev ProgressEvent) { r.mu.Lock() defer r.mu.Unlock() r.events = append(r.events, ev) diff --git a/internal/ingestion/service/doc_state.go b/internal/ingestion/service/doc_state.go index 21e5ad7dae..64530f60a9 100644 --- a/internal/ingestion/service/doc_state.go +++ b/internal/ingestion/service/doc_state.go @@ -17,6 +17,7 @@ package service import ( + "context" "fmt" "ragflow/internal/common" @@ -29,9 +30,9 @@ import ( // can inject a stub without constructing a real DocumentService (which depends // on initialized server config). type docStateSvc interface { - GetDocumentMetadataByID(docID string) (map[string]any, error) - SetDocumentMetadata(docID string, meta map[string]any) error - IncrementChunkNum(docID, kbID string, chunkNum, tokenNum int, duration float64) error + GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]any, error) + SetDocumentMetadata(ctx context.Context, docID string, meta map[string]any) error + IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error } // docStateUpdater applies a pipeline run's results to document state: it @@ -50,16 +51,16 @@ func newDocStateUpdater() *docStateUpdater { } } -func (u *docStateUpdater) apply(r *taskpkg.PipelineResult) { +func (u *docStateUpdater) apply(ctx context.Context, r *taskpkg.PipelineResult) { if r == nil { return } if len(r.Metadata) > 0 { - if err := mergeDocMetadata(u.docSvc, r.DocID, r.Metadata); err != nil { + if err := mergeDocMetadata(ctx, u.docSvc, r.DocID, r.Metadata); err != nil { common.Warn(fmt.Sprintf("failed to update document metadata: %v", err)) } } - if err := u.docSvc.IncrementChunkNum(r.DocID, r.KbID, r.ChunkCount, r.TokenConsumption, r.Duration); err != nil { + if err := u.docSvc.IncrementChunkNum(ctx, r.DocID, r.KbID, r.ChunkCount, r.TokenConsumption, r.Duration); err != nil { common.Warn(fmt.Sprintf("failed to increment chunk num: %v", err)) } } @@ -68,8 +69,8 @@ func (u *docStateUpdater) apply(r *taskpkg.PipelineResult) { // (existing keys are preserved, not overwritten), and writes the merged map back. // A read failure aborts the merge: SetDocumentMetadata is a full overwrite, so // writing with an empty baseline would destroy existing keys. -func mergeDocMetadata(svc docStateSvc, docID string, metadata map[string]any) error { - existing, err := svc.GetDocumentMetadataByID(docID) +func mergeDocMetadata(ctx context.Context, svc docStateSvc, docID string, metadata map[string]any) error { + existing, err := svc.GetDocumentMetadataByID(ctx, docID) if err != nil { return err } @@ -81,5 +82,5 @@ func mergeDocMetadata(svc docStateSvc, docID string, metadata map[string]any) er existing[k] = v } } - return svc.SetDocumentMetadata(docID, existing) + return svc.SetDocumentMetadata(ctx, docID, existing) } diff --git a/internal/ingestion/service/doc_state_test.go b/internal/ingestion/service/doc_state_test.go index 77503d19da..f403d3a090 100644 --- a/internal/ingestion/service/doc_state_test.go +++ b/internal/ingestion/service/doc_state_test.go @@ -17,6 +17,7 @@ package service import ( + "context" "testing" taskpkg "ragflow/internal/ingestion/task" @@ -33,14 +34,14 @@ type stubDocStateSvc struct { incrementCalled bool } -func (s *stubDocStateSvc) GetDocumentMetadataByID(docID string) (map[string]any, error) { +func (s *stubDocStateSvc) GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]any, error) { if s.metaData == nil { return make(map[string]any), nil } return s.metaData, nil } -func (s *stubDocStateSvc) SetDocumentMetadata(docID string, meta map[string]any) error { +func (s *stubDocStateSvc) SetDocumentMetadata(ctx context.Context, docID string, meta map[string]any) error { s.setCalled = true if s.metaData == nil { s.metaData = make(map[string]any) @@ -51,7 +52,7 @@ func (s *stubDocStateSvc) SetDocumentMetadata(docID string, meta map[string]any) return nil } -func (s *stubDocStateSvc) IncrementChunkNum(docID, kbID string, chunkNum, tokenNum int, duration float64) error { +func (s *stubDocStateSvc) IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error { s.incrementCalled = true s.gotDocID = docID s.gotKbID = kbID @@ -64,7 +65,8 @@ func (s *stubDocStateSvc) IncrementChunkNum(docID, kbID string, chunkNum, tokenN func TestDocStateUpdater_NilResultIsNoop(t *testing.T) { svc := &stubDocStateSvc{} u := &docStateUpdater{docSvc: svc} - u.apply(nil) + ctx := t.Context() + u.apply(ctx, nil) if svc.setCalled || svc.incrementCalled { t.Fatal("nil result must not touch document state") } @@ -73,7 +75,8 @@ func TestDocStateUpdater_NilResultIsNoop(t *testing.T) { func TestDocStateUpdater_EmptyMetadataSkipsMerge(t *testing.T) { svc := &stubDocStateSvc{} u := &docStateUpdater{docSvc: svc} - u.apply(&taskpkg.PipelineResult{DocID: "doc-1", KbID: "kb-1", ChunkCount: 3, TokenConsumption: 100}) + ctx := t.Context() + u.apply(ctx, &taskpkg.PipelineResult{DocID: "doc-1", KbID: "kb-1", ChunkCount: 3, TokenConsumption: 100}) if svc.setCalled { t.Fatal("empty metadata must not call SetDocumentMetadata") } @@ -88,7 +91,8 @@ func TestDocStateUpdater_EmptyMetadataSkipsMerge(t *testing.T) { func TestDocStateUpdater_MergesNewKeys(t *testing.T) { svc := &stubDocStateSvc{metaData: map[string]any{"existing": "old"}} u := &docStateUpdater{docSvc: svc} - u.apply(&taskpkg.PipelineResult{DocID: "doc-1", Metadata: map[string]any{"new_key": "value"}, ChunkCount: 1, TokenConsumption: 10}) + ctx := t.Context() + u.apply(ctx, &taskpkg.PipelineResult{DocID: "doc-1", Metadata: map[string]any{"new_key": "value"}, ChunkCount: 1, TokenConsumption: 10}) if svc.metaData["existing"] != "old" { t.Errorf("existing key should be preserved: got %q", svc.metaData["existing"]) } @@ -100,7 +104,8 @@ func TestDocStateUpdater_MergesNewKeys(t *testing.T) { func TestDocStateUpdater_PreservesExistingKey(t *testing.T) { svc := &stubDocStateSvc{metaData: map[string]any{"author": "Alice"}} u := &docStateUpdater{docSvc: svc} - u.apply(&taskpkg.PipelineResult{DocID: "doc-1", Metadata: map[string]any{"author": "Bob"}, ChunkCount: 1, TokenConsumption: 10}) + ctx := t.Context() + u.apply(ctx, &taskpkg.PipelineResult{DocID: "doc-1", Metadata: map[string]any{"author": "Bob"}, ChunkCount: 1, TokenConsumption: 10}) if svc.metaData["author"] != "Alice" { t.Errorf("existing key must NOT be overwritten: got %q", svc.metaData["author"]) } @@ -109,7 +114,8 @@ func TestDocStateUpdater_PreservesExistingKey(t *testing.T) { func TestDocStateUpdater_IncrementArgs(t *testing.T) { svc := &stubDocStateSvc{} u := &docStateUpdater{docSvc: svc} - u.apply(&taskpkg.PipelineResult{DocID: "doc-1", KbID: "kb-1", ChunkCount: 10, TokenConsumption: 100}) + ctx := t.Context() + u.apply(ctx, &taskpkg.PipelineResult{DocID: "doc-1", KbID: "kb-1", ChunkCount: 10, TokenConsumption: 100}) if svc.gotDocID != "doc-1" || svc.gotKbID != "kb-1" { t.Fatalf("docID=%q kbID=%q, want doc-1/kb-1", svc.gotDocID, svc.gotKbID) } diff --git a/internal/ingestion/service/execute_task_test.go b/internal/ingestion/service/execute_task_test.go index cb1af2e664..5016473c33 100644 --- a/internal/ingestion/service/execute_task_test.go +++ b/internal/ingestion/service/execute_task_test.go @@ -229,7 +229,8 @@ func TestExecuteTask_CancelBeforePipeline(t *testing.T) { t.Fatal("expected runDocumentTask to NOT be called when cancel is detected before pipeline") } - doc, err := dao.NewDocumentDAO().GetByID(docID) + ctx := t.Context() + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, docID) if err != nil { t.Fatalf("load document: %v", err) } diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index 465e6bc966..dab2fb8c60 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -183,7 +183,7 @@ func (e *Ingestor) processMessage(handle common.TaskHandle) { return } - task, err := e.ingestionTaskSvc.StartRunning(taskMessage.TaskID) + task, err := e.ingestionTaskSvc.StartRunning(e.ctx, taskMessage.TaskID) if err != nil { if errors.Is(err, common.ErrTaskNotFound) { common.Warn(fmt.Sprintf("task %s not found, skipping", taskMessage.TaskID)) @@ -609,7 +609,7 @@ func (e *Ingestor) pollCancel(taskID string, cancel context.CancelFunc, done <-c // appended timestamped cancel message (progress_msg += cancelMsg). func (e *Ingestor) markCancelProgress(task *entity.IngestionTask) { svc := documentpkg.NewDocumentService() - doc, err := svc.GetDocumentByID(task.DocumentID) + doc, err := svc.GetDocumentByID(e.ctx, task.DocumentID) if err != nil { common.Error(fmt.Sprintf("markCancelProgress: load document %s: %v", task.DocumentID, err), err) return @@ -619,7 +619,7 @@ func (e *Ingestor) markCancelProgress(task *entity.IngestionTask) { if doc.ProgressMsg != nil { existingMsg = *doc.ProgressMsg } - _ = svc.UpdateRunProgress(task.DocumentID, -1.0, string(entity.TaskStatusCancel), existingMsg+cancelMsg) + _ = svc.UpdateRunProgress(e.ctx, task.DocumentID, -1.0, string(entity.TaskStatusCancel), existingMsg+cancelMsg) } // markTimeoutProgress writes the timeout-progress markers to the document @@ -627,7 +627,7 @@ func (e *Ingestor) markCancelProgress(task *entity.IngestionTask) { // failure rather than a user-initiated stop. func (e *Ingestor) markTimeoutProgress(task *entity.IngestionTask) { svc := documentpkg.NewDocumentService() - doc, err := svc.GetDocumentByID(task.DocumentID) + doc, err := svc.GetDocumentByID(e.ctx, task.DocumentID) if err != nil { common.Error(fmt.Sprintf("markTimeoutProgress: load document %s: %v", task.DocumentID, err), err) return @@ -637,7 +637,7 @@ func (e *Ingestor) markTimeoutProgress(task *entity.IngestionTask) { if doc.ProgressMsg != nil { existingMsg = *doc.ProgressMsg } - _ = svc.UpdateRunProgress(task.DocumentID, -1.0, string(entity.TaskStatusFail), existingMsg+timeoutMsg) + _ = svc.UpdateRunProgress(e.ctx, task.DocumentID, -1.0, string(entity.TaskStatusFail), existingMsg+timeoutMsg) } // claimTask registers a worker claim on a task ID. Returns false if another @@ -697,7 +697,7 @@ func (e *Ingestor) releaseTask(taskID string) { } func (e *Ingestor) defaultRunDocumentTask(ctx context.Context, ingestionTask *entity.IngestionTask) error { - docTaskCtx, err := taskpkg.LoadFromIngestionTask(ingestionTask) + docTaskCtx, err := taskpkg.LoadFromIngestionTask(ctx, ingestionTask) if err != nil { return fmt.Errorf("load task context for %s: %w", ingestionTask.ID, err) } @@ -732,11 +732,11 @@ func (e *Ingestor) defaultRunDocumentTask(ctx context.Context, ingestionTask *en return dsl, parserID, nil }) } - result, err := executor.WithRequireResume().WithProgressSink(newProgressSink(e.ingestionTaskSvc)).Execute(docTaskCtx.Ctx) + result, err := executor.WithRequireResume().WithProgressSink(newProgressSink(ctx, e.ingestionTaskSvc)).Execute(docTaskCtx.Ctx) if err != nil { return err } - e.docState.apply(result) + e.docState.apply(ctx, result) return nil } diff --git a/internal/ingestion/service/progress_sink.go b/internal/ingestion/service/progress_sink.go index 3f0f350434..6932350864 100644 --- a/internal/ingestion/service/progress_sink.go +++ b/internal/ingestion/service/progress_sink.go @@ -17,6 +17,7 @@ package service import ( + "context" "fmt" "sync/atomic" @@ -50,10 +51,10 @@ type progressSink struct { // tests can inject a stub and assert the mirror call without depending on the // full DocumentService surface. type docProgressSvc interface { - UpdateRunProgress(docID string, progress float64, run, progressMsg string) error + UpdateRunProgress(ctx context.Context, docID string, progress float64, run, progressMsg string) error } -func newProgressSink(taskSvc *servicepkg.IngestionTaskService) *progressSink { +func newProgressSink(ctx context.Context, taskSvc *servicepkg.IngestionTaskService) *progressSink { // Eagerly construct the DocumentService so docSvc is immutable after this // point. eino's compose graph runs parallel branches concurrently, so // OnComponentProgress (and thus docSvc) can fire from multiple goroutines; @@ -65,14 +66,14 @@ func newProgressSink(taskSvc *servicepkg.IngestionTaskService) *progressSink { } } -func (s *progressSink) OnComponentTotal(taskID string, total int) { +func (s *progressSink) OnComponentTotal(ctx context.Context, taskID string, total int) { s.total.Store(int64(total)) if err := s.taskSvc.UpdateComponentTotal(taskID, total); err != nil { common.Error(fmt.Sprintf("progressSink: update component_total for task %s failed: %v", taskID, err), err) } } -func (s *progressSink) OnComponentProgress(ev pipeline.ProgressEvent) { +func (s *progressSink) OnComponentProgress(ctx context.Context, ev pipeline.ProgressEvent) { if err := s.taskSvc.RecordComponentProgress(ev.TaskID, ev.Component, ev.Phase, ev.Message); err != nil { common.Error(fmt.Sprintf("progressSink: record component progress for task %s failed: %v", ev.TaskID, err), err) } @@ -89,7 +90,7 @@ func (s *progressSink) OnComponentProgress(ev pipeline.ProgressEvent) { return } progress, run := deriveDocumentProgress(agg, int(total)) - if err := s.docSvc.UpdateRunProgress(ev.DocumentID, progress, run, ev.Message); err != nil { + if err = s.docSvc.UpdateRunProgress(ctx, ev.DocumentID, progress, run, ev.Message); err != nil { common.Error(fmt.Sprintf("progressSink: mirror progress to document %s for task %s failed: %v", ev.DocumentID, ev.TaskID, err), err) } } diff --git a/internal/ingestion/service/progress_sink_test.go b/internal/ingestion/service/progress_sink_test.go index 710d8df10c..8980830c47 100644 --- a/internal/ingestion/service/progress_sink_test.go +++ b/internal/ingestion/service/progress_sink_test.go @@ -17,6 +17,7 @@ package service import ( + "context" "runtime" "sync" "testing" @@ -56,7 +57,8 @@ func TestProgressSink_EagerlyConstructsDocumentService(t *testing.T) { cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() - sink := newProgressSink(servicepkg.NewIngestionTaskService()) + ctx := t.Context() + sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService()) if sink.docSvc == nil { t.Fatal("expected sink to eagerly construct its DocumentService, got nil (lazy)") } @@ -76,10 +78,11 @@ func TestProgressSink_DocService_NoDataRace(t *testing.T) { cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() + ctx := t.Context() // Deliberately do NOT inject a stub docSvc: the sink's own DocumentService // must already be constructed (not lazily built mid-call) when the // goroutines below race into docSvc. - sink := newProgressSink(servicepkg.NewIngestionTaskService()) + sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService()) const n = 30 var wg sync.WaitGroup @@ -108,7 +111,8 @@ func TestProgressSink_Total_NoDataRace(t *testing.T) { defer cleanup() _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - sink := newProgressSink(servicepkg.NewIngestionTaskService()) + ctx := t.Context() + sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService()) const n = 30 var wg sync.WaitGroup @@ -118,7 +122,7 @@ func TestProgressSink_Total_NoDataRace(t *testing.T) { go func() { defer wg.Done() <-start - sink.OnComponentTotal(taskID, 5) // writes s.total + sink.OnComponentTotal(ctx, taskID, 5) // writes s.total }() } for i := 0; i < n; i++ { @@ -142,7 +146,7 @@ type stubDocProgressSvc struct { calls int } -func (s *stubDocProgressSvc) UpdateRunProgress(docID string, progress float64, run, progressMsg string) error { +func (s *stubDocProgressSvc) UpdateRunProgress(ctx context.Context, docID string, progress float64, run, progressMsg string) error { s.calls++ s.gotDocID = docID s.gotProgress = progress @@ -160,11 +164,12 @@ func TestProgressSinkPersistsViaService(t *testing.T) { defer cleanup() _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - sink := newProgressSink(servicepkg.NewIngestionTaskService()) + ctx := t.Context() + sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService()) stub := &stubDocProgressSvc{} sink.docSvc = stub - sink.OnComponentTotal(taskID, 2) + sink.OnComponentTotal(ctx, taskID, 2) task, err := dao.NewIngestionTaskDAO().GetByID(taskID) if err != nil { t.Fatalf("load task: %v", err) @@ -173,7 +178,7 @@ func TestProgressSinkPersistsViaService(t *testing.T) { t.Fatalf("component_total = %d, want 2", task.ComponentTotal) } - sink.OnComponentProgress(pipeline.ProgressEvent{ + sink.OnComponentProgress(ctx, pipeline.ProgressEvent{ TaskID: taskID, DocumentID: docID, Component: "Parser", @@ -218,11 +223,12 @@ func TestProgressSinkEmptyDocumentIDSkipsMirror(t *testing.T) { defer cleanup() _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - sink := newProgressSink(servicepkg.NewIngestionTaskService()) + ctx := t.Context() + sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService()) stub := &stubDocProgressSvc{} sink.docSvc = stub - sink.OnComponentProgress(pipeline.ProgressEvent{ + sink.OnComponentProgress(ctx, pipeline.ProgressEvent{ TaskID: taskID, Component: "Chunker", Phase: 1, diff --git a/internal/ingestion/task/pipeline_e2e_test.go b/internal/ingestion/task/pipeline_e2e_test.go index 7a9d487af7..aafe4d72ce 100644 --- a/internal/ingestion/task/pipeline_e2e_test.go +++ b/internal/ingestion/task/pipeline_e2e_test.go @@ -176,6 +176,7 @@ func TestPipelineE2E_PipelineExecutor(t *testing.T) { }, } + ctx := t.Context() for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { t.Logf("Running Pipeline E2E test with engine: %s", tc.name) @@ -205,7 +206,7 @@ func TestPipelineE2E_PipelineExecutor(t *testing.T) { if err != nil { t.Fatalf("GetByID failed: %v", err) } - taskCtx, err := LoadFromIngestionTask(ingestionTask) + taskCtx, err := LoadFromIngestionTask(ctx, ingestionTask) if err != nil { t.Fatalf("LoadFromIngestionTask failed: %v", err) } diff --git a/internal/ingestion/task/pipeline_executor_test.go b/internal/ingestion/task/pipeline_executor_test.go index 8e504ac6fe..3c63367963 100644 --- a/internal/ingestion/task/pipeline_executor_test.go +++ b/internal/ingestion/task/pipeline_executor_test.go @@ -417,14 +417,14 @@ type recordingProgressSink struct { events []pipelinepkg.ProgressEvent } -func (r *recordingProgressSink) OnComponentTotal(taskID string, total int) { +func (r *recordingProgressSink) OnComponentTotal(ctx context.Context, taskID string, total int) { r.mu.Lock() defer r.mu.Unlock() r.total = total r.totalSet = true } -func (r *recordingProgressSink) OnComponentProgress(ev pipelinepkg.ProgressEvent) { +func (r *recordingProgressSink) OnComponentProgress(ctx context.Context, ev pipelinepkg.ProgressEvent) { r.mu.Lock() defer r.mu.Unlock() r.events = append(r.events, ev) diff --git a/internal/ingestion/task/task_context.go b/internal/ingestion/task/task_context.go index a0e2fcf276..e3759694ac 100644 --- a/internal/ingestion/task/task_context.go +++ b/internal/ingestion/task/task_context.go @@ -57,8 +57,8 @@ func NewTaskContextForScheduling(ctx context.Context, task *entity.IngestionTask // LoadFromIngestionTask loads the full task context from an IngestionTask. // It follows the FK chain: ingestion task -> document -> knowledgebase -> tenant. -func LoadFromIngestionTask(ingestionTask *entity.IngestionTask) (*TaskContext, error) { - doc, err := dao.NewDocumentDAO().GetByID(ingestionTask.DocumentID) +func LoadFromIngestionTask(ctx context.Context, ingestionTask *entity.IngestionTask) (*TaskContext, error) { + doc, err := dao.NewDocumentDAO().GetByID(ctx, dao.DB, ingestionTask.DocumentID) if err != nil { return nil, fmt.Errorf("load document %s: %w", ingestionTask.DocumentID, err) } @@ -79,6 +79,7 @@ func LoadFromIngestionTask(ingestionTask *entity.IngestionTask) (*TaskContext, e pipelineID := resolvePipelineID(doc, kb) return &TaskContext{ + Ctx: ctx, IngestionTask: ingestionTask, PipelineID: pipelineID, Doc: *doc, diff --git a/internal/ingestion/task/task_context_test.go b/internal/ingestion/task/task_context_test.go index 37380de0ea..b8970445fe 100644 --- a/internal/ingestion/task/task_context_test.go +++ b/internal/ingestion/task/task_context_test.go @@ -56,7 +56,8 @@ func TestLoadFromIngestionTask_FallsBackToKnowledgebasePipelineID(t *testing.T) t.Fatalf("create tenant: %v", err) } - ctx, err := LoadFromIngestionTask(&entity.IngestionTask{ + ctx := t.Context() + taskCtx, err := LoadFromIngestionTask(ctx, &entity.IngestionTask{ ID: "task-1", DocumentID: "doc-1", DatasetID: "kb-1", @@ -64,8 +65,8 @@ func TestLoadFromIngestionTask_FallsBackToKnowledgebasePipelineID(t *testing.T) if err != nil { t.Fatalf("LoadFromIngestionTask: %v", err) } - if ctx.PipelineID != kbPipelineID { - t.Fatalf("PipelineID = %q, want %q", ctx.PipelineID, kbPipelineID) + if taskCtx.PipelineID != kbPipelineID { + t.Fatalf("PipelineID = %q, want %q", taskCtx.PipelineID, kbPipelineID) } } @@ -104,7 +105,8 @@ func TestLoadFromIngestionTask_PrefersDocumentPipelineID(t *testing.T) { t.Fatalf("create tenant: %v", err) } - ctx, err := LoadFromIngestionTask(&entity.IngestionTask{ + ctx := t.Context() + taskCtx, err := LoadFromIngestionTask(ctx, &entity.IngestionTask{ ID: "task-1", DocumentID: "doc-1", DatasetID: "kb-1", @@ -112,7 +114,7 @@ func TestLoadFromIngestionTask_PrefersDocumentPipelineID(t *testing.T) { if err != nil { t.Fatalf("LoadFromIngestionTask: %v", err) } - if ctx.PipelineID != docPipelineID { - t.Fatalf("PipelineID = %q, want %q", ctx.PipelineID, docPipelineID) + if taskCtx.PipelineID != docPipelineID { + t.Fatalf("PipelineID = %q, want %q", taskCtx.PipelineID, docPipelineID) } } diff --git a/internal/parser/parser/office_parsers_nocgo_test.go b/internal/parser/parser/office_parsers_nocgo_test.go index c3cacbf690..19cb033221 100644 --- a/internal/parser/parser/office_parsers_nocgo_test.go +++ b/internal/parser/parser/office_parsers_nocgo_test.go @@ -9,7 +9,7 @@ import ( ) func TestOfficeParsers_ParseWithResult_NoCGO(t *testing.T) { - ctx := context.Background() + ctx := t.Context() cases := []struct { name string res ParseResult diff --git a/internal/parser/parser/pdf_parser_nocgo_test.go b/internal/parser/parser/pdf_parser_nocgo_test.go index a1b00d3724..59b8d88b2f 100644 --- a/internal/parser/parser/pdf_parser_nocgo_test.go +++ b/internal/parser/parser/pdf_parser_nocgo_test.go @@ -9,9 +9,8 @@ import ( ) func TestPDFParser_ParseWithResult_NoCGO(t *testing.T) { - ctx := context.Background() pdf := NewPDFParser() - + ctx := t.Context() empty := pdf.ParseWithResult(ctx, "empty.pdf", nil) if empty.Err != nil { t.Fatalf("empty input: want nil err, got %v", empty.Err) diff --git a/internal/service/agent.go b/internal/service/agent.go index 3f94df99a8..9d08ee4f2c 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -796,7 +796,7 @@ func (s *AgentService) PublishAgent(ctx context.Context, userID, canvasID string opts := s.saveOrReplaceVersionOptions(ctx, userID, canvasID, dsl, titleStr, description, true) var row *entity.UserCanvasVersion if err = dao.DB.Transaction(func(tx *gorm.DB) error { - if err := s.canvasDAO.UpdateTx(tx, canvasInstance); err != nil { + if err = s.canvasDAO.UpdateTx(tx, canvasInstance); err != nil { return fmt.Errorf("publish agent %s: update parent: %w", canvasID, err) } saved, err := s.versionDAO.SaveOrReplaceLatestTx(tx, opts) diff --git a/internal/service/ask_service.go b/internal/service/ask_service.go index 9a50b17a7c..942b82d4ce 100644 --- a/internal/service/ask_service.go +++ b/internal/service/ask_service.go @@ -72,7 +72,7 @@ type AskStreamOptions struct { // Retriever abstracts chunk retrieval for AskService. type Retriever interface { - RetrievalTest(req *RetrievalTestRequest, userID string) (*RetrievalTestResponse, error) + RetrievalTest(ctx context.Context, req *RetrievalTestRequest, userID string) (*RetrievalTestResponse, error) } // StreamingLLM abstracts streaming chat for AskService. @@ -159,7 +159,7 @@ func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question req.Page = &page req.Size = &ps - result, err := s.retriever.RetrievalTest(req, userID) + result, err := s.retriever.RetrievalTest(ctx, req, userID) if err != nil { common.Warn("AskService retrieval failed", zap.Error(err)) s.sendOrCancel(out, AskDelta{Kind: AskDeltaError, Value: "retrieval failed"}, ctx) diff --git a/internal/service/ask_service_test.go b/internal/service/ask_service_test.go index 9c7ce3eb19..6b6deb5dfe 100644 --- a/internal/service/ask_service_test.go +++ b/internal/service/ask_service_test.go @@ -32,7 +32,7 @@ type fakeRetriever struct { err error } -func (r *fakeRetriever) RetrievalTest(req *RetrievalTestRequest, userID string) (*RetrievalTestResponse, error) { +func (r *fakeRetriever) RetrievalTest(ctx context.Context, req *RetrievalTestRequest, userID string) (*RetrievalTestResponse, error) { if r.err != nil { return nil, r.err } diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index 8cd8c8f9cb..64a748fd45 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -87,11 +87,11 @@ type ChunkService struct { // startParseDocumentsFunc overrides the DSL start-parse flow. Production // uses service.DocumentService.StartParseDocuments; tests inject a fake // to avoid the MQ publisher. - startParseDocumentsFunc func(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts document.StartParseOptions) error + startParseDocumentsFunc func(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, userID string, opts document.StartParseOptions) error // cancelIngestionTaskFunc overrides the document-parsing cancellation. // Production uses service.DocumentService.CancelDocParse; tests inject // a fake to avoid the MQ publisher. - cancelIngestionTaskFunc func(doc *entity.Document) error + cancelIngestionTaskFunc func(ctx context.Context, doc *entity.Document) error getEmbeddingModelFunc func(string, string) (*models.EmbeddingModel, error) incrementChunkStatsFunc func(string, string, int64, int64, float64) error decrementChunkStatsFunc func(string, string, int64, int64, float64) error @@ -129,7 +129,7 @@ func NewChunkService() *ChunkService { // - Builds doc_aggs by aggregating chunks per document // 7. knowledge graph retrieval (not implemented) // 8. Apply retrieval by children to group child chunks under parent chunks -func (s *ChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { +func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { common.Info("RetrievalTest started", zap.String("userID", userID), zap.Any("kbID", req.Datasets), zap.String("question", req.Question)) common.Debug(fmt.Sprintf("RetrievalTest request:\n"+ @@ -160,8 +160,6 @@ func (s *ChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID s return nil, fmt.Errorf("dataset_ids is required") } - ctx := context.Background() - tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { return nil, fmt.Errorf("failed to get user tenants: %w", err) @@ -504,7 +502,7 @@ func hydrateChunkVectors(ctx context.Context, engine engine.DocEngine, chunks [] } // Get retrieves a chunk by ID -func (s *ChunkService) Get(req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error) { +func (s *ChunkService) Get(ctx context.Context, req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error) { if s.docEngine == nil { return nil, fmt.Errorf("doc engine not initialized") } @@ -513,8 +511,6 @@ func (s *ChunkService) Get(req *service.GetChunkRequest, userID string) (*servic return nil, fmt.Errorf("chunk_id is required") } - ctx := context.Background() - // Get user's tenants tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -599,17 +595,17 @@ const ( docStopParsingInvalidStateErrorCode = "DOC_STOP_PARSING_INVALID_STATE" ) -func (s *ChunkService) cancelAllTasksOfDoc(doc *entity.Document) error { +func (s *ChunkService) cancelAllTasksOfDoc(ctx context.Context, doc *entity.Document) error { cancel := s.cancelIngestionTaskFunc if cancel == nil { cancel = document.NewDocumentService().CancelDocParse } - return cancel(doc) + return cancel(ctx, doc) } -func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { +func (s *ChunkService) StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { if !s.kbDAO.Accessible(datasetID, userID) { - return nil, common.CodeDataError, fmt.Errorf("You don't own the dataset %s", datasetID) + return nil, common.CodeDataError, fmt.Errorf("you don't own the dataset %s", datasetID) } if len(req.DocumentIDs) == 0 { @@ -618,16 +614,16 @@ func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopPar _, err := s.kbDAO.GetByID(datasetID) if err != nil { - return nil, common.CodeDataError, fmt.Errorf("You don't own the dataset %s", datasetID) + return nil, common.CodeDataError, fmt.Errorf("you don't own the dataset %s", datasetID) } docIDs, duplicateMessages := service.CheckDuplicateIDs(req.DocumentIDs, "document") successCount := 0 for _, docID := range docIDs { - doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(docID, datasetID) + doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(ctx, dao.DB, docID, datasetID) if err != nil || doc == nil { - return nil, common.CodeDataError, fmt.Errorf("You don't own the document %s", docID) + return nil, common.CodeDataError, fmt.Errorf("you don't own the document %s", docID) } task, err := dao.NewIngestionTaskDAO().GetByDocumentID(docID) @@ -641,7 +637,7 @@ func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopPar }, common.CodeDataError, fmt.Errorf("%s", docStopParsingInvalidStateMessage) } - if err := s.cancelAllTasksOfDoc(doc); err != nil { + if err = s.cancelAllTasksOfDoc(ctx, doc); err != nil { return nil, common.CodeServerError, err } // CancelDocParse (inside cancelAllTasksOfDoc) already issues @@ -699,16 +695,16 @@ func (s *ChunkService) getKnowledgebaseByID(datasetID string) (*entity.Knowledge return s.kbDAO.GetByID(datasetID) } -func (s *ChunkService) getDocumentsByIDs(docIDs []string) ([]*entity.Document, error) { +func (s *ChunkService) getDocumentsByIDs(ctx context.Context, docIDs []string) ([]*entity.Document, error) { if s.getDocumentsByIDsFunc != nil { return s.getDocumentsByIDsFunc(docIDs) } - return s.documentDAO.GetByIDs(docIDs) + return s.documentDAO.GetByIDs(ctx, dao.DB, docIDs) } -func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) { +func (s *ChunkService) Parse(ctx context.Context, userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) { if !s.accessible(datasetID, userID) { - return nil, common.CodeOperatingError, fmt.Errorf("You don't own the dataset %s.", datasetID) + return nil, common.CodeOperatingError, fmt.Errorf("you don't own the dataset %s", datasetID) } if req == nil || len(req.DocumentIDs) == 0 { return nil, common.CodeDataError, fmt.Errorf("`document_ids` is required") @@ -722,7 +718,7 @@ func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileReq docIDs, duplicateMessages := checkDuplicateIDs(req.DocumentIDs, "document") notFound := make([]string, 0) - docs, err := s.getDocumentsByIDs(docIDs) + docs, err := s.getDocumentsByIDs(ctx, docIDs) if err != nil { return nil, common.CodeServerError, err } @@ -748,7 +744,7 @@ func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileReq // Batch pre-check: refuse the whole request if any document's ingestion // task is non-terminal (RUNNING/STOPPING), so we never partially clean. - if err := (document.NewDocumentService().AssertIngestionTasksTerminal(docIDs)); err != nil { + if err = (document.NewDocumentService().AssertIngestionTasksTerminal(docIDs)); err != nil { return nil, common.CodeDataError, err } @@ -762,7 +758,7 @@ func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileReq for _, docID := range docIDs { doc := docByID[docID] - if err := startParse(doc, kb, userID, document.StartParseOptions{RerunWithDelete: true}); err != nil { + if err = startParse(ctx, doc, kb, userID, document.StartParseOptions{RerunWithDelete: true}); err != nil { return nil, common.CodeServerError, err } successCount++ @@ -781,7 +777,7 @@ func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileReq } // List retrieves chunks for a document -func (s *ChunkService) List(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { +func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { if s.docEngine == nil { return nil, fmt.Errorf("doc engine not initialized") } @@ -790,8 +786,6 @@ func (s *ChunkService) List(req *service.ListChunksRequest, userID string) (*ser return nil, fmt.Errorf("doc_id is required") } - ctx := context.Background() - // Get user's tenants tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -803,7 +797,7 @@ func (s *ChunkService) List(req *service.ListChunksRequest, userID string) (*ser // Get document to find its tenant docDAO := dao.NewDocumentDAO() - doc, err := docDAO.GetByID(req.DocID) + doc, err := docDAO.GetByID(ctx, dao.DB, req.DocID) if err != nil || doc == nil { return nil, fmt.Errorf("document not found") } @@ -1005,7 +999,7 @@ func (s *ChunkService) List(req *service.ListChunksRequest, userID string) (*ser }, nil } -func (s *ChunkService) SwitchChunks(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error { +func (s *ChunkService) SwitchChunks(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error { if s.docEngine == nil { return fmt.Errorf("doc engine not initialized") } @@ -1018,9 +1012,6 @@ func (s *ChunkService) SwitchChunks(userID, datasetID, documentID string, availa return fmt.Errorf("req is null") } - ctx := context.Background() - defer ctx.Done() - // Get user's tenants tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -1044,7 +1035,7 @@ func (s *ChunkService) SwitchChunks(userID, datasetID, documentID string, availa } docDAO := dao.NewDocumentDAO() - doc, err := docDAO.GetByID(documentID) + doc, err := docDAO.GetByID(ctx, dao.DB, documentID) if err != nil || doc == nil { return fmt.Errorf("document not found") } @@ -1069,7 +1060,7 @@ func (s *ChunkService) SwitchChunks(userID, datasetID, documentID string, availa return nil } -func (s *ChunkService) UpdateChunk(req *service.UpdateChunkRequest, userID string) error { +func (s *ChunkService) UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest, userID string) error { if s.docEngine == nil { return fmt.Errorf("doc engine not initialized") } @@ -1078,8 +1069,6 @@ func (s *ChunkService) UpdateChunk(req *service.UpdateChunkRequest, userID strin return fmt.Errorf("chunk_id is required") } - ctx := context.Background() - // Get user's tenants tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -1104,7 +1093,7 @@ func (s *ChunkService) UpdateChunk(req *service.UpdateChunkRequest, userID strin // Verify document belongs to dataset docDAO := dao.NewDocumentDAO() - doc, err := docDAO.GetByID(req.DocumentID) + doc, err := docDAO.GetByID(ctx, dao.DB, req.DocumentID) if err != nil || doc == nil { return fmt.Errorf("document not found") } @@ -1210,7 +1199,7 @@ func (s *ChunkService) UpdateChunk(req *service.UpdateChunkRequest, userID strin return nil } -func (s *ChunkService) RemoveChunks(req *service.RemoveChunksRequest, userID string) (int64, error) { +func (s *ChunkService) RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest, userID string) (int64, error) { if s.docEngine == nil { return 0, fmt.Errorf("doc engine not initialized") } @@ -1219,8 +1208,6 @@ func (s *ChunkService) RemoveChunks(req *service.RemoveChunksRequest, userID str return 0, fmt.Errorf("doc_id is required") } - ctx := context.Background() - // Get user's tenants tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -1232,7 +1219,7 @@ func (s *ChunkService) RemoveChunks(req *service.RemoveChunksRequest, userID str // Verify document exists and belongs to a dataset (do this first to get doc.KbID) docDAO := dao.NewDocumentDAO() - doc, err := docDAO.GetByID(req.DocID) + doc, err := docDAO.GetByID(ctx, dao.DB, req.DocID) if err != nil || doc == nil { return 0, fmt.Errorf("document not found") } @@ -1278,7 +1265,7 @@ func (s *ChunkService) RemoveChunks(req *service.RemoveChunksRequest, userID str } if deletedCount > 0 { - if err := s.decrementChunkStats(req.DocID, doc.KbID, 0, deletedCount, 0); err != nil { + if err = s.decrementChunkStats(req.DocID, doc.KbID, 0, deletedCount, 0); err != nil { return deletedCount, fmt.Errorf("failed to update chunk stats: %w", err) } } @@ -1302,7 +1289,7 @@ func (s *ChunkService) AddChunk(ctx context.Context, req *service.AddChunkReques return nil, addChunkError{code: common.CodeDataError, message: fmt.Sprintf("You don't own the dataset %s.", req.DatasetID)} } - doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(req.DocumentID, req.DatasetID) + doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(ctx, dao.DB, req.DocumentID, req.DatasetID) if err != nil || doc == nil { return nil, addChunkError{code: common.CodeDataError, message: fmt.Sprintf("You don't own the document %s.", req.DocumentID)} } diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index 6127ece923..3eef8aa154 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -146,7 +146,8 @@ func TestParsePrevalidatesDocumentsBeforeMutating(t *testing.T) { documentDAO: dao.NewDocumentDAO(), } - _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{ + ctx := t.Context() + _, code, err := svc.Parse(ctx, userID, datasetID, &service.ParseFileRequest{ DocumentIDs: []string{"doc-1", "missing-doc", "doc-2"}, }) if err == nil { @@ -160,14 +161,14 @@ func TestParsePrevalidatesDocumentsBeforeMutating(t *testing.T) { } var taskCount int64 - if err = dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil { + if err = db.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil { t.Fatalf("count tasks: %v", err) } if taskCount != 1 { t.Fatalf("expected existing task to remain, got %d tasks", taskCount) } - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -185,8 +186,8 @@ func TestParsePrevalidatesDocumentsBeforeMutating(t *testing.T) { func TestParseRejectsInaccessibleDataset(t *testing.T) { svc := newParseTestService(t) svc.accessibleFunc = func(string, string) bool { return false } - - _, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + ctx := t.Context() + _, code, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) if err == nil { t.Fatal("expected parse to fail") } @@ -201,9 +202,9 @@ func TestParseRejectsInaccessibleDataset(t *testing.T) { func TestParseRequiresDocumentIDs(t *testing.T) { svc := newParseTestService(t) svc.accessibleFunc = func(string, string) bool { return true } - + ctx := t.Context() for _, req := range []*service.ParseFileRequest{nil, {DocumentIDs: nil}, {DocumentIDs: []string{}}} { - _, code, err := svc.Parse("user-1", "kb-1", req) + _, code, err := svc.Parse(ctx, "user-1", "kb-1", req) if err == nil { t.Fatal("expected parse to fail") } @@ -220,8 +221,8 @@ func TestParseReturnsDataErrorWhenDatasetMissing(t *testing.T) { svc := newParseTestService(t) svc.accessibleFunc = func(string, string) bool { return true } svc.getKnowledgebaseByIDFunc = func(string) (*entity.Knowledgebase, error) { return nil, nil } - - _, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + ctx := t.Context() + _, code, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) if err == nil { t.Fatal("expected parse to fail") } @@ -243,8 +244,8 @@ func TestParseReturnsServerErrorWhenDocumentsQueryFails(t *testing.T) { svc.getDocumentsByIDsFunc = func([]string) ([]*entity.Document, error) { return nil, queryErr } - - _, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + ctx := t.Context() + _, code, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) if !errors.Is(err, queryErr) { t.Fatalf("expected query error, got %v", err) } @@ -262,12 +263,13 @@ func TestParseRejectsRunningDocument(t *testing.T) { insertChunkTestKB(t, datasetID, userID) insertChunkTestDoc(t, "doc-1", datasetID) running := string(entity.TaskStatusRunning) - if err := dao.DB.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("run", running).Error; err != nil { + if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("run", running).Error; err != nil { t.Fatalf("mark doc running: %v", err) } svc := newParseTestService(t) - _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + ctx := t.Context() + _, code, err := svc.Parse(ctx, userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) if err == nil { t.Fatal("expected parse to fail") } @@ -299,9 +301,10 @@ func TestListSortsChunksByDocumentPosition(t *testing.T) { documentDAO: dao.NewDocumentDAO(), } + ctx := t.Context() page := 1 size := 30 - if _, err := svc.List(&service.ListChunksRequest{ + if _, err := svc.List(ctx, &service.ListChunksRequest{ DatasetID: datasetID, DocID: documentID, Page: &page, @@ -346,10 +349,11 @@ func TestListBuildsMatchTextExprForKeywords(t *testing.T) { userTenantDAO: dao.NewUserTenantDAO(), documentDAO: dao.NewDocumentDAO(), } + ctx := t.Context() page := 2 size := 5 - resp, err := svc.List(&service.ListChunksRequest{ + resp, err := svc.List(ctx, &service.ListChunksRequest{ DatasetID: datasetID, DocID: documentID, Page: &page, @@ -717,8 +721,8 @@ func TestRemoveChunksDecrementsStatsAfterDelete(t *testing.T) { kbDAO: dao.NewKnowledgebaseDAO(), userTenantDAO: dao.NewUserTenantDAO(), } - - deletedCount, err := svc.RemoveChunks(&service.RemoveChunksRequest{ + ctx := t.Context() + deletedCount, err := svc.RemoveChunks(ctx, &service.RemoveChunksRequest{ DocID: "doc-1", ChunkIDs: []string{"chunk-1", "chunk-2", "chunk-3"}, }, "user-1") @@ -741,7 +745,7 @@ func TestRemoveChunksDecrementsStatsAfterDelete(t *testing.T) { t.Fatalf("delete id condition = %#v", engine.deleteChunksCondition["id"]) } - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -782,8 +786,8 @@ func TestRemoveChunksSkipsStatsWhenNothingDeleted(t *testing.T) { return nil }, } - - deletedCount, err := svc.RemoveChunks(&service.RemoveChunksRequest{ + ctx := t.Context() + deletedCount, err := svc.RemoveChunks(ctx, &service.RemoveChunksRequest{ DocID: "doc-1", DeleteAll: true, }, "user-1") @@ -816,8 +820,8 @@ func TestRemoveChunksReturnsStatsError(t *testing.T) { return errors.New("stats update failed") }, } - - deletedCount, err := svc.RemoveChunks(&service.RemoveChunksRequest{ + ctx := t.Context() + deletedCount, err := svc.RemoveChunks(ctx, &service.RemoveChunksRequest{ DocID: "doc-1", ChunkIDs: []string{"chunk-1", "chunk-2"}, }, "user-1") @@ -845,8 +849,8 @@ func TestDecrementChunkStatsClampsCounters(t *testing.T) { if err := svc.decrementChunkStats("doc-1", "kb-1", 5, 7, -1); err != nil { t.Fatalf("decrementChunkStats() error = %v", err) } - - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -873,8 +877,8 @@ func TestDecrementChunkStatsRollsBackWhenKnowledgebaseMissing(t *testing.T) { if err == nil || !strings.Contains(err.Error(), "knowledgebase not found") { t.Fatalf("expected missing knowledgebase error, got %v", err) } - - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -1356,7 +1360,8 @@ func TestSwitchChunksUpdatesDocEngineWithAvailableInt(t *testing.T) { userTenantDAO: dao.NewUserTenantDAO(), } - if err := svc.SwitchChunks("user-1", "kb-1", "doc-1", 0, []string{"chunk-1", "chunk-2"}); err != nil { + ctx := t.Context() + if err = svc.SwitchChunks(ctx, "user-1", "kb-1", "doc-1", 0, []string{"chunk-1", "chunk-2"}); err != nil { t.Fatalf("SwitchChunks() error = %v", err) } @@ -1493,7 +1498,8 @@ func TestChunkServiceParse_RejectsBatchWithRunningIngestionTask(t *testing.T) { svc := newParseTestService(t) svc.accessibleFunc = func(string, string) bool { return true } - _, _, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1", "doc-2"}}) + ctx := t.Context() + _, _, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1", "doc-2"}}) if err == nil { t.Fatal("expected error for RUNNING ingestion task, got nil") } @@ -1510,13 +1516,13 @@ func TestChunkServiceParse_CallsStartParseDocumentsWithRerunWithDelete(t *testin svc.accessibleFunc = func(string, string) bool { return true } var calledOpts document.StartParseOptions var calledDocID string - svc.startParseDocumentsFunc = func(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts document.StartParseOptions) error { + svc.startParseDocumentsFunc = func(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, userID string, opts document.StartParseOptions) error { calledOpts = opts calledDocID = doc.ID return nil } - - _, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + ctx := t.Context() + _, code, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) if err != nil { t.Fatalf("Parse: %v", err) } @@ -1539,11 +1545,12 @@ func TestChunkServiceParse_ReturnsPartialSuccessForDuplicateDocumentIDs(t *testi insertChunkTestDoc(t, "doc-1", datasetID) svc := newParseTestService(t) - svc.startParseDocumentsFunc = func(*entity.Document, *entity.Knowledgebase, string, document.StartParseOptions) error { + svc.startParseDocumentsFunc = func(context.Context, *entity.Document, *entity.Knowledgebase, string, document.StartParseOptions) error { return nil } - result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{ + ctx := t.Context() + result, code, err := svc.Parse(ctx, userID, datasetID, &service.ParseFileRequest{ DocumentIDs: []string{"doc-1", "doc-1"}, }) if err == nil { @@ -1571,12 +1578,12 @@ func TestStopParsing_CallsCancelIngestionTask(t *testing.T) { svc := newParseTestService(t) var calledDocID string - svc.cancelIngestionTaskFunc = func(doc *entity.Document) error { + svc.cancelIngestionTaskFunc = func(ctx context.Context, doc *entity.Document) error { calledDocID = doc.ID return nil } - - _, _, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) + ctx := t.Context() + _, _, err := svc.StopParsing(ctx, "user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) if err != nil { t.Fatalf("StopParsing: %v", err) } @@ -1594,9 +1601,9 @@ func TestStopParsing_RejectsDocumentWithoutIngestionTask(t *testing.T) { svc := newParseTestService(t) var called bool - svc.cancelIngestionTaskFunc = func(doc *entity.Document) error { called = true; return nil } - - _, code, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) + svc.cancelIngestionTaskFunc = func(ctx context.Context, doc *entity.Document) error { called = true; return nil } + ctx := t.Context() + _, code, err := svc.StopParsing(ctx, "user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) if err == nil { t.Fatal("expected error for document without ingestion task") } @@ -1623,21 +1630,21 @@ func TestStopParsing_DoesNotDeleteChunksOrResetCountersAfterCancel(t *testing.T) insertChunkTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1", common.RUNNING) svc := newParseTestService(t) - svc.cancelIngestionTaskFunc = func(doc *entity.Document) error { + svc.cancelIngestionTaskFunc = func(ctx context.Context, doc *entity.Document) error { // Simulate CancelDocParse: set doc.run=CANCEL. return dao.DB.Model(&entity.Document{}).Where("id = ?", doc.ID). Update("run", string(entity.TaskStatusCancel)).Error } engine := &parseTestDocEngine{chunkStoreExists: true} svc.docEngine = engine - - _, _, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) + ctx := t.Context() + _, _, err := svc.StopParsing(ctx, "user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) if err != nil { t.Fatalf("StopParsing: %v", err) } // Counters must NOT be reset. - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("reload doc: %v", err) } diff --git a/internal/service/chunk_types.go b/internal/service/chunk_types.go index 606265de93..bd112ad36c 100644 --- a/internal/service/chunk_types.go +++ b/internal/service/chunk_types.go @@ -253,7 +253,7 @@ func IndexName(uid string) string { return fmt.Sprintf("ragflow_%s", uid) } -func (s *ChunkService) StopParsing(userID, datasetID string, req StopParsingRequest) (map[string]interface{}, common.ErrorCode, error) { +func (s *ChunkService) StopParsing(ctx context.Context, userID, datasetID string, req StopParsingRequest) (map[string]interface{}, common.ErrorCode, error) { if !s.kbDAO.Accessible(datasetID, userID) { return nil, common.CodeAuthenticationError, fmt.Errorf("You don't own the dataset %s", datasetID) } @@ -271,7 +271,8 @@ func (s *ChunkService) StopParsing(userID, datasetID string, req StopParsingRequ successCount := 0 for _, id := range docList { - doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(id, datasetID) + var doc *entity.Document + doc, err = s.documentDAO.GetByDocumentIDAndDatasetID(ctx, dao.DB, id, datasetID) if err != nil { return nil, common.CodeDataError, fmt.Errorf("You don't own the document %s", id) } @@ -293,7 +294,7 @@ func (s *ChunkService) StopParsing(userID, datasetID string, req StopParsingRequ "progress": 0, "chunk_num": 0, } - if err := s.documentDAO.UpdateByID(doc.ID, info); err != nil { + if err = s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, info); err != nil { return nil, common.CodeServerError, fmt.Errorf("failed to update document %s: %w", doc.ID, err) } @@ -343,7 +344,7 @@ type ListChunksResponse struct { } // List retrieves chunks for a document -func (s *ChunkService) List(req *ListChunksRequest, userID string) (*ListChunksResponse, error) { +func (s *ChunkService) List(ctx context.Context, req *ListChunksRequest, userID string) (*ListChunksResponse, error) { if s.docEngine == nil { return nil, fmt.Errorf("doc engine not initialized") } @@ -352,8 +353,6 @@ func (s *ChunkService) List(req *ListChunksRequest, userID string) (*ListChunksR return nil, fmt.Errorf("doc_id is required") } - ctx := context.Background() - // Get user's tenants tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -365,7 +364,7 @@ func (s *ChunkService) List(req *ListChunksRequest, userID string) (*ListChunksR // Get document to find its tenant docDAO := dao.NewDocumentDAO() - doc, err := docDAO.GetByID(req.DocID) + doc, err := docDAO.GetByID(ctx, dao.DB, req.DocID) if err != nil || doc == nil { return nil, fmt.Errorf("document not found") } @@ -516,7 +515,7 @@ type UpdateChunkRequest struct { } // UpdateChunk updates a chunk fields -func (s *ChunkService) UpdateChunk(req *UpdateChunkRequest, userID string) error { +func (s *ChunkService) UpdateChunk(ctx context.Context, req *UpdateChunkRequest, userID string) error { if s.docEngine == nil { return fmt.Errorf("doc engine not initialized") } @@ -525,8 +524,6 @@ func (s *ChunkService) UpdateChunk(req *UpdateChunkRequest, userID string) error return fmt.Errorf("chunk_id is required") } - ctx := context.Background() - // Get user's tenants tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -551,7 +548,7 @@ func (s *ChunkService) UpdateChunk(req *UpdateChunkRequest, userID string) error // Verify document belongs to dataset docDAO := dao.NewDocumentDAO() - doc, err := docDAO.GetByID(req.DocumentID) + doc, err := docDAO.GetByID(ctx, dao.DB, req.DocumentID) if err != nil || doc == nil { return fmt.Errorf("document not found") } @@ -719,7 +716,7 @@ type RemoveChunksRequest struct { // RemoveChunks removes chunks from the dataset table. // If ChunkIDs is empty and DeleteAll is true, removes all chunks for the document. // Otherwise removes only the specified chunks. -func (s *ChunkService) RemoveChunks(req *RemoveChunksRequest, userID string) (int64, error) { +func (s *ChunkService) RemoveChunks(ctx context.Context, req *RemoveChunksRequest, userID string) (int64, error) { if s.docEngine == nil { return 0, fmt.Errorf("doc engine not initialized") } @@ -728,8 +725,6 @@ func (s *ChunkService) RemoveChunks(req *RemoveChunksRequest, userID string) (in return 0, fmt.Errorf("doc_id is required") } - ctx := context.Background() - // Get user's tenants tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { @@ -741,7 +736,7 @@ func (s *ChunkService) RemoveChunks(req *RemoveChunksRequest, userID string) (in // Verify document exists and belongs to a dataset (do this first to get doc.KbID) docDAO := dao.NewDocumentDAO() - doc, err := docDAO.GetByID(req.DocID) + doc, err := docDAO.GetByID(ctx, dao.DB, req.DocID) if err != nil || doc == nil { return 0, fmt.Errorf("document not found") } diff --git a/internal/service/dataset/crud.go b/internal/service/dataset/crud.go index 5fd07e5b07..4347f51839 100644 --- a/internal/service/dataset/crud.go +++ b/internal/service/dataset/crud.go @@ -18,7 +18,7 @@ import ( func (d *DatasetService) CreateDataset(req *service.CreateDatasetRequest, tenantID string) (map[string]interface{}, common.ErrorCode, error) { if !common.IsValidString(req.Name) { - return nil, common.CodeDataError, errors.New("Dataset name must be string.") + return nil, common.CodeDataError, errors.New("dataset name must be string") } name := strings.TrimSpace(req.Name) @@ -31,7 +31,7 @@ func (d *DatasetService) CreateDataset(req *service.CreateDatasetRequest, tenant tenant, err := d.tenantDAO.GetByID(tenantID) if err != nil || tenant == nil { - return nil, common.CodeDataError, errors.New("Tenant not found.") + return nil, common.CodeDataError, errors.New("tenant not found") } if req.ParserID != nil || req.PipelineID != nil || req.ParseType != nil { @@ -120,10 +120,10 @@ func (d *DatasetService) CreateDataset(req *service.CreateDatasetRequest, tenant // Reject duplicate name within tenant to match the established API contract. existing, err := d.kbDAO.GetByName(name, tenantID) if err != nil && !dao.IsNotFoundErr(err) { - return nil, common.CodeServerError, errors.New("Database operation failed") + return nil, common.CodeServerError, errors.New("database operation failed") } if existing != nil { - return nil, common.CodeDataError, fmt.Errorf("Dataset name '%s' already exists", name) + return nil, common.CodeDataError, fmt.Errorf("dataset name '%s' already exists", name) } kb := &entity.Knowledgebase{ @@ -142,14 +142,14 @@ func (d *DatasetService) CreateDataset(req *service.CreateDatasetRequest, tenant if err = d.kbDAO.Create(kb); err != nil { if dao.IsDuplicateKeyErr(err) { - return nil, common.CodeDataError, fmt.Errorf("Dataset name '%s' already exists", name) + return nil, common.CodeDataError, fmt.Errorf("dataset name '%s' already exists", name) } - return nil, common.CodeServerError, errors.New("Failed to save dataset") + return nil, common.CodeServerError, errors.New("failed to save dataset") } createdKB, err := d.kbDAO.GetByID(kbID) if err != nil || createdKB == nil { - return nil, common.CodeServerError, errors.New("Dataset created failed") + return nil, common.CodeServerError, errors.New("dataset created failed") } return datasetToMap(createdKB), common.CodeSuccess, nil @@ -178,7 +178,7 @@ func (d *DatasetService) GetDataset(ctx context.Context, datasetID, userID strin data := datasetToMap(kb) - size, err := d.documentDAO.SumSizeByDatasetID(datasetID) + size, err := d.documentDAO.SumSizeByDatasetID(ctx, dao.DB, datasetID) if err != nil { return nil, common.CodeServerError, errors.New("database operation failed") } @@ -215,7 +215,7 @@ func (d *DatasetService) DeleteDatasets(ids []string, deleteAll bool, tenantID s } kbs, err := d.kbDAO.Query(map[string]interface{}{"tenant_id": tenantID}) if err != nil { - return nil, common.CodeServerError, errors.New("Database operation failed") + return nil, common.CodeServerError, errors.New("database operation failed") } for _, kb := range kbs { normalizedIDs = append(normalizedIDs, kb.ID) @@ -235,7 +235,7 @@ func (d *DatasetService) DeleteDatasets(ids []string, deleteAll bool, tenantID s } if len(unauthorizedIDs) > 0 { return nil, common.CodeDataError, - fmt.Errorf("User '%s' lacks permission for datasets: '%s'", tenantID, strings.Join(unauthorizedIDs, ", ")) + fmt.Errorf("user '%s' lacks permission for datasets: '%s'", tenantID, strings.Join(unauthorizedIDs, ", ")) } successCount := 0 @@ -260,7 +260,7 @@ func (d *DatasetService) deleteDataset(tenantID string, kb *entity.Knowledgebase // transaction (engine ops are not transactional). var documents []entity.Document if err := dao.DB.Where("kb_id = ?", kb.ID).Find(&documents).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } docIDs := extractDocIDs(documents) if len(docIDs) > 0 { @@ -271,30 +271,30 @@ func (d *DatasetService) deleteDataset(tenantID string, kb *entity.Knowledgebase // Delete index tasks referencing this KB. if taskIDs := datasetIndexTaskIDs(kb); len(taskIDs) > 0 { if err := tx.Where("id IN ?", taskIDs).Delete(&entity.Task{}).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } } if len(docIDs) > 0 { var mappings []entity.File2Document if err := tx.Where("document_id IN ?", docIDs).Find(&mappings).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } fileIDs := extractUniqueFileIDs(mappings) if err := tx.Where("doc_id IN ?", docIDs).Delete(&entity.Task{}).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } if err := tx.Where("document_id IN ?", docIDs).Delete(&entity.File2Document{}).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } if len(fileIDs) > 0 { if err := tx.Unscoped().Where("id IN ?", fileIDs).Delete(&entity.File{}).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } } if err := tx.Where("id IN ?", docIDs).Delete(&entity.Document{}).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } } @@ -303,11 +303,11 @@ func (d *DatasetService) deleteDataset(tenantID string, kb *entity.Knowledgebase Where("source_type = ? AND type = ? AND name = ? AND tenant_id = ?", string(entity.FileSourceKnowledgebase), "folder", kb.Name, tenantID). Delete(&entity.File{}).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } if err := tx.Where("id = ?", kb.ID).Delete(&entity.Knowledgebase{}).Error; err != nil { - return fmt.Errorf("Delete dataset error for %s", kb.ID) + return fmt.Errorf("delete dataset error for %s", kb.ID) } return nil }) @@ -324,10 +324,10 @@ func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, order kbs, err := d.kbDAO.GetKBByIDAndUserID(id, userID) if err != nil { - return nil, 0, common.CodeServerError, errors.New("Database operation failed") + return nil, 0, common.CodeServerError, errors.New("database operation failed") } if len(kbs) == 0 { - return nil, 0, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, id) + return nil, 0, common.CodeDataError, fmt.Errorf("user '%s' lacks permission for dataset '%s'", userID, id) } } @@ -335,10 +335,10 @@ func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, order if name != "" { kbs, err := d.kbDAO.GetKBByNameAndUserID(name, userID) if err != nil { - return nil, 0, common.CodeServerError, errors.New("Database operation failed") + return nil, 0, common.CodeServerError, errors.New("database operation failed") } if len(kbs) == 0 { - return nil, 0, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, name) + return nil, 0, common.CodeDataError, fmt.Errorf("user '%s' lacks permission for dataset '%s'", userID, name) } } @@ -367,7 +367,7 @@ func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, order if len(tenantIDs) == 0 { joinedTenants, err := d.tenantDAO.GetJoinedTenantsByUserID(userID) if err != nil { - return nil, 0, common.CodeServerError, errors.New("Database operation failed") + return nil, 0, common.CodeServerError, errors.New("database operation failed") } for _, joinedTenant := range joinedTenants { if joinedTenant == nil || joinedTenant.TenantID == "" { @@ -379,7 +379,7 @@ func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, order kbs, total, err := d.kbDAO.GetByTenantIDs(tenantIDs, userID, page, pageSize, orderby, desc, keywords, parserID, id, name) if err != nil { - return nil, 0, common.CodeServerError, errors.New("Database operation failed") + return nil, 0, common.CodeServerError, errors.New("database operation failed") } data := make([]map[string]interface{}, 0, len(kbs)) diff --git a/internal/service/dataset/index.go b/internal/service/dataset/index.go index 0e5c1cf75e..47ce1aca86 100644 --- a/internal/service/dataset/index.go +++ b/internal/service/dataset/index.go @@ -35,7 +35,7 @@ func checkType(indexType string) bool { return haveType } -func (d *DatasetService) newRaptorOrGraphRagTask(sampleDoc *entity.Document, taskType string, taskDocID string, queueDocID string, docIDs []string) (*entity.Task, map[string]interface{}, error) { +func (d *DatasetService) newRaptorOrGraphRagTask(ctx context.Context, sampleDoc *entity.Document, taskType string, taskDocID string, queueDocID string, docIDs []string) (*entity.Task, map[string]interface{}, error) { if docIDs == nil || len(docIDs) == 0 { docIDs = make([]string, 0) } @@ -43,7 +43,7 @@ func (d *DatasetService) newRaptorOrGraphRagTask(sampleDoc *entity.Document, tas return nil, nil, errors.New("type should be graphrag, raptor or mindmap") } - chunkingConfig, err := d.documentDAO.GetChunkingConfig(sampleDoc.ID) + chunkingConfig, err := d.documentDAO.GetChunkingConfig(ctx, dao.DB, sampleDoc.ID) if err != nil { return nil, nil, err } @@ -124,7 +124,7 @@ func createDatasetIndexTaskInTx(tx *gorm.DB, task *entity.Task, queueDocID strin if task.BeginAt != nil { beginAt = *task.BeginAt } - if err := tx.Model(&entity.Document{}).Where("id = ?", queueDocID).Updates(map[string]interface{}{ + if err = tx.Model(&entity.Document{}).Where("id = ?", queueDocID).Updates(map[string]interface{}{ "progress_msg": "Task is queued...", "process_begin_at": beginAt, }).Error; err != nil { @@ -262,30 +262,30 @@ func clearGraphPhaseMarkers(redisClient *redisengine.Client, datasetID string) { } } -func (d *DatasetService) RunIndex(userID, datasetID, indexType string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) RunIndex(ctx context.Context, userID, datasetID, indexType string) (map[string]interface{}, common.ErrorCode, error) { if !checkType(indexType) { - return nil, common.CodeDataError, fmt.Errorf("Invalid index type '%s'. Must be one of %v", indexType, validIndexTypes) + return nil, common.CodeDataError, fmt.Errorf("invalid index type '%s'. Must be one of %v", indexType, validIndexTypes) } if 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(datasetID, userID) { - return nil, common.CodeDataError, errors.New("No authorization.") + return nil, common.CodeDataError, errors.New("no authorization") } kb, err := d.kbDAO.GetByID(datasetID) if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("Invalid Dataset ID") + return nil, common.CodeDataError, errors.New("invalid Dataset ID") } - return nil, common.CodeDataError, errors.New("Internal server error") + return nil, common.CodeDataError, errors.New("internal server error") } taskType := indexTypeToTaskType[indexType] displayName := indexTypeToDisplayName[indexType] - documents, code, err := d.getDocumentsByDatasetForIndex(datasetID) + documents, code, err := d.getDocumentsByDatasetForIndex(ctx, datasetID) if err != nil { return nil, code, err } @@ -298,17 +298,17 @@ func (d *DatasetService) RunIndex(userID, datasetID, indexType string) (map[stri documentIDs[i] = doc.ID } - task, queueMessage, err := d.newRaptorOrGraphRagTask(sampleDocument, taskType, sampleDocument.ID, graphRaptorQueueDocID, documentIDs) + task, queueMessage, err := d.newRaptorOrGraphRagTask(ctx, sampleDocument, taskType, sampleDocument.ID, graphRaptorQueueDocID, documentIDs) if err != nil { common.Warn("Failed to build dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err)) - return nil, common.CodeDataError, errors.New("Internal server error") + return nil, common.CodeDataError, errors.New("internal server error") } var updatedDocument *entity.Document var dataErr error err = dao.DB.Transaction(func(tx *gorm.DB) error { var lockedKB entity.Knowledgebase - if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + if err = tx.Clauses(clause.Locking{Strength: "UPDATE"}). Where("id = ? AND status = ?", kb.ID, string(entity.StatusValid)). First(&lockedKB).Error; err != nil { return err @@ -324,7 +324,7 @@ func (d *DatasetService) RunIndex(userID, datasetID, indexType string) (map[stri return taskErr } } else if existingTask.Progress != 1 && existingTask.Progress != -1 { - dataErr = fmt.Errorf("Task %s in progress with status %v. A %s Task is already running.", existingTaskID, existingTask.Progress, displayName) + dataErr = fmt.Errorf("task %s in progress with status %v. A %s Task is already running", existingTaskID, existingTask.Progress, displayName) return dataErr } } @@ -340,22 +340,22 @@ func (d *DatasetService) RunIndex(userID, datasetID, indexType string) (map[stri return nil, common.CodeDataError, dataErr } common.Warn("Failed to create dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err)) - return nil, common.CodeDataError, errors.New("Internal server error") + return nil, common.CodeDataError, errors.New("internal server error") } - if err := enqueueDatasetIndexTask(0, queueMessage); err != nil { + if err = enqueueDatasetIndexTask(0, queueMessage); err != nil { if cleanupErr := cleanupFailedDatasetIndexTask(task.ID, updatedDocument, kb.ID, indexType); cleanupErr != nil { err = errors.Join(err, cleanupErr) } common.Warn("Failed to queue dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err)) - return nil, common.CodeDataError, errors.New("Internal server error") + return nil, common.CodeDataError, errors.New("internal server error") } return map[string]interface{}{"task_id": task.ID}, common.CodeSuccess, nil } -func (d *DatasetService) getDocumentsByDatasetForIndex(datasetID string) ([]*entity.Document, common.ErrorCode, error) { - documents, _, err := d.documentDAO.GetByKBID(datasetID) +func (d *DatasetService) getDocumentsByDatasetForIndex(ctx context.Context, datasetID string) ([]*entity.Document, common.ErrorCode, error) { + documents, _, err := d.documentDAO.GetByKBID(ctx, dao.DB, datasetID) if err != nil { common.Warn("Failed to load dataset documents for index", zap.String("dataset_id", datasetID), zap.Error(err)) return nil, common.CodeDataError, errors.New("Internal server error") diff --git a/internal/service/dataset/ingestion.go b/internal/service/dataset/ingestion.go index c6a327ab85..23a4cb3e03 100644 --- a/internal/service/dataset/ingestion.go +++ b/internal/service/dataset/ingestion.go @@ -1,6 +1,7 @@ package dataset import ( + "context" "errors" "fmt" @@ -9,7 +10,7 @@ import ( "ragflow/internal/entity" ) -func (d *DatasetService) GetIngestionSummary(datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) GetIngestionSummary(ctx context.Context, datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) { if datasetID == "" { return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`) } @@ -25,7 +26,7 @@ func (d *DatasetService) GetIngestionSummary(datasetID, userID string) (map[stri return nil, common.CodeServerError, errors.New("Database operation failed") } - status, err := d.documentDAO.GetParsingStatusByKBID(datasetID) + status, err := d.documentDAO.GetParsingStatusByKBID(ctx, dao.DB, datasetID) if err != nil { return nil, common.CodeServerError, errors.New("Database operation failed") } diff --git a/internal/service/dataset/metadata.go b/internal/service/dataset/metadata.go index f3fef6000f..a18b7c68c7 100644 --- a/internal/service/dataset/metadata.go +++ b/internal/service/dataset/metadata.go @@ -1,6 +1,7 @@ package dataset import ( + "context" "errors" "fmt" "strings" @@ -12,14 +13,14 @@ import ( ) // UpdateDocumentMetadataConfig updates the metadata config for a document in a dataset. -func (d *DatasetService) UpdateDocumentMetadataConfig(userID, datasetID, documentID string, req map[string]interface{}) (*entity.Document, common.ErrorCode, error) { +func (d *DatasetService) UpdateDocumentMetadataConfig(ctx context.Context, userID, datasetID, documentID string, req map[string]interface{}) (*entity.Document, common.ErrorCode, error) { userID = strings.TrimSpace(userID) datasetID = strings.TrimSpace(datasetID) if !d.Accessible(datasetID, userID) { return nil, common.CodeDataError, errors.New("You don't own the dataset.") } - doc, err := d.documentDAO.GetByDocumentIDAndDatasetID(documentID, datasetID) + doc, err := d.documentDAO.GetByDocumentIDAndDatasetID(ctx, dao.DB, documentID, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, fmt.Errorf("Document %s not found in dataset %s", documentID, datasetID) @@ -38,11 +39,11 @@ func (d *DatasetService) UpdateDocumentMetadataConfig(userID, datasetID, documen } parserConfig["metadata"] = metadata - if err = d.documentDAO.UpdateByID(doc.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { + if err = d.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { return nil, common.CodeServerError, errors.New("Database operation failed") } - doc, err = d.documentDAO.GetByID(doc.ID) + doc, err = d.documentDAO.GetByID(ctx, dao.DB, doc.ID) if err != nil { return nil, common.CodeServerError, errors.New("Database operation failed") } diff --git a/internal/service/dataset/metadata_config_test.go b/internal/service/dataset/metadata_config_test.go index fbb6d1b35c..3c8f597559 100644 --- a/internal/service/dataset/metadata_config_test.go +++ b/internal/service/dataset/metadata_config_test.go @@ -88,8 +88,10 @@ func TestDatasetServiceUpdateDocumentMetadataConfig(t *testing.T) { insertDatasetMetadataConfigKB(t, "kb-1", "user-1") insertDatasetMetadataConfigDoc(t, "doc-1", "kb-1", entity.JSONMap{"pages": []interface{}{1, 2}}) + ctx := t.Context() metadata := map[string]interface{}{"author": "Alice", "year": float64(2026)} doc, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig( + ctx, "user-1", "kb-1", "doc-1", @@ -116,7 +118,7 @@ func TestDatasetServiceUpdateDocumentMetadataConfig(t *testing.T) { t.Fatalf("unexpected metadata: %#v", updatedMetadata) } - persisted, err := dao.NewDocumentDAO().GetByID("doc-1") + persisted, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("failed to fetch persisted document: %v", err) } @@ -131,7 +133,9 @@ func TestDatasetServiceUpdateDocumentMetadataConfigRequiresMetadata(t *testing.T insertDatasetMetadataConfigKB(t, "kb-1", "user-1") insertDatasetMetadataConfigDoc(t, "doc-1", "kb-1", entity.JSONMap{}) + ctx := t.Context() _, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig( + ctx, "user-1", "kb-1", "doc-1", @@ -154,7 +158,9 @@ func TestDatasetServiceUpdateDocumentMetadataConfigRejectsNonOwner(t *testing.T) insertDatasetMetadataConfigKB(t, "kb-1", "owner-1") insertDatasetMetadataConfigDoc(t, "doc-1", "kb-1", entity.JSONMap{}) + ctx := t.Context() _, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig( + ctx, "user-1", "kb-1", "doc-1", @@ -183,7 +189,9 @@ func TestDatasetServiceUpdateDocumentMetadataConfigAllowsTeamMember(t *testing.T insertDatasetMetadataConfigTeamMember(t, "user-1", "owner-1") insertDatasetMetadataConfigDoc(t, "doc-1", "kb-1", entity.JSONMap{}) + ctx := t.Context() doc, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig( + ctx, "user-1", "kb-1", "doc-1", diff --git a/internal/service/dataset/task_cleanup_test.go b/internal/service/dataset/task_cleanup_test.go index 5ddc0b9abe..5c40f5190a 100644 --- a/internal/service/dataset/task_cleanup_test.go +++ b/internal/service/dataset/task_cleanup_test.go @@ -89,7 +89,8 @@ func TestCleanupFailedDatasetIndexTaskDeletesTaskAndRestoresDocument(t *testing. t.Fatalf("expected task to be deleted, got err=%v task=%#v", err, persistedTask) } - persistedDoc, err := dao.NewDocumentDAO().GetByID(doc.ID) + ctx := t.Context() + persistedDoc, err := dao.NewDocumentDAO().GetByID(ctx, db, doc.ID) if err != nil { t.Fatalf("fetch document: %v", err) } @@ -101,7 +102,7 @@ func TestCleanupFailedDatasetIndexTaskDeletesTaskAndRestoresDocument(t *testing. } var persistedKB entity.Knowledgebase - if err := dao.DB.Where("id = ?", kb.ID).First(&persistedKB).Error; err != nil { + if err = dao.DB.Where("id = ?", kb.ID).First(&persistedKB).Error; err != nil { t.Fatalf("fetch kb: %v", err) } if persistedKB.GraphragTaskID != nil { diff --git a/internal/service/dataset/update.go b/internal/service/dataset/update.go index cde3a7113a..88340d00fd 100644 --- a/internal/service/dataset/update.go +++ b/internal/service/dataset/update.go @@ -135,10 +135,10 @@ func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, tenantID } if req.ParserConfig != nil { - if err := validateDatasetParserConfigSize(req.ParserConfig); err != nil { + if err = validateDatasetParserConfigSize(req.ParserConfig); err != nil { return nil, common.CodeDataError, err } - if err := pipelinepkg.NormalizeParserConfigPages(map[string]any(req.ParserConfig)); err != nil { + if err = pipelinepkg.NormalizeParserConfigPages(req.ParserConfig); err != nil { return nil, common.CodeDataError, err } } @@ -148,7 +148,7 @@ func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, tenantID if pagerankRequested { requestedPagerank = *req.Pagerank if *req.Pagerank < 0 || *req.Pagerank > 100 { - return nil, common.CodeDataError, errors.New("Input should be less than or equal to 100") + return nil, common.CodeDataError, errors.New("input should be less than or equal to 100") } if d.docEngine == nil { return nil, common.CodeServerError, errors.New("document engine is not initialized") diff --git a/internal/service/document/document_artifact.go b/internal/service/document/document_artifact.go index 6a2707c881..8f70b14a8b 100644 --- a/internal/service/document/document_artifact.go +++ b/internal/service/document/document_artifact.go @@ -1,6 +1,7 @@ package document import ( + "context" "fmt" "path/filepath" "strings" @@ -13,7 +14,7 @@ import ( ) // GetDocumentImage retrieves an image object from storage. -func (s *DocumentService) GetDocumentImage(imageID string) ([]byte, error) { +func (s *DocumentService) GetDocumentImage(ctx context.Context, imageID string) ([]byte, error) { parts := strings.SplitN(imageID, "-", 2) if len(parts) != 2 || parts[0] == "" || parts[1] == "" { return nil, fmt.Errorf("Image not found.") @@ -36,7 +37,7 @@ func (s *DocumentService) GetDocumentImage(imageID string) ([]byte, error) { // gate runs BEFORE the storage read so a probe of an unknown // filename cannot distinguish "you cannot see it" from "it // exists" — both return ErrArtifactNotFound. Mirrors PR #16169. -func (s *DocumentService) GetDocumentArtifact(filename, userID string) (*ArtifactResponse, error) { +func (s *DocumentService) GetDocumentArtifact(ctx context.Context, filename, userID string) (*ArtifactResponse, error) { basename := filepath.Base(filename) if basename != filename || strings.Contains(filename, "/") || strings.Contains(filename, "\\") { return nil, ErrArtifactInvalidFilename @@ -192,8 +193,8 @@ func shouldForceArtifactAttachment(ext, contentType string) bool { return ok } -func (s *DocumentService) GetDocumentPreview(docID string) (*DocumentPreview, error) { - doc, err := s.documentDAO.GetByID(docID) +func (s *DocumentService) GetDocumentPreview(ctx context.Context, docID string) (*DocumentPreview, error) { + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil { return nil, err } diff --git a/internal/service/document/document_crud.go b/internal/service/document/document_crud.go index 7c310009e9..a00a564d6c 100644 --- a/internal/service/document/document_crud.go +++ b/internal/service/document/document_crud.go @@ -18,11 +18,11 @@ import ( // the caller has access to. Returns false on any lookup failure or // empty inputs so callers can treat a denial as a 404-equivalent // and avoid leaking whether the document exists at all. -func (s *DocumentService) Accessible(docID, userID string) bool { +func (s *DocumentService) Accessible(ctx context.Context, docID, userID string) bool { if docID == "" || userID == "" { return false } - doc, err := s.documentDAO.GetByID(docID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil || doc == nil { return false } @@ -62,11 +62,11 @@ func (s *DocumentService) GetDocumentStorageAddress(doc *entity.Document) (strin return doc.KbID, *doc.Location, nil } -func (s *DocumentService) DownloadDocument(datasetID, docID string) (*DownloadDocumentResp, error) { +func (s *DocumentService) DownloadDocument(ctx context.Context, datasetID, docID string) (*DownloadDocumentResp, error) { if docID == "" { return nil, fmt.Errorf("Specify document_id please.") } - doc, err := s.documentDAO.GetByID(docID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil || doc.KbID != datasetID { return nil, fmt.Errorf("Document not found!") } @@ -101,8 +101,8 @@ func (s *DocumentService) DownloadDocument(datasetID, docID string) (*DownloadDo } // GetDocumentByID get document by ID -func (s *DocumentService) GetDocumentByID(id string) (*DocumentResponse, error) { - document, err := s.documentDAO.GetByID(id) +func (s *DocumentService) GetDocumentByID(ctx context.Context, id string) (*DocumentResponse, error) { + document, err := s.documentDAO.GetByID(ctx, dao.DB, id) if err != nil { return nil, err } @@ -111,8 +111,8 @@ func (s *DocumentService) GetDocumentByID(id string) (*DocumentResponse, error) } // UpdateDocument update document -func (s *DocumentService) UpdateDocument(id string, req *UpdateDocumentRequest) error { - document, err := s.documentDAO.GetByID(id) +func (s *DocumentService) UpdateDocument(ctx context.Context, id string, req *UpdateDocumentRequest) error { + document, err := s.documentDAO.GetByID(ctx, dao.DB, id) if err != nil { return err } @@ -136,14 +136,14 @@ func (s *DocumentService) UpdateDocument(id string, req *UpdateDocumentRequest) document.ProgressMsg = req.ProgressMsg } - return s.documentDAO.Update(document) + return s.documentDAO.Update(ctx, dao.DB, document) } // IncrementChunkNum atomically increments chunk/token counters on the document and its knowledge base in a transaction -func (s *DocumentService) IncrementChunkNum(docID, kbID string, chunkNum, tokenNum int, duration float64) error { - return dao.DB.Transaction(func(tx *gorm.DB) error { +func (s *DocumentService) IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error { + return dao.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { // Update document - if err := tx.Model(&entity.Document{}). + if err := tx.WithContext(ctx).Model(&entity.Document{}). Where("id = ? AND kb_id = ?", docID, kbID). Updates(map[string]interface{}{ "chunk_num": gorm.Expr("chunk_num + ?", int64(chunkNum)), @@ -154,7 +154,7 @@ func (s *DocumentService) IncrementChunkNum(docID, kbID string, chunkNum, tokenN } // Update knowledgebase - if err := tx.Model(&entity.Knowledgebase{}). + if err := tx.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ "chunk_num": gorm.Expr("chunk_num + ?", int64(chunkNum)), @@ -171,8 +171,8 @@ func (s *DocumentService) IncrementChunkNum(docID, kbID string, chunkNum, tokenN // row so the document-list endpoint (which reads document.progress/run/ // progress_msg) reflects in-flight Go pipeline progress. Best-effort by // design; callers log and continue on error. -func (s *DocumentService) UpdateRunProgress(docID string, progress float64, run, progressMsg string) error { - return s.documentDAO.UpdateByID(docID, map[string]interface{}{ +func (s *DocumentService) UpdateRunProgress(ctx context.Context, docID string, progress float64, run, progressMsg string) error { + return s.documentDAO.UpdateByID(ctx, dao.DB, docID, map[string]interface{}{ "progress": progress, "run": run, "progress_msg": progressMsg, @@ -180,18 +180,18 @@ func (s *DocumentService) UpdateRunProgress(docID string, progress float64, run, } // DeleteDocument delete document — delegates to full cleanup logic. -func (s *DocumentService) DeleteDocument(id string) error { - return s.deleteDocumentFull(id) +func (s *DocumentService) DeleteDocument(ctx context.Context, id string) error { + return s.deleteDocumentFull(ctx, id) } // DeleteDocuments deletes multiple documents under a dataset. // // ids: specific document IDs; deleteAll: delete all docs in the dataset. // Returns the number of successfully deleted documents. -func (s *DocumentService) DeleteDocuments(ids []string, deleteAll bool, datasetID, userID string) (int, error) { +func (s *DocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error) { // 1. Check dataset is accessible by the user if !s.kbDAO.Accessible(datasetID, userID) { - return 0, fmt.Errorf("You don't own the dataset %s.", datasetID) + return 0, fmt.Errorf("you don't own the dataset %s", datasetID) } // 2. Resolve document IDs @@ -211,7 +211,7 @@ func (s *DocumentService) DeleteDocuments(ids []string, deleteAll bool, datasetI // 4. Validate IDs belong to this dataset (only for explicit ids; deleteAll is already scoped) if !deleteAll { - if _, err := s.validateDocsInDataset(ids, datasetID); err != nil { + if _, err := s.validateDocsInDataset(ctx, ids, datasetID); err != nil { return 0, err } } @@ -219,7 +219,7 @@ func (s *DocumentService) DeleteDocuments(ids []string, deleteAll bool, datasetI // 5. Delete each document (non-critical failures are tolerated per doc) deleted := 0 for _, docID := range ids { - if err := s.deleteDocumentFull(docID); err != nil { + if err := s.deleteDocumentFull(ctx, docID); err != nil { common.Warn(fmt.Sprintf("DeleteDocuments: failed to delete %s: %v", docID, err)) continue } @@ -232,8 +232,8 @@ func (s *DocumentService) DeleteDocuments(ids []string, deleteAll bool, datasetI // deleteDocumentFull performs full document cleanup. Non-critical failures // are tolerated (logged and continue). Critical failures (e.g. document or // KB not found) return an error immediately. -func (s *DocumentService) deleteDocumentFull(docID string) error { - doc, kb, err := s.resolveDocAndKB(docID) +func (s *DocumentService) deleteDocumentFull(ctx context.Context, docID string) error { + doc, kb, err := s.resolveDocAndKB(ctx, docID) if err != nil { return err } @@ -254,7 +254,7 @@ func (s *DocumentService) deleteDocumentFull(docID string) error { } s.deleteDocEngineData(docID, kb.TenantID, doc.KbID) - if err := s.deleteDocRecordWithCounters(doc, kb.ID); err != nil { + if err = s.deleteDocRecordWithCounters(ctx, doc, kb.ID); err != nil { return err } s.cleanupFileReferences(docID) @@ -267,8 +267,8 @@ func (s *DocumentService) deleteDocumentFull(docID string) error { // deleting the underlying file record, its storage blob, or its file2document // mappings. Mirrors Python DocumentService.remove_document — the caller is // responsible for cleaning up the file2document mappings separately. -func (s *DocumentService) RemoveDocumentKeepFile(docID string) error { - doc, kb, err := s.resolveDocAndKB(docID) +func (s *DocumentService) RemoveDocumentKeepFile(ctx context.Context, docID string) error { + doc, kb, err := s.resolveDocAndKB(ctx, docID) if err != nil { return err } @@ -276,7 +276,7 @@ func (s *DocumentService) RemoveDocumentKeepFile(docID string) error { common.Logger.Warn(fmt.Sprintf("RemoveDocumentKeepFile: failed to delete tasks for %s: %v", docID, delErr)) } s.deleteDocEngineData(docID, kb.TenantID, doc.KbID) - return s.deleteDocRecordWithCounters(doc, kb.ID) + return s.deleteDocRecordWithCounters(ctx, doc, kb.ID) } // InsertDocument creates a document row and increments the owning KB's doc_num @@ -310,8 +310,8 @@ func (s *DocumentService) InsertDocument(doc *entity.Document) error { // mirroring Python duplicate_name(DocumentService.query, name=..., kb_id=...). // Returns an error when the existing-name lookup fails so callers never write // blind and risk duplicated document names. -func (s *DocumentService) UniqueDocumentName(kbID, name string) (string, error) { - names, err := s.documentDAO.ListNamesByKbID(kbID) +func (s *DocumentService) UniqueDocumentName(ctx context.Context, kbID, name string) (string, error) { + names, err := s.documentDAO.ListNamesByKbID(ctx, dao.DB, kbID) if err != nil { return "", err } @@ -324,8 +324,8 @@ func (s *DocumentService) UniqueDocumentName(kbID, name string) (string, error) // resolveDocAndKB loads the document and its knowledgebase, returning both or // an error. -func (s *DocumentService) resolveDocAndKB(docID string) (*entity.Document, *entity.Knowledgebase, error) { - doc, err := s.documentDAO.GetByID(docID) +func (s *DocumentService) resolveDocAndKB(ctx context.Context, docID string) (*entity.Document, *entity.Knowledgebase, error) { + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil { return nil, nil, fmt.Errorf("document not found: %w", err) } @@ -348,7 +348,7 @@ func (s *DocumentService) deleteDocEngineData(docID, tenantID, kbID string) { common.Logger.Warn(fmt.Sprintf("deleteDocEngineData: failed to delete chunks for %s: %v", docID, delErr)) } if s.metadataSvc != nil { - _ = s.DeleteDocumentAllMetadata(docID) // logs internally + _ = s.DeleteDocumentAllMetadata(ctx, docID) // logs internally } } @@ -356,7 +356,7 @@ func (s *DocumentService) deleteDocEngineData(docID, tenantID, kbID string) { // KB counters in a single transaction. Counters are only decremented when a // document row was actually removed (RowsAffected > 0), guarding against // double-decrement on retries or concurrent deletes. -func (s *DocumentService) deleteDocRecordWithCounters(doc *entity.Document, kbID string) error { +func (s *DocumentService) deleteDocRecordWithCounters(ctx context.Context, doc *entity.Document, kbID string) error { return dao.DB.Transaction(func(tx *gorm.DB) error { result := tx.Where("id = ?", doc.ID).Delete(&entity.Document{}) if result.Error != nil { @@ -383,8 +383,8 @@ func (s *DocumentService) deleteDocRecordWithCounters(doc *entity.Document, kbID }) } -func (s *DocumentService) rollbackAddFileFromKBError(doc *entity.Document, kbID string, err error) error { - if cleanupErr := s.deleteDocRecordWithCounters(doc, kbID); cleanupErr != nil { +func (s *DocumentService) rollbackAddFileFromKBError(ctx context.Context, doc *entity.Document, kbID string, err error) error { + if cleanupErr := s.deleteDocRecordWithCounters(ctx, doc, kbID); cleanupErr != nil { return fmt.Errorf("%w; rollback cleanup failed: %w", err, cleanupErr) } return err diff --git a/internal/service/document/document_dataset_update.go b/internal/service/document/document_dataset_update.go index 696157934b..edc0c19881 100644 --- a/internal/service/document/document_dataset_update.go +++ b/internal/service/document/document_dataset_update.go @@ -19,7 +19,7 @@ import ( "go.uber.org/zap" ) -func (s *DocumentService) BatchUpdateDocumentStatus(userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) { +func (s *DocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) { kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, userID) if err != nil { return nil, common.CodeDataError, fmt.Errorf("You don't own the dataset.") @@ -32,7 +32,7 @@ func (s *DocumentService) BatchUpdateDocumentStatus(userID, datasetID, status st result := make(map[string]interface{}, len(documentIDs)) hasError := false - documents, err := s.documentDAO.GetByIDs(documentIDs) + documents, err := s.documentDAO.GetByIDs(ctx, dao.DB, documentIDs) if err != nil { return nil, common.CodeServerError, fmt.Errorf("failed to fetch documents: %w", err) } @@ -67,7 +67,7 @@ func (s *DocumentService) BatchUpdateDocumentStatus(userID, datasetID, status st if doc.Status != nil { previousStatus = *doc.Status } - if err := s.documentDAO.UpdateByID(docID, map[string]interface{}{"status": status}); err != nil { + if err = s.documentDAO.UpdateByID(ctx, dao.DB, docID, map[string]interface{}{"status": status}); err != nil { result[docID] = map[string]string{"error": "Database error (Document update)!"} hasError = true continue @@ -75,7 +75,7 @@ func (s *DocumentService) BatchUpdateDocumentStatus(userID, datasetID, status st if doc.ChunkNum > 0 { if s.docEngine == nil { - _ = s.documentDAO.UpdateByID(docID, map[string]interface{}{"status": previousStatus}) + _ = s.documentDAO.UpdateByID(ctx, dao.DB, docID, map[string]interface{}{"status": previousStatus}) result[docID] = map[string]string{"error": "Document store update failed: document engine not initialized"} hasError = true continue @@ -88,7 +88,7 @@ func (s *DocumentService) BatchUpdateDocumentStatus(userID, datasetID, status st doc.KbID, ) if err != nil { - _ = s.documentDAO.UpdateByID(docID, map[string]interface{}{"status": previousStatus}) + _ = s.documentDAO.UpdateByID(ctx, dao.DB, docID, map[string]interface{}{"status": previousStatus}) msg := err.Error() if strings.Contains(msg, "3022") { result[docID] = map[string]string{"error": "Document store table missing."} @@ -103,22 +103,22 @@ func (s *DocumentService) BatchUpdateDocumentStatus(userID, datasetID, status st } if hasError { - return result, common.CodeServerError, fmt.Errorf("Partial failure") + return result, common.CodeServerError, fmt.Errorf("partial failure") } return result, common.CodeSuccess, nil } -func (s *DocumentService) UpdateDatasetDocument(userID, datasetID, documentID string, req *UpdateDatasetDocumentRequest, present map[string]bool) (*UpdateDatasetDocumentResponse, common.ErrorCode, error) { +func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *UpdateDatasetDocumentRequest, present map[string]bool) (*UpdateDatasetDocumentResponse, common.ErrorCode, error) { tenantID := userID kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, tenantID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("You don't own the dataset.") } - return nil, common.CodeDataError, errors.New("Can't find this dataset!") + return nil, common.CodeDataError, errors.New("can't find this dataset") } - doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(documentID, datasetID) + doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(ctx, dao.DB, documentID, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("The dataset doesn't own the document.") @@ -126,18 +126,18 @@ func (s *DocumentService) UpdateDatasetDocument(userID, datasetID, documentID st return nil, common.CodeServerError, err } - if code, err := s.validateDatasetDocumentUpdate(datasetID, documentID, userID, doc, req, present); err != nil { + if code, err := s.validateDatasetDocumentUpdate(ctx, datasetID, documentID, userID, doc, req, present); err != nil { return nil, code, err } if present["meta_fields"] { - if err := s.replaceDocumentMetadata(documentID, req.MetaFields); err != nil { + if err = s.replaceDocumentMetadata(ctx, documentID, req.MetaFields); err != nil { return nil, common.CodeDataError, err } } if present["name"] && req.Name != nil && (doc.Name == nil || *req.Name != *doc.Name) { - if err := s.updateDocumentNameOnly(doc, kb.TenantID, *req.Name); err != nil { + if err = s.updateDocumentNameOnly(ctx, doc, kb.TenantID, *req.Name); err != nil { return nil, common.CodeDataError, err } } @@ -153,19 +153,20 @@ func (s *DocumentService) UpdateDatasetDocument(userID, datasetID, documentID st if present["parser_config"] && req.ParserConfig != nil { // Normalize "pages" ranges before persistence. Invalid ranges are // rejected (fail-fast) — the request is aborted with a clear error. - if err := pipelinepkg.NormalizeParserConfigPages(map[string]any(req.ParserConfig)); err != nil { + if err = pipelinepkg.NormalizeParserConfigPages(req.ParserConfig); err != nil { return nil, common.CodeDataError, err } - dslJSON, err := service.LoadPipelineDSL(isPipeline, effParserID, effPipelineID) + var dslJSON []byte + dslJSON, err = service.LoadPipelineDSL(isPipeline, effParserID, effPipelineID) if err != nil { common.Warn("cleanAndUpdateDocumentParserConfig: failed to load DSL, falling back to merge", zap.Error(err)) - if err := s.updateDocumentParserConfig(doc.ID, req.ParserConfig); err != nil { + if err = s.updateDocumentParserConfig(ctx, doc.ID, req.ParserConfig); err != nil { return nil, common.CodeDataError, err } } else { - cleaned := pipelinepkg.BuildParserConfig(dslJSON, map[string]interface{}(req.ParserConfig)) - if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{ + cleaned := pipelinepkg.BuildParserConfig(dslJSON, req.ParserConfig) + if err = s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{ "parser_config": cleaned, }); err != nil { return nil, common.CodeDataError, err @@ -196,49 +197,49 @@ func (s *DocumentService) UpdateDatasetDocument(userID, datasetID, documentID st } } if reparseParserID != nil || reparsePipelineID != nil { - if err := s.resetDocumentForReparse(doc, kb.TenantID, reparseParserID, reparsePipelineID); err != nil { + if err = s.resetDocumentForReparse(ctx, doc, kb.TenantID, reparseParserID, reparsePipelineID); err != nil { return nil, common.CodeDataError, err } } if present["enabled"] && req.Enabled != nil { - if err := s.updateDocumentStatusOnly(doc, kb, *req.Enabled); err != nil { + if err = s.updateDocumentStatusOnly(ctx, doc, kb, *req.Enabled); err != nil { return nil, common.CodeServerError, err } } - updatedDoc, err := s.documentDAO.GetByID(doc.ID) + updatedDoc, err := s.documentDAO.GetByID(ctx, dao.DB, doc.ID) if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, fmt.Errorf("Can not get document by id:%s", doc.ID) + return nil, common.CodeDataError, fmt.Errorf("can not get document by id:%s", doc.ID) } - return nil, common.CodeDataError, errors.New("Database operation failed") + return nil, common.CodeDataError, errors.New("database operation failed") } metaFields := map[string]interface{}{} if s.docEngine != nil && s.metadataSvc != nil { - metaFields, _ = s.GetDocumentMetadataByID(updatedDoc.ID) + metaFields, _ = s.GetDocumentMetadataByID(ctx, updatedDoc.ID) } return s.toUpdateDatasetDocumentResponse(updatedDoc, metaFields), common.CodeSuccess, nil } -func (s *DocumentService) validateDatasetDocumentUpdate(datasetID, documentID, userID string, doc *entity.Document, req *UpdateDatasetDocumentRequest, present map[string]bool) (common.ErrorCode, error) { +func (s *DocumentService) validateDatasetDocumentUpdate(ctx context.Context, datasetID, documentID, userID string, doc *entity.Document, req *UpdateDatasetDocumentRequest, present map[string]bool) (common.ErrorCode, error) { if req == nil { - return common.CodeDataError, errors.New("Invalid request payload") + return common.CodeDataError, errors.New("invalid request payload") } if present["chunk_count"] && req.ChunkCount != nil && *req.ChunkCount != 0 && *req.ChunkCount != doc.ChunkNum { - return common.CodeDataError, errors.New("Can't change `chunk_count`.") + return common.CodeDataError, errors.New("can't change `chunk_count`") } if present["token_count"] && req.TokenCount != nil && *req.TokenCount != 0 && *req.TokenCount != doc.TokenNum { - return common.CodeDataError, errors.New("Can't change `token_count`.") + return common.CodeDataError, errors.New("can't change `token_count`") } if present["progress"] && req.Progress != nil { if *req.Progress > 1 { return common.CodeDataError, fmt.Errorf("Field: - Message: - Value: <%v>", *req.Progress) } if *req.Progress != 0 && math.Abs(*req.Progress-doc.Progress) > 1e-9 { - return common.CodeDataError, errors.New("Can't change `progress`.") + return common.CodeDataError, errors.New("can't change `progress`") } } @@ -258,12 +259,12 @@ func (s *DocumentService) validateDatasetDocumentUpdate(datasetID, documentID, u if isBuiltin && present["parser_id"] && req.ParserID != nil { parserID := strings.TrimSpace(*req.ParserID) if (doc.Type == "visual" && parserID != "picture") || (isPresentationFile(doc.Name) && parserID != "presentation") { - return common.CodeDataError, errors.New("Not supported yet!") + return common.CodeDataError, errors.New("not supported yet") } } } if present["name"] && req.Name != nil { - if err := s.validateDocumentName(doc, *req.Name); err != nil { + if err := s.validateDocumentName(ctx, doc, *req.Name); err != nil { return common.CodeDataError, err } } @@ -277,12 +278,12 @@ func (s *DocumentService) validateDatasetDocumentUpdate(datasetID, documentID, u return common.CodeSuccess, nil } -func (s *DocumentService) validateDocumentName(doc *entity.Document, newName string) error { +func (s *DocumentService) validateDocumentName(ctx context.Context, doc *entity.Document, newName string) error { if strings.TrimSpace(newName) == "" { - return errors.New("File name can't be empty.") + return errors.New("file name can't be empty") } if len([]byte(newName)) > 255 { - return errors.New("File name must be 255 bytes or less.") + return errors.New("file name must be 255 bytes or less") } oldName := "" @@ -291,16 +292,16 @@ func (s *DocumentService) validateDocumentName(doc *entity.Document, newName str } if strings.ToLower(filepath.Ext(newName)) != strings.ToLower(filepath.Ext(oldName)) { - return errors.New("The extension of file can't be changed") + return errors.New("the extension of file can't be changed") } - docs, err := s.documentDAO.GetByNameAndKBID(newName, doc.KbID) + docs, err := s.documentDAO.GetByNameAndKBID(ctx, dao.DB, newName, doc.KbID) if err != nil { return err } for _, d := range docs { if d.ID != doc.ID && d.Name != nil && *d.Name == newName { - return errors.New("Duplicated document name in the same dataset.") + return errors.New("duplicated document name in the same dataset") } } @@ -330,20 +331,20 @@ func validateMetaFields(meta map[string]any) error { case string, float64, int, int64, float32: continue default: - return fmt.Errorf("The type is not supported in list: %v", typed) + return fmt.Errorf("the type is not supported in list: %v", typed) } } default: - return fmt.Errorf("The type is not supported: %v", v) + return fmt.Errorf("the type is not supported: %v", v) } } return nil } -func (s *DocumentService) updateDocumentNameOnly(doc *entity.Document, tenantID, newName string) error { - if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{"name": newName}); err != nil { - return errors.New("Database error (Document rename)!") +func (s *DocumentService) updateDocumentNameOnly(ctx context.Context, doc *entity.Document, tenantID, newName string) error { + if err := s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{"name": newName}); err != nil { + return errors.New("database error (Document rename)") } mappings, err := s.file2DocumentDAO.GetByDocumentID(doc.ID) @@ -371,22 +372,22 @@ func (s *DocumentService) updateDocumentNameOnly(doc *entity.Document, tenantID, ) } -func (s *DocumentService) updateDocumentParserConfig(documentID string, config map[string]any) error { +func (s *DocumentService) updateDocumentParserConfig(ctx context.Context, documentID string, config map[string]any) error { if len(config) == 0 { return nil } - doc, err := s.documentDAO.GetByID(documentID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, documentID) if err != nil { - return fmt.Errorf("Document(%s) not found.", documentID) + return fmt.Errorf("document(%s) not found", documentID) } - merged := common.DeepMergeMaps(map[string]interface{}(doc.ParserConfig), map[string]interface{}(config)) + merged := common.DeepMergeMaps(doc.ParserConfig, config) if _, ok := config["raptor"]; !ok { delete(merged, "raptor") } - return s.documentDAO.UpdateByID(documentID, map[string]interface{}{ + return s.documentDAO.UpdateByID(ctx, dao.DB, documentID, map[string]interface{}{ "parser_config": entity.JSONMap(merged), }) } @@ -401,7 +402,7 @@ func (s *DocumentService) toUpdateDatasetDocumentResponse(doc *entity.Document, DatasetID: doc.KbID, ParserID: doc.ParserID, PipelineID: doc.PipelineID, - ParserConfig: map[string]interface{}(doc.ParserConfig), + ParserConfig: doc.ParserConfig, SourceType: doc.SourceType, Type: doc.Type, CreatedBy: doc.CreatedBy, diff --git a/internal/service/document/document_ingest.go b/internal/service/document/document_ingest.go index 92bed3ed02..29109b03c7 100644 --- a/internal/service/document/document_ingest.go +++ b/internal/service/document/document_ingest.go @@ -3,18 +3,19 @@ package document import ( "context" "fmt" + "ragflow/internal/dao" "ragflow/internal/service" "ragflow/internal/common" "ragflow/internal/entity" ) -func (s *DocumentService) ListIngestionTasks(userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) { +func (s *DocumentService) ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) { return s.ingestionTaskSvc.ListByUser(userID, datasetID, page, pageSize) } -func (s *DocumentService) IngestDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) { - responses, err := s.ingestionTaskSvc.CreateForDocuments(datasetID, userID, docIDs) +func (s *DocumentService) IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) { + responses, err := s.ingestionTaskSvc.CreateForDocuments(ctx, datasetID, userID, docIDs) if err != nil { return nil, err } @@ -22,18 +23,18 @@ func (s *DocumentService) IngestDocuments(datasetID, userID string, docIDs []str return responses, nil } -func (s *DocumentService) StopIngestionTasks(tasks []string, userID string) ([]*entity.IngestionTask, error) { +func (s *DocumentService) StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error) { return s.ingestionTaskSvc.RequestStopMany(tasks, &userID) } -func (s *DocumentService) RemoveIngestionTasks(tasks []string, userID string) ([]map[string]string, error) { +func (s *DocumentService) RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error) { return s.ingestionTaskSvc.RemoveMany(tasks, &userID) } -func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (common.ErrorCode, error) { +func (s *DocumentService) Ingest(ctx context.Context, userID string, req *IngestDocumentRequest) (common.ErrorCode, error) { run := fmt.Sprint(req.Run) - docs, err := s.documentDAO.GetByIDs(req.DocIDs) + docs, err := s.documentDAO.GetByIDs(ctx, dao.DB, req.DocIDs) if err != nil { return common.CodeExceptionError, fmt.Errorf("fail to get documents: %s", err.Error()) } @@ -72,7 +73,7 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com // Batch pre-check for re-parse with delete: use the validated doc IDs // so we don't silently skip non-existent or unauthorized documents. if run == string(entity.TaskStatusRunning) && req.Delete { - if err := s.AssertIngestionTasksTerminal(validatedIDs); err != nil { + if err = s.AssertIngestionTasksTerminal(validatedIDs); err != nil { return common.CodeDataError, err } } @@ -85,7 +86,7 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com // document run status is set by service.IngestionTaskService.StartRunning // when the task transitions from CREATED, not here. if run == string(entity.TaskStatusRunning) { - if err := s.StartParseDocuments(doc, kb, userID, StartParseOptions{ + if err = s.StartParseDocuments(ctx, doc, kb, userID, StartParseOptions{ ApplyKB: req.ApplyKB, RerunWithDelete: req.Delete, }); err != nil { @@ -101,11 +102,11 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com // worker detects STOPPING and transitions to STOPPED, the task // is terminal and can be safely cleaned up. if run == string(entity.TaskStatusCancel) { - if err := s.CancelDocParse(doc); err != nil { + if err = s.CancelDocParse(ctx, doc); err != nil { common.Error(fmt.Sprintf("go side, start to process %s, run is cancel", doc.ID), err) return common.CodeDataError, err } - if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{ + if err = s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{ "run": string(entity.TaskStatusCancel), "progress": 0, }); err != nil { @@ -117,7 +118,7 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com // Delete-only: user asked to remove prior parse results without // starting a new parse. RUNNING already continued above. - if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{ + if err = s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{ "run": run, "progress": 0, }); err != nil { diff --git a/internal/service/document/document_list.go b/internal/service/document/document_list.go index 580095fefc..85a9b34de0 100644 --- a/internal/service/document/document_list.go +++ b/internal/service/document/document_list.go @@ -1,6 +1,7 @@ package document import ( + "context" "fmt" "strings" "time" @@ -21,9 +22,9 @@ func escapeSQLLikePattern(s string) string { } // ListDocuments list documents -func (s *DocumentService) ListDocuments(page, pageSize int) ([]*DocumentResponse, int64, error) { +func (s *DocumentService) ListDocuments(ctx context.Context, page, pageSize int) ([]*DocumentResponse, int64, error) { offset := (page - 1) * pageSize - documents, total, err := s.documentDAO.List(offset, pageSize) + documents, total, err := s.documentDAO.List(ctx, dao.DB, offset, pageSize) if err != nil { return nil, 0, err } @@ -36,7 +37,7 @@ func (s *DocumentService) ListDocuments(page, pageSize int) ([]*DocumentResponse return responses, total, nil } -func (s *DocumentService) GetThumbnails(userID string, docIDs []string) (map[string]string, error) { +func (s *DocumentService) GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error) { if len(docIDs) == 0 { return map[string]string{}, nil } @@ -50,7 +51,7 @@ func (s *DocumentService) GetThumbnails(userID string, docIDs []string) (map[str tenantIDs = append(tenantIDs, ids...) } - documents, err := s.documentDAO.GetByIDsAndTenantIDs(docIDs, tenantIDs) + documents, err := s.documentDAO.GetByIDsAndTenantIDs(ctx, dao.DB, docIDs, tenantIDs) if err != nil { return nil, fmt.Errorf("failed to fetch document thumbnails: %w", err) } @@ -81,8 +82,8 @@ func (s *DocumentService) GetThumbnails(userID string, docIDs []string) (map[str } // ListDocumentsByDatasetID list documents by knowledge base ID -func (s *DocumentService) ListDocumentsByDatasetID(kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { - return s.ListDocumentsByDatasetIDWithOptions(dao.DocumentListOptions{ +func (s *DocumentService) ListDocumentsByDatasetID(ctx context.Context, kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { + return s.ListDocumentsByDatasetIDWithOptions(ctx, dao.DocumentListOptions{ KbID: kbID, Keywords: keywords, OrderBy: "create_time", @@ -91,13 +92,13 @@ func (s *DocumentService) ListDocumentsByDatasetID(kbID, keywords string, page, } // ListDocumentsByDatasetIDWithOptions lists documents by knowledge base ID with filters. -func (s *DocumentService) ListDocumentsByDatasetIDWithOptions(opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { +func (s *DocumentService) ListDocumentsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { opts.Offset = (page - 1) * pageSize opts.Limit = pageSize if opts.OrderBy == "" { opts.OrderBy = "create_time" } - documents, total, err := s.documentDAO.ListByKBIDWithOptions(opts) + documents, total, err := s.documentDAO.ListByKBIDWithOptions(ctx, dao.DB, opts) if err != nil { return nil, 0, err } @@ -111,16 +112,16 @@ func (s *DocumentService) ListDocumentsByDatasetIDWithOptions(opts dao.DocumentL } // GetDocumentFiltersByDatasetID returns aggregate filter values for documents in a dataset. -func (s *DocumentService) GetDocumentFiltersByDatasetID(opts dao.DocumentListOptions) (map[string]interface{}, int64, error) { - filters, total, err := s.documentDAO.GetFilterByKBID(opts) +func (s *DocumentService) GetDocumentFiltersByDatasetID(ctx context.Context, opts dao.DocumentListOptions) (map[string]interface{}, int64, error) { + filters, total, err := s.documentDAO.GetFilterByKBID(ctx, dao.DB, opts) if err != nil { return nil, 0, err } - docIDs, err := s.documentDAO.ListIDsByKBIDWithOptions(opts) + docIDs, err := s.documentDAO.ListIDsByKBIDWithOptions(ctx, dao.DB, opts) if err != nil { return nil, 0, err } - metadataFilter, err := s.getDocumentMetadataFilter(opts.KbID, docIDs) + metadataFilter, err := s.getDocumentMetadataFilter(ctx, opts.KbID, docIDs) if err != nil { return nil, 0, err } @@ -128,8 +129,8 @@ func (s *DocumentService) GetDocumentFiltersByDatasetID(opts dao.DocumentListOpt return filters, total, nil } -func (s *DocumentService) getDocumentMetadataFilter(kbID string, docIDs []string) (map[string]interface{}, error) { - metadataByKey, err := s.GetMetadataByKBs([]string{kbID}) +func (s *DocumentService) getDocumentMetadataFilter(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error) { + metadataByKey, err := s.GetMetadataByKBs(ctx, []string{kbID}) if err != nil { return nil, err } @@ -164,14 +165,14 @@ func (s *DocumentService) getDocumentMetadataFilter(kbID string, docIDs []string } // ListDocumentIDsByDatasetIDWithOptions lists matching document IDs without pagination. -func (s *DocumentService) ListDocumentIDsByDatasetIDWithOptions(opts dao.DocumentListOptions) ([]string, error) { - return s.documentDAO.ListIDsByKBIDWithOptions(opts) +func (s *DocumentService) ListDocumentIDsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions) ([]string, error) { + return s.documentDAO.ListIDsByKBIDWithOptions(ctx, dao.DB, opts) } // GetDocumentsByAuthorID get documents by author ID -func (s *DocumentService) GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*DocumentResponse, int64, error) { +func (s *DocumentService) GetDocumentsByAuthorID(ctx context.Context, authorID, page, pageSize int) ([]*DocumentResponse, int64, error) { offset := (page - 1) * pageSize - documents, total, err := s.documentDAO.GetByAuthorID(fmt.Sprintf("%d", authorID), offset, pageSize) + documents, total, err := s.documentDAO.GetByAuthorID(ctx, dao.DB, fmt.Sprintf("%d", authorID), offset, pageSize) if err != nil { return nil, 0, err } diff --git a/internal/service/document/document_metadata.go b/internal/service/document/document_metadata.go index 8c70d4e658..14fa8cdbd7 100644 --- a/internal/service/document/document_metadata.go +++ b/internal/service/document/document_metadata.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "ragflow/internal/dao" "ragflow/internal/service" "reflect" "regexp" @@ -18,7 +19,7 @@ import ( ) // GetMetadataSummary get metadata summary for documents -func (s *DocumentService) GetMetadataSummary(kbID string, docIDs []string) (map[string]interface{}, error) { +func (s *DocumentService) GetMetadataSummary(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error) { tenantID, err := s.metadataSvc.GetTenantIDByKBID(kbID) if err != nil { return nil, err @@ -34,9 +35,9 @@ func (s *DocumentService) GetMetadataSummary(kbID string, docIDs []string) (map[ } // SetDocumentMetadata sets metadata for a document in the document engine -func (s *DocumentService) SetDocumentMetadata(docID string, meta map[string]interface{}) error { +func (s *DocumentService) SetDocumentMetadata(ctx context.Context, docID string, meta map[string]interface{}) error { // Get document to find kb_id - doc, err := s.documentDAO.GetByID(docID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil { return fmt.Errorf("document not found: %w", err) } @@ -47,7 +48,7 @@ func (s *DocumentService) SetDocumentMetadata(docID string, meta map[string]inte return fmt.Errorf("failed to get tenant ID: %w", err) } - if err := s.docEngine.UpdateMetadata(context.Background(), docID, doc.KbID, meta, tenantID); err != nil { + if err = s.docEngine.UpdateMetadata(ctx, docID, doc.KbID, meta, tenantID); err != nil { return fmt.Errorf("failed to update metadata: %w", err) } @@ -55,9 +56,9 @@ func (s *DocumentService) SetDocumentMetadata(docID string, meta map[string]inte } // DeleteDocumentMetadata deletes metadata keys for a document in the document engine -func (s *DocumentService) DeleteDocumentMetadata(docID string, keys []string) error { +func (s *DocumentService) DeleteDocumentMetadata(ctx context.Context, docID string, keys []string) error { // Get document to find kb_id - doc, err := s.documentDAO.GetByID(docID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil { return fmt.Errorf("document not found: %w", err) } @@ -78,9 +79,9 @@ func (s *DocumentService) DeleteDocumentMetadata(docID string, keys []string) er } // DeleteDocumentAllMetadata deletes all metadata for a document in the document engine -func (s *DocumentService) DeleteDocumentAllMetadata(docID string) error { +func (s *DocumentService) DeleteDocumentAllMetadata(ctx context.Context, docID string) error { // Get document to find kb_id - doc, err := s.documentDAO.GetByID(docID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil { return fmt.Errorf("document not found: %w", err) } @@ -107,9 +108,9 @@ func (s *DocumentService) DeleteDocumentAllMetadata(docID string) error { } // GetDocumentMetadataByID get metadata for a specific document -func (s *DocumentService) GetDocumentMetadataByID(docID string) (map[string]interface{}, error) { +func (s *DocumentService) GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]interface{}, error) { // Get document to find kb_id - doc, err := s.documentDAO.GetByID(docID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil { return nil, fmt.Errorf("document not found: %w", err) } @@ -134,7 +135,7 @@ func (s *DocumentService) GetDocumentMetadataByID(docID string) (map[string]inte } // GetMetadataByKBs get metadata for knowledge bases -func (s *DocumentService) GetMetadataByKBs(kbIDs []string) (map[string]interface{}, error) { +func (s *DocumentService) GetMetadataByKBs(ctx context.Context, kbIDs []string) (map[string]interface{}, error) { if len(kbIDs) == 0 { return make(map[string]interface{}), nil } @@ -499,17 +500,17 @@ func isTimeString(s string) bool { return matched } -func (s *DocumentService) replaceDocumentMetadata(docID string, meta map[string]any) error { +func (s *DocumentService) replaceDocumentMetadata(ctx context.Context, docID string, meta map[string]any) error { if s.docEngine == nil || s.metadataSvc == nil { return nil } - if err := s.DeleteDocumentAllMetadata(docID); err != nil { + if err := s.DeleteDocumentAllMetadata(ctx, docID); err != nil { return err } - return s.SetDocumentMetadata(docID, map[string]interface{}(meta)) + return s.SetDocumentMetadata(ctx, docID, meta) } -func (s *DocumentService) patchDocumentMetadata(docID string, before, after map[string]interface{}) error { +func (s *DocumentService) patchDocumentMetadata(ctx context.Context, docID string, before, after map[string]interface{}) error { if s.docEngine == nil || s.metadataSvc == nil { return nil } @@ -521,7 +522,7 @@ func (s *DocumentService) patchDocumentMetadata(docID string, before, after map[ } } if len(deleteKeys) > 0 { - if err := s.DeleteDocumentMetadata(docID, deleteKeys); err != nil { + if err := s.DeleteDocumentMetadata(ctx, docID, deleteKeys); err != nil { return err } } @@ -535,13 +536,14 @@ func (s *DocumentService) patchDocumentMetadata(docID string, before, after map[ if len(updateFields) == 0 { return nil } - return s.SetDocumentMetadata(docID, updateFields) + return s.SetDocumentMetadata(ctx, docID, updateFields) } // BatchUpdateDocumentMetadatas implements the shared logic for // PATCH /datasets/:dataset_id/documents/metadatas and // POST /datasets/:dataset_id/metadata/update. func (s *DocumentService) BatchUpdateDocumentMetadatas( + ctx context.Context, datasetID string, selector *DocumentMetadataSelector, updates []DocumentMetadataUpdate, @@ -559,7 +561,7 @@ func (s *DocumentService) BatchUpdateDocumentMetadatas( if len(selector.DocumentIDs) > 0 { // Validate that supplied IDs actually belong to this dataset. - allRows, err := s.documentDAO.GetAllDocIDsByKBIDs([]string{datasetID}) + allRows, err := s.documentDAO.GetAllDocIDsByKBIDs(ctx, dao.DB, []string{datasetID}) if err != nil { return nil, common.CodeServerError, fmt.Errorf("failed to list dataset documents: %w", err) } @@ -628,7 +630,7 @@ func (s *DocumentService) BatchUpdateDocumentMetadatas( // semantics instead of a simple merge-then-delete. updated := 0 for _, docID := range ids { - currentMeta, err := s.GetDocumentMetadataByID(docID) + currentMeta, err := s.GetDocumentMetadataByID(ctx, docID) if err != nil { common.Warn("BatchUpdateDocumentMetadata: get metadata failed", zap.String("docID", docID), zap.Error(err)) @@ -647,7 +649,7 @@ func (s *DocumentService) BatchUpdateDocumentMetadatas( continue } - if err := s.patchDocumentMetadata(docID, originalMeta, meta); err != nil { + if err = s.patchDocumentMetadata(ctx, docID, originalMeta, meta); err != nil { common.Warn("BatchUpdateDocumentMetadata: patch metadata failed", zap.String("docID", docID), zap.Error(err)) continue diff --git a/internal/service/document/document_parse.go b/internal/service/document/document_parse.go index d556ed48d8..ad5ea0a529 100644 --- a/internal/service/document/document_parse.go +++ b/internal/service/document/document_parse.go @@ -25,7 +25,7 @@ import ( // sets it to RUNNING when the worker picks up the task and transitions it from // CREATED. Extracted from Ingest so other entry points (e.g. ChunkService.Parse) // can reuse the same start-parse flow. -func (s *DocumentService) StartParseDocuments(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts StartParseOptions) error { +func (s *DocumentService) StartParseDocuments(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, userID string, opts StartParseOptions) error { // 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. @@ -39,7 +39,7 @@ func (s *DocumentService) StartParseDocuments(doc *entity.Document, kb *entity.K } } - if _, err := s.IngestDocuments(doc.KbID, userID, []string{doc.ID}); err != nil { + if _, err := s.IngestDocuments(ctx, doc.KbID, userID, []string{doc.ID}); err != nil { return err } return nil @@ -171,7 +171,7 @@ func (s *DocumentService) clearKBChunkNumWhenRerun(doc *entity.Document) error { }).Error } -func (s *DocumentService) ParseDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) { +func (s *DocumentService) ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) { // create document parse id // save to task table // send to message queue @@ -186,7 +186,7 @@ func (s *DocumentService) ParseDocuments(datasetID, userID string, docIDs []stri // query database, if the document ids are valid for _, docID := range uniqueDocIDs { - doc, err := s.documentDAO.GetByID(docID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil { errorMessage := err.Error() responses = append(responses, &service.ParseDocumentResponse{ @@ -242,30 +242,30 @@ func (s *DocumentService) ParseDocuments(datasetID, userID string, docIDs []stri // StopParseDocuments stops parsing for the given documents in a dataset. // It sets Redis cancel signals for associated tasks and updates doc.run to CANCEL. // Returns a map with success_count and optionally errors. -func (s *DocumentService) StopParseDocuments(datasetID string, docIDs []string) (map[string]interface{}, error) { +func (s *DocumentService) StopParseDocuments(ctx context.Context, datasetID string, docIDs []string) (map[string]interface{}, error) { deduped := common.Deduplicate(docIDs) if len(deduped) == 0 { return nil, fmt.Errorf("no document IDs provided") } - docs, err := s.validateDocsInDataset(deduped, datasetID) + docs, err := s.validateDocsInDataset(ctx, deduped, datasetID) if err != nil { return nil, err } - var errors []string + var errs []string successCount := 0 for _, doc := range docs { - if cancelErr := s.CancelDocParse(doc); cancelErr != nil { - errors = append(errors, cancelErr.Error()) + if cancelErr := s.CancelDocParse(ctx, doc); cancelErr != nil { + errs = append(errs, cancelErr.Error()) continue } successCount++ } result := map[string]interface{}{"success_count": successCount} - if len(errors) > 0 { - result["errors"] = errors + if len(errs) > 0 { + result["errors"] = errs } return result, nil } @@ -273,8 +273,8 @@ func (s *DocumentService) StopParseDocuments(datasetID string, docIDs []string) // validateDocsInDataset deduplicates IDs, fetches the documents, and ensures // every document exists and belongs to the given dataset. Returns the resolved // documents. -func (s *DocumentService) validateDocsInDataset(docIDs []string, datasetID string) ([]*entity.Document, error) { - docs, err := s.documentDAO.GetByIDs(docIDs) +func (s *DocumentService) validateDocsInDataset(ctx context.Context, docIDs []string, datasetID string) ([]*entity.Document, error) { + docs, err := s.documentDAO.GetByIDs(ctx, dao.DB, docIDs) if err != nil { return nil, fmt.Errorf("failed to fetch documents: %w", err) } @@ -295,7 +295,7 @@ func (s *DocumentService) validateDocsInDataset(docIDs []string, datasetID strin // CancelDocParse stops the ingestion task for the document by calling // RequestStop (STOPPING), then marks the document run status as CANCEL. -func (s *DocumentService) CancelDocParse(doc *entity.Document) error { +func (s *DocumentService) CancelDocParse(ctx context.Context, doc *entity.Document) error { task, err := s.ingestionTaskDAO.GetByDocumentID(doc.ID) if err != nil { return fmt.Errorf("failed to get ingestion task for %s: %v", doc.ID, err) @@ -304,17 +304,17 @@ func (s *DocumentService) CancelDocParse(doc *entity.Document) error { return fmt.Errorf("no ingestion task found for document %s", doc.ID) } - if _, err := s.ingestionTaskSvc.RequestStop(task.ID); err != nil { + if _, err = s.ingestionTaskSvc.RequestStop(task.ID); err != nil { return fmt.Errorf("failed to stop ingestion task %s: %v", task.ID, err) } - if upErr := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{"run": string(entity.TaskStatusCancel)}); upErr != nil { + if upErr := s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{"run": string(entity.TaskStatusCancel)}); upErr != nil { return fmt.Errorf("failed to update document %s: %v", doc.ID, upErr) } return nil } -func (s *DocumentService) resetDocumentForReparse(doc *entity.Document, tenantID string, parserID *string, pipelineID *string) error { +func (s *DocumentService) resetDocumentForReparse(ctx context.Context, doc *entity.Document, tenantID string, parserID *string, pipelineID *string) error { progressMsg := "" run := string(entity.TaskStatusUnstart) updates := map[string]interface{}{ @@ -329,14 +329,14 @@ func (s *DocumentService) resetDocumentForReparse(doc *entity.Document, tenantID updates["pipeline_id"] = *pipelineID } - if err := s.documentDAO.UpdateByID(doc.ID, updates); err != nil { - return errors.New("Document not found!") + if err := s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, updates); err != nil { + return errors.New("document not found") } if doc.TokenNum > 0 { decremented, err := s.decrementDocumentAndKBCountersForReparse(doc) if err != nil { - return errors.New("Document not found!") + return errors.New("document not found") } if !decremented { return nil @@ -344,7 +344,7 @@ func (s *DocumentService) resetDocumentForReparse(doc *entity.Document, tenantID 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 { + if _, err = s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil { return err } } @@ -446,14 +446,14 @@ func (s *DocumentService) decrementDocumentAndKBCountersForReparse(doc *entity.D return decremented, err } -func (s *DocumentService) updateDocumentStatusOnly(doc *entity.Document, kb *entity.Knowledgebase, status int) error { +func (s *DocumentService) updateDocumentStatusOnly(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, status int) error { statusStr := strconv.Itoa(status) if doc.Status != nil && *doc.Status == statusStr { return nil } - if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{"status": statusStr}); err != nil { - return errors.New("Database error (Document update)!") + if err := s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{"status": statusStr}); err != nil { + return errors.New("database error (Document update)") } if s.docEngine == nil { diff --git a/internal/service/document/document_test.go b/internal/service/document/document_test.go index 8115eab4a4..574685ec1b 100644 --- a/internal/service/document/document_test.go +++ b/internal/service/document/document_test.go @@ -589,14 +589,15 @@ func TestDeleteDocumentFull_Basic(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) + ctx := t.Context() - err := svc.deleteDocumentFull("doc-1") + err := svc.deleteDocumentFull(ctx, "doc-1") if err != nil { t.Fatalf("deleteDocumentFull failed: %v", err) } // Verify document deleted - _, err = dao.NewDocumentDAO().GetByID("doc-1") + _, err = dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err == nil { t.Fatal("document should be deleted but it still exists") } @@ -628,8 +629,9 @@ func TestDeleteDocumentFull_NotFound(t *testing.T) { pushServiceDB(t, db) svc := testDocumentService(t) + ctx := t.Context() - err := svc.deleteDocumentFull("nonexistent") + err := svc.deleteDocumentFull(ctx, "nonexistent") if err == nil { t.Fatal("expected error for nonexistent document") } @@ -647,8 +649,9 @@ func TestDeleteDocumentFull_CleansUpFile2Document(t *testing.T) { insertTestFile2Document(t, "f2d-1", "file-1", "doc-1") svc := testDocumentService(t) + ctx := t.Context() - err := svc.deleteDocumentFull("doc-1") + err := svc.deleteDocumentFull(ctx, "doc-1") if err != nil { t.Fatalf("deleteDocumentFull failed: %v", err) } @@ -683,9 +686,10 @@ func TestDeleteDocumentFull_SharedFilePreserved(t *testing.T) { insertTestFile2Document(t, "f2d-2", "file-shared", "doc-2") svc := testDocumentService(t) + ctx := t.Context() // Delete doc-1; file-shared should survive because doc-2 still references it - err := svc.deleteDocumentFull("doc-1") + err := svc.deleteDocumentFull(ctx, "doc-1") if err != nil { t.Fatalf("deleteDocumentFull failed: %v", err) } @@ -761,9 +765,10 @@ func TestUploadLocalDocuments_MirrorsPythonCoreFields(t *testing.T) { t.Fatalf("insert existing doc: %v", err) } + ctx := t.Context() svc := testDocumentService(t) fh := makeTestFileHeader(t, "file", "deck.pptx", []byte("abc")) - got, errs := svc.UploadLocalDocuments(kb, "user-1", []*multipart.FileHeader{fh}, "nested/path", map[string]interface{}{ + got, errs := svc.UploadLocalDocuments(ctx, kb, "user-1", []*multipart.FileHeader{fh}, "nested/path", map[string]interface{}{ "table_column_mode": "assist", }) if len(errs) != 0 { @@ -819,8 +824,9 @@ func TestUploadEmptyDocument_CreatesVirtualDocumentAndFileLink(t *testing.T) { t.Fatalf("insert kb: %v", err) } + ctx := t.Context() svc := testDocumentService(t) - got, code, err := svc.UploadEmptyDocument(kb, "user-1", "draft.md") + got, code, err := svc.UploadEmptyDocument(ctx, kb, "user-1", "draft.md") if err != nil { t.Fatalf("UploadEmptyDocument error: %v", err) } @@ -900,8 +906,9 @@ func TestDeleteDocuments_DeleteAll(t *testing.T) { insertTestIngestionTask(t, "task-3", "user-1", "doc-3", "kb-1") svc := testDocumentService(t) + ctx := t.Context() - deleted, err := svc.DeleteDocuments(nil, true, "kb-1", "user-1") + deleted, err := svc.DeleteDocuments(ctx, nil, true, "kb-1", "user-1") if err != nil { t.Fatalf("DeleteDocuments failed: %v", err) } @@ -930,8 +937,8 @@ func TestDeleteDocuments_ByIDs(t *testing.T) { insertTestIngestionTask(t, "task-3", "user-1", "doc-3", "kb-1") svc := testDocumentService(t) - - deleted, err := svc.DeleteDocuments([]string{"doc-1", "doc-2"}, false, "kb-1", "user-1") + ctx := t.Context() + deleted, err := svc.DeleteDocuments(ctx, []string{"doc-1", "doc-2"}, false, "kb-1", "user-1") if err != nil { t.Fatalf("DeleteDocuments failed: %v", err) } @@ -940,7 +947,7 @@ func TestDeleteDocuments_ByIDs(t *testing.T) { } // doc-3 should still exist - _, err = dao.NewDocumentDAO().GetByID("doc-3") + _, err = dao.NewDocumentDAO().GetByID(ctx, db, "doc-3") if err != nil { t.Fatal("doc-3 should not have been deleted") } @@ -957,8 +964,9 @@ func TestDeleteDocuments_WrongDataset(t *testing.T) { insertTestDoc(t, "doc-1", "kb-2", 10, 5) // belongs to kb-2, not kb-1 svc := testDocumentService(t) + ctx := t.Context() - _, err := svc.DeleteDocuments([]string{"doc-1"}, false, "kb-1", "user-1") + _, err := svc.DeleteDocuments(ctx, []string{"doc-1"}, false, "kb-1", "user-1") if err == nil { t.Fatal("expected error for doc not belonging to dataset") } @@ -971,9 +979,10 @@ func TestDeleteDocuments_NotAccessible(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) svc := testDocumentService(t) + ctx := t.Context() // user-1 has no user_tenant entry → accessible returns false - _, err := svc.DeleteDocuments([]string{"doc-1"}, false, "kb-1", "user-1") + _, err := svc.DeleteDocuments(ctx, []string{"doc-1"}, false, "kb-1", "user-1") if err == nil { t.Fatal("expected error for inaccessible dataset") } @@ -988,7 +997,8 @@ func TestDeleteDocuments_EmptyIDs(t *testing.T) { svc := testDocumentService(t) // Empty ids, no deleteAll → returns 0, no error - deleted, err := svc.DeleteDocuments([]string{}, false, "kb-1", "user-1") + ctx := t.Context() + deleted, err := svc.DeleteDocuments(ctx, []string{}, false, "kb-1", "user-1") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -1007,8 +1017,9 @@ func TestDeleteDocuments_Deduplicate(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) + ctx := t.Context() - deleted, err := svc.DeleteDocuments([]string{"doc-1", "doc-1", "doc-1"}, false, "kb-1", "user-1") + deleted, err := svc.DeleteDocuments(ctx, []string{"doc-1", "doc-1", "doc-1"}, false, "kb-1", "user-1") if err != nil { t.Fatalf("DeleteDocuments failed: %v", err) } @@ -1064,17 +1075,18 @@ func TestStartParseDocuments_EnqueuesIngestionTask(t *testing.T) { publisher := &recordingTaskPublisher{} svc := testDocumentService(t) svc.ingestionTaskSvc.SetTaskPublisher(publisher) + ctx := t.Context() kb, err := svc.kbDAO.GetByID("kb-1") if err != nil { t.Fatalf("load kb: %v", err) } - doc, err := svc.documentDAO.GetByID("doc-1") + doc, err := svc.documentDAO.GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("load doc: %v", err) } - if err := svc.StartParseDocuments(doc, kb, "user-1", StartParseOptions{}); err != nil { + if err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{}); err != nil { t.Fatalf("StartParseDocuments: %v", err) } @@ -1128,8 +1140,8 @@ func TestStopParseDocuments_Success(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) - - result, err := svc.StopParseDocuments("kb-1", []string{"doc-1"}) + ctx := t.Context() + result, err := svc.StopParseDocuments(ctx, "kb-1", []string{"doc-1"}) if err != nil { t.Fatalf("StopParseDocuments failed: %v", err) } @@ -1143,7 +1155,7 @@ func TestStopParseDocuments_Success(t *testing.T) { } // Verify document run status updated to CANCEL - doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + doc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if doc == nil || doc.Run == nil { t.Fatal("doc not found or run is nil") } @@ -1162,8 +1174,8 @@ func TestStopParseDocuments_CancelStatus(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) - - result, err := svc.StopParseDocuments("kb-1", []string{"doc-1"}) + ctx := t.Context() + result, err := svc.StopParseDocuments(ctx, "kb-1", []string{"doc-1"}) if err != nil { t.Fatalf("StopParseDocuments failed: %v", err) } @@ -1183,8 +1195,8 @@ func TestStopParseDocuments_NotRunningOrCancel(t *testing.T) { insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusUnstart), 10, 5) svc := testDocumentService(t) - - result, err := svc.StopParseDocuments("kb-1", []string{"doc-1"}) + ctx := t.Context() + result, err := svc.StopParseDocuments(ctx, "kb-1", []string{"doc-1"}) if err != nil { t.Fatalf("StopParseDocuments failed: %v", err) } @@ -1209,8 +1221,8 @@ func TestStopParseDocuments_UnfinishedTask(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) - - result, err := svc.StopParseDocuments("kb-1", []string{"doc-1"}) + ctx := t.Context() + result, err := svc.StopParseDocuments(ctx, "kb-1", []string{"doc-1"}) if err != nil { t.Fatalf("StopParseDocuments failed: %v", err) } @@ -1230,8 +1242,8 @@ func TestStopParseDocuments_WrongDataset(t *testing.T) { insertTestDocWithRun(t, "doc-1", "kb-2", string(entity.TaskStatusRunning), 10, 5) svc := testDocumentService(t) - - _, err := svc.StopParseDocuments("kb-1", []string{"doc-1"}) + ctx := t.Context() + _, err := svc.StopParseDocuments(ctx, "kb-1", []string{"doc-1"}) if err == nil { t.Fatal("expected error for doc not belonging to dataset") } @@ -1244,8 +1256,8 @@ func TestStopParseDocuments_NotFound(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) svc := testDocumentService(t) - - _, err := svc.StopParseDocuments("kb-1", []string{"nonexistent"}) + ctx := t.Context() + _, err := svc.StopParseDocuments(ctx, "kb-1", []string{"nonexistent"}) if err == nil { t.Fatal("expected error for nonexistent document IDs") } @@ -1258,8 +1270,8 @@ func TestStopParseDocuments_EmptyIDs(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) svc := testDocumentService(t) - - _, err := svc.StopParseDocuments("kb-1", []string{}) + ctx := t.Context() + _, err := svc.StopParseDocuments(ctx, "kb-1", []string{}) if err == nil { t.Fatal("expected error for empty doc IDs") } @@ -1274,8 +1286,8 @@ func TestStopParseDocuments_Deduplicate(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) - - result, err := svc.StopParseDocuments("kb-1", []string{"doc-1", "doc-1", "doc-1"}) + ctx := t.Context() + result, err := svc.StopParseDocuments(ctx, "kb-1", []string{"doc-1", "doc-1", "doc-1"}) if err != nil { t.Fatalf("StopParseDocuments failed: %v", err) } @@ -1296,14 +1308,15 @@ func TestDeleteDocument_DeligatesToFullCleanup(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) + ctx := t.Context() // Public DeleteDocument should delegate to deleteDocumentFull - err := svc.DeleteDocument("doc-1") + err := svc.DeleteDocument(ctx, "doc-1") if err != nil { t.Fatalf("DeleteDocument failed: %v", err) } - _, err = dao.NewDocumentDAO().GetByID("doc-1") + _, err = dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err == nil { t.Fatal("document should be deleted") } @@ -1319,8 +1332,9 @@ func TestResolveDocAndKB_Success(t *testing.T) { insertTestDoc(t, "doc-1", "kb-1", 10, 5) svc := testDocumentService(t) + ctx := t.Context() - doc, kb, err := svc.resolveDocAndKB("doc-1") + doc, kb, err := svc.resolveDocAndKB(ctx, "doc-1") if err != nil { t.Fatalf("resolveDocAndKB: %v", err) } @@ -1340,8 +1354,9 @@ func TestResolveDocAndKB_DocNotFound(t *testing.T) { pushServiceDB(t, db) svc := testDocumentService(t) + ctx := t.Context() - _, _, err := svc.resolveDocAndKB("nonexistent") + _, _, err := svc.resolveDocAndKB(ctx, "nonexistent") if err == nil { t.Fatal("expected error for nonexistent doc") } @@ -1361,8 +1376,9 @@ func TestResolveDocAndKB_KBNotFound(t *testing.T) { } svc := testDocumentService(t) + ctx := t.Context() - _, _, err := svc.resolveDocAndKB("orphan-doc") + _, _, err := svc.resolveDocAndKB(ctx, "orphan-doc") if err == nil { t.Fatal("expected error for nonexistent KB") } @@ -1375,16 +1391,17 @@ func TestDeleteDocRecordWithCounters_Success(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 3, 100, 50) insertTestDoc(t, "doc-1", "kb-1", 30, 10) - doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") svc := testDocumentService(t) - err := svc.deleteDocRecordWithCounters(doc, "kb-1") + err := svc.deleteDocRecordWithCounters(ctx, doc, "kb-1") if err != nil { t.Fatalf("deleteDocRecordWithCounters: %v", err) } // Doc gone - _, err = dao.NewDocumentDAO().GetByID("doc-1") + _, err = dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err == nil { t.Fatal("document should be deleted") } @@ -1409,16 +1426,17 @@ func TestDeleteDocRecordWithCounters_DocAlreadyDeleted(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDoc(t, "doc-1", "kb-1", 10, 5) - doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") svc := testDocumentService(t) // First delete: row removed, counters decremented - if err := svc.deleteDocRecordWithCounters(doc, "kb-1"); err != nil { + if err := svc.deleteDocRecordWithCounters(ctx, doc, "kb-1"); err != nil { t.Fatalf("first delete: %v", err) } // Second delete: RowsAffected==0 → counters NOT decremented again - if err := svc.deleteDocRecordWithCounters(doc, "kb-1"); err != nil { + if err := svc.deleteDocRecordWithCounters(ctx, doc, "kb-1"); err != nil { t.Fatalf("second delete should not error: %v", err) } @@ -1441,14 +1459,15 @@ func TestDeleteDocRecordWithCounters_KBUpdateFailureRollsBackDocumentDelete(t *t insertTestDoc(t, "doc-1", "missing-kb", 10, 5) - doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") svc := testDocumentService(t) - if err := svc.deleteDocRecordWithCounters(doc, "missing-kb"); err == nil { + if err := svc.deleteDocRecordWithCounters(ctx, doc, "missing-kb"); err == nil { t.Fatal("expected missing KB counter update to return an error") } - if _, err := dao.NewDocumentDAO().GetByID("doc-1"); err != nil { + if _, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1"); err != nil { t.Fatalf("expected document delete to roll back, got: %v", err) } } @@ -1532,16 +1551,18 @@ func TestArtifactHelpers(t *testing.T) { func TestGetDocumentArtifact_InvalidFilename(t *testing.T) { svc := testDocumentService(t) - _, err := svc.GetDocumentArtifact("../test.txt", "user-1") - if err != ErrArtifactInvalidFilename { + ctx := t.Context() + _, err := svc.GetDocumentArtifact(ctx, "../test.txt", "user-1") + if !errors.Is(err, ErrArtifactInvalidFilename) { t.Errorf("expected ErrArtifactInvalidFilename, got %v", err) } } func TestGetDocumentArtifact_InvalidFileType(t *testing.T) { svc := testDocumentService(t) - _, err := svc.GetDocumentArtifact("test.exe", "user-1") - if err != ErrArtifactInvalidFileType { + ctx := t.Context() + _, err := svc.GetDocumentArtifact(ctx, "test.exe", "user-1") + if !errors.Is(err, ErrArtifactInvalidFileType) { t.Errorf("expected ErrArtifactInvalidFileType, got %v", err) } } @@ -1551,7 +1572,8 @@ func TestGetDocumentPreview_DocumentNotFound(t *testing.T) { pushServiceDB(t, db) svc := testDocumentService(t) - _, err := svc.GetDocumentPreview("nonexistent") + ctx := t.Context() + _, err := svc.GetDocumentPreview(ctx, "nonexistent") if err == nil { t.Error("expected error for nonexistent document") } @@ -1559,7 +1581,8 @@ func TestGetDocumentPreview_DocumentNotFound(t *testing.T) { func TestDownloadDocument_MissingDocID(t *testing.T) { svc := testDocumentService(t) - _, err := svc.DownloadDocument("ds-1", "") + ctx := t.Context() + _, err := svc.DownloadDocument(ctx, "ds-1", "") if err == nil { t.Error("expected error for missing docID") } @@ -1572,7 +1595,8 @@ func TestDownloadDocument_WrongDataset(t *testing.T) { insertTestDoc(t, "doc-1", "kb-1", 5, 2) svc := testDocumentService(t) - _, err := svc.DownloadDocument("wrong-ds", "doc-1") + ctx := t.Context() + _, err := svc.DownloadDocument(ctx, "wrong-ds", "doc-1") if err == nil { t.Error("expected error for wrong dataset") } @@ -1585,7 +1609,8 @@ func TestUpdateDatasetDocumentRejectsNonOwner(t *testing.T) { insertTestDoc(t, "doc-1", "kb-1", 0, 0) svc := testDocumentService(t) - _, code, err := svc.UpdateDatasetDocument("tenant-2", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{}, map[string]bool{}) + ctx := t.Context() + _, code, err := svc.UpdateDatasetDocument(ctx, "tenant-2", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{}, map[string]bool{}) if err == nil { t.Fatal("expected ownership error") } @@ -1605,7 +1630,8 @@ func TestUpdateDatasetDocumentRejectsCounterMutation(t *testing.T) { chunkCount := int64(6) svc := testDocumentService(t) - _, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + _, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ ChunkCount: &chunkCount, }, map[string]bool{"chunk_count": true}) if err == nil { @@ -1614,7 +1640,7 @@ func TestUpdateDatasetDocumentRejectsCounterMutation(t *testing.T) { if code != common.CodeDataError { t.Fatalf("code = %v, want %v", code, common.CodeDataError) } - if err.Error() != "Can't change `chunk_count`." { + if err.Error() != "can't change `chunk_count`" { t.Fatalf("err = %q", err.Error()) } } @@ -1627,7 +1653,8 @@ func TestUpdateDatasetDocumentAllowsZeroCounterLikePythonTruthyCheck(t *testing. chunkCount := int64(0) svc := testDocumentService(t) - _, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + _, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ ChunkCount: &chunkCount, }, map[string]bool{"chunk_count": true}) if err != nil { @@ -1647,7 +1674,8 @@ func TestUpdateDatasetDocumentRejectsUnsupportedParserIDForVisualDoc(t *testing. parserID := "naive" parseType := 1 svc := testDocumentService(t) - _, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + _, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ ParserID: &parserID, ParseType: &parseType, }, map[string]bool{"parser_id": true, "parse_type": true}) @@ -1657,7 +1685,7 @@ func TestUpdateDatasetDocumentRejectsUnsupportedParserIDForVisualDoc(t *testing. if code != common.CodeDataError { t.Fatalf("code = %v, want %v", code, common.CodeDataError) } - if err.Error() != "Not supported yet!" { + if err.Error() != "not supported yet" { t.Fatalf("err = %q", err.Error()) } } @@ -1672,7 +1700,8 @@ func TestUpdateDatasetDocumentRenameUpdatesDocumentAndFile(t *testing.T) { newName := "new.pdf" svc := testDocumentService(t) - resp, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + resp, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ Name: &newName, }, map[string]bool{"name": true}) if err != nil { @@ -1682,7 +1711,7 @@ func TestUpdateDatasetDocumentRenameUpdatesDocumentAndFile(t *testing.T) { t.Fatalf("response name = %#v, want %q", resp, newName) } - doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + doc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if doc.Name == nil || *doc.Name != newName { t.Fatalf("document name = %v, want %q", doc.Name, newName) } @@ -1701,7 +1730,8 @@ func TestUpdateDatasetDocumentParserIDResetsForReparse(t *testing.T) { parseType := 1 chunkMethod := "manual" svc := testDocumentService(t) - resp, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + resp, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ ParserID: &chunkMethod, ParseType: &parseType, }, map[string]bool{"parser_id": true, "parse_type": true}) @@ -1712,7 +1742,7 @@ func TestUpdateDatasetDocumentParserIDResetsForReparse(t *testing.T) { t.Fatalf("response = %+v, want method=%s run=UNSTART counts=0", resp, chunkMethod) } - doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + doc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if doc.ParserID != chunkMethod { t.Fatalf("parser_id = %q, want %q", doc.ParserID, chunkMethod) } @@ -1731,21 +1761,22 @@ func TestResetDocumentForReparseSkipsSecondCounterDecrement(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertNamedTestDoc(t, "doc-1", "kb-1", "doc.txt", 10, 5) - staleDoc, err := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + staleDoc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } svc := testDocumentService(t) parserID := "manual" - if err := svc.resetDocumentForReparse(staleDoc, "tenant-1", &parserID, nil); err != nil { + if err = svc.resetDocumentForReparse(ctx, staleDoc, "tenant-1", &parserID, nil); err != nil { t.Fatalf("first resetDocumentForReparse failed: %v", err) } - if err := svc.resetDocumentForReparse(staleDoc, "tenant-1", &parserID, nil); err != nil { + if err = svc.resetDocumentForReparse(ctx, staleDoc, "tenant-1", &parserID, nil); err != nil { t.Fatalf("second resetDocumentForReparse failed: %v", err) } - doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + doc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if doc.TokenNum != 0 || doc.ChunkNum != 0 { t.Fatalf("doc counters = token:%d chunk:%d, want zero", doc.TokenNum, doc.ChunkNum) } @@ -1762,7 +1793,8 @@ func TestClearDocumentParseResultsClearsCountersTasksAndChunks(t *testing.T) { insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED) - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -1771,11 +1803,11 @@ func TestClearDocumentParseResultsClearsCountersTasksAndChunks(t *testing.T) { svc := testDocumentService(t) svc.docEngine = engine - if err := svc.clearDocumentParseResults(doc, "tenant-1"); err != nil { + if err = svc.clearDocumentParseResults(doc, "tenant-1"); err != nil { t.Fatalf("clearDocumentParseResults failed: %v", err) } - updatedDoc, _ := dao.NewDocumentDAO().GetByID("doc-1") + updatedDoc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if updatedDoc.TokenNum != 0 || updatedDoc.ChunkNum != 0 { t.Fatalf("doc counters = token:%d chunk:%d, want zero", updatedDoc.TokenNum, updatedDoc.ChunkNum) } @@ -1802,7 +1834,8 @@ func TestClearDocumentParseResultsIsIdempotentForStaleDocSnapshot(t *testing.T) insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) - staleDoc, err := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + staleDoc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -1815,7 +1848,7 @@ func TestClearDocumentParseResultsIsIdempotentForStaleDocSnapshot(t *testing.T) t.Fatalf("second clearDocumentParseResults failed: %v", err) } - doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + doc, _ := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if doc.TokenNum != 0 || doc.ChunkNum != 0 { t.Fatalf("doc counters = token:%d chunk:%d, want zero", doc.TokenNum, doc.ChunkNum) } @@ -1837,8 +1870,8 @@ func TestClearDocumentParseResults_RejectsNonTerminalIngestionTask(t *testing.T) insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusRunning), 10, 5) insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", status) - - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -1867,8 +1900,8 @@ func TestClearDocumentParseResults_DeletesTerminalIngestionTask(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", status) - - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -1929,8 +1962,9 @@ func TestIngest_RerunWithDelete_RejectsBatchWithRunningTask(t *testing.T) { insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED) insertTestIngestionTaskWithStatus(t, "task-2", "user-1", "doc-2", "kb-1", common.RUNNING) + ctx := t.Context() svc := testDocumentService(t) - _, err := svc.Ingest("user-1", &IngestDocumentRequest{ + _, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{ DocIDs: []string{"doc-1", "doc-2"}, Run: string(entity.TaskStatusRunning), Delete: true, @@ -1955,8 +1989,8 @@ func TestUpdateDatasetDocumentPropagatesMetadataDeleteFailure(t *testing.T) { svc := testDocumentService(t) svc.docEngine = engine svc.metadataSvc = service.NewMetadataServiceForTest(nil, nil) - - _, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + _, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ MetaFields: map[string]any{"new": "value"}, }, map[string]bool{"meta_fields": true}) if err == nil { @@ -1989,7 +2023,8 @@ func TestSetDocumentMetadataMergesMetadataRow(t *testing.T) { svc.docEngine = engine svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine) - if err := svc.SetDocumentMetadata("doc-1", map[string]interface{}{"category": "tech", "year": 2026}); err != nil { + ctx := t.Context() + if err := svc.SetDocumentMetadata(ctx, "doc-1", map[string]interface{}{"category": "tech", "year": 2026}); err != nil { t.Fatalf("SetDocumentMetadata failed: %v", err) } if got := engine.records["doc-1"]["author"]; got != "alice" { @@ -2060,8 +2095,8 @@ func TestBatchUpdateDocumentMetadatasMatchesPythonSemantics(t *testing.T) { svc := testDocumentService(t) svc.docEngine = engine svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine) - - resp, code, err := svc.BatchUpdateDocumentMetadatas("kb-1", &DocumentMetadataSelector{ + ctx := t.Context() + resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{ DocumentIDs: []string{"doc-1", "doc-2", "doc-3"}, }, []DocumentMetadataUpdate{ {Key: "tags", Value: "new", Match: "old"}, @@ -2121,8 +2156,8 @@ func TestBatchUpdateDocumentMetadatasDoesNotReplaceWhenCurrentSearchIsStale(t *t svc := testDocumentService(t) svc.docEngine = engine svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine) - - resp, code, err := svc.BatchUpdateDocumentMetadatas("kb-1", &DocumentMetadataSelector{ + ctx := t.Context() + resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{ DocumentIDs: []string{"doc-1"}, }, []DocumentMetadataUpdate{ {Key: "category", Value: "paper"}, @@ -2158,8 +2193,8 @@ func TestBatchUpdateDocumentMetadatasDeletesEmptyMetadataAndNoOps(t *testing.T) svc := testDocumentService(t) svc.docEngine = engine svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine) - - resp, code, err := svc.BatchUpdateDocumentMetadatas("kb-1", &DocumentMetadataSelector{ + ctx := t.Context() + resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{ DocumentIDs: []string{"doc-1", "doc-2"}, }, nil, []DocumentMetadataDelete{{Key: "status", Value: "draft"}}) if err != nil || code != common.CodeSuccess { @@ -2187,8 +2222,8 @@ func TestBatchUpdateDocumentMetadatasNormalizesNumberValues(t *testing.T) { svc := testDocumentService(t) svc.docEngine = engine svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine) - - resp, code, err := svc.BatchUpdateDocumentMetadatas("kb-1", &DocumentMetadataSelector{ + ctx := t.Context() + resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{ DocumentIDs: []string{"doc-1"}, }, []DocumentMetadataUpdate{ {Key: "score", Value: "42", ValueType: "number"}, @@ -2217,7 +2252,8 @@ func TestBatchUpdateDocumentMetadatasNormalizesNumberValues(t *testing.T) { func TestBatchUpdateDocumentMetadatasRejectsMissingValue(t *testing.T) { svc := testDocumentService(t) - resp, code, err := svc.BatchUpdateDocumentMetadatas("kb-1", &DocumentMetadataSelector{}, []DocumentMetadataUpdate{ + ctx := t.Context() + resp, code, err := svc.BatchUpdateDocumentMetadatas(ctx, "kb-1", &DocumentMetadataSelector{}, []DocumentMetadataUpdate{ {Key: "status"}, }, nil) if err == nil { @@ -2276,7 +2312,8 @@ func TestUpdateDatasetDocumentPipelineIDTakesPrecedenceOverParserID(t *testing.T chunkMethod := "manual" parseType := 2 svc := testDocumentService(t) - resp, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + resp, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ PipelineID: &pipelineID, ParserID: &chunkMethod, ParseType: &parseType, @@ -2311,7 +2348,8 @@ func TestUpdateDatasetDocumentParseTypeBuiltin(t *testing.T) { parserID := "manual" pipelineIDEmpty := "" svc := testDocumentService(t) - resp, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + resp, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ ParseType: &parseType, ParserID: &parserID, PipelineID: &pipelineIDEmpty, @@ -2328,7 +2366,7 @@ func TestUpdateDatasetDocumentParseTypeBuiltin(t *testing.T) { // Assert the persisted pipeline_id was cleared, not merely the response. var updated entity.Document - if err := db.First(&updated, "id = ?", "doc-1").Error; err != nil { + if err = db.First(&updated, "id = ?", "doc-1").Error; err != nil { t.Fatalf("read back doc: %v", err) } if updated.PipelineID != nil && *updated.PipelineID != "" { @@ -2345,7 +2383,8 @@ func TestUpdateDatasetDocumentParseTypePipeline(t *testing.T) { parseType := 2 pipelineID := "1234567890abcdef1234567890abcdef" svc := testDocumentService(t) - resp, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + resp, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ ParseType: &parseType, PipelineID: &pipelineID, }, map[string]bool{"parser_id": false, "pipeline_id": true, "parse_type": true}) @@ -2381,7 +2420,8 @@ func TestUpdateDatasetDocumentParseTypePipelineIgnoresDirtyParserID(t *testing.T parserID := "manual" pipelineID := "1234567890abcdef1234567890abcdef" svc := testDocumentService(t) - resp, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + resp, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ ParseType: &parseType, ParserID: &parserID, PipelineID: &pipelineID, @@ -2418,7 +2458,8 @@ func TestUpdateDatasetDocumentRejectsInvalidPages(t *testing.T) { parseType := 1 parserID := "manual" svc := testDocumentService(t) - _, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + _, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ ParseType: &parseType, ParserID: &parserID, ParserConfig: map[string]any{ @@ -2445,7 +2486,8 @@ func TestUpdateDatasetDocumentEnabledUpdatesStatus(t *testing.T) { enabled := 0 svc := testDocumentService(t) - resp, code, err := svc.UpdateDatasetDocument("tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ + ctx := t.Context() + resp, code, err := svc.UpdateDatasetDocument(ctx, "tenant-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{ Enabled: &enabled, }, map[string]bool{"enabled": true}) if err != nil { @@ -2603,14 +2645,15 @@ func TestGetDocumentArtifact_AuthGate(t *testing.T) { } svc := testDocumentService(t) + ctx := t.Context() // Case 1: empty user -> not allowed. - if _, err := svc.GetDocumentArtifact("result.png", ""); !errors.Is(err, ErrArtifactNotFound) { + if _, err := svc.GetDocumentArtifact(ctx, "result.png", ""); !errors.Is(err, ErrArtifactNotFound) { t.Errorf("empty user: want ErrArtifactNotFound, got %v", err) } // Case 2: another user without any session reference -> not allowed. - if _, err := svc.GetDocumentArtifact("result.png", "user-2"); !errors.Is(err, ErrArtifactNotFound) { + if _, err := svc.GetDocumentArtifact(ctx, "result.png", "user-2"); !errors.Is(err, ErrArtifactNotFound) { t.Errorf("unrelated user: want ErrArtifactNotFound, got %v", err) } @@ -2632,7 +2675,7 @@ func TestGetDocumentArtifact_AuthGate(t *testing.T) { }).Error; err != nil { t.Fatalf("seed conv 2: %v", err) } - if _, err := svc.GetDocumentArtifact("result.png", "user-2"); !errors.Is(err, ErrArtifactNotFound) { + if _, err := svc.GetDocumentArtifact(ctx, "result.png", "user-2"); !errors.Is(err, ErrArtifactNotFound) { t.Errorf("user-2 with unrelated session: want ErrArtifactNotFound, got %v", err) } } @@ -2703,7 +2746,8 @@ func TestGetThumbnails_AlignsWithPythonFormatting(t *testing.T) { } svc := testDocumentService(t) - got, err := svc.GetThumbnails("user-1", []string{"doc-file", "doc-base64", "doc-other", "missing-doc"}) + ctx := t.Context() + got, err := svc.GetThumbnails(ctx, "user-1", []string{"doc-file", "doc-base64", "doc-other", "missing-doc"}) if err != nil { t.Fatalf("GetThumbnails failed: %v", err) } @@ -2732,7 +2776,8 @@ func TestStartParseDocuments_FailsBeforeClearing(t *testing.T) { insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED) - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + ctx := t.Context() + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("get doc: %v", err) } @@ -2747,7 +2792,7 @@ func TestStartParseDocuments_FailsBeforeClearing(t *testing.T) { } svc := testDocumentService(t) - err = svc.StartParseDocuments(doc, kb, "user-1", StartParseOptions{RerunWithDelete: true}) + err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{RerunWithDelete: true}) if err == nil { t.Fatal("expected error from GetDocumentStorageAddress, got nil") } @@ -2772,7 +2817,8 @@ func TestIngest_CancelDoesNotDeleteIngestionTask(t *testing.T) { insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.RUNNING) svc := testDocumentService(t) - _, err := svc.Ingest("user-1", &IngestDocumentRequest{ + ctx := t.Context() + _, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{ DocIDs: []string{"doc-1"}, Run: string(entity.TaskStatusCancel), Delete: true, @@ -2794,10 +2840,11 @@ func TestUpdateRunProgressMirrorsFields(t *testing.T) { insertTestDoc(t, "doc-1", "kb-1", 0, 0) svc := testDocumentService(t) - if err := svc.UpdateRunProgress("doc-1", 0.5, "1", "halfway"); err != nil { + ctx := t.Context() + if err := svc.UpdateRunProgress(ctx, "doc-1", 0.5, "1", "halfway"); err != nil { t.Fatalf("UpdateRunProgress failed: %v", err) } - doc, err := dao.NewDocumentDAO().GetByID("doc-1") + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err != nil { t.Fatalf("load document: %v", err) } @@ -2815,6 +2862,7 @@ func TestUpdateRunProgressMirrorsFields(t *testing.T) { func TestFileDeleteRemovesLinkedDocument(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) + ctx := t.Context() insertTestKB(t, "kb-1", "tenant-1", 1, 30, 10) insertTestDoc(t, "doc-1", "kb-1", 30, 10) @@ -2843,12 +2891,12 @@ func TestFileDeleteRemovesLinkedDocument(t *testing.T) { docSvc, ) - success, msg := fileSvc.DeleteFiles(context.Background(), "tenant-1", []string{"file-1"}) + success, msg := fileSvc.DeleteFiles(ctx, "tenant-1", []string{"file-1"}) if !success { t.Fatalf("DeleteFiles failed: %s", msg) } - _, err := dao.NewDocumentDAO().GetByID("doc-1") + _, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1") if err == nil { t.Fatal("document should have been deleted but still exists") } diff --git a/internal/service/document/document_upload.go b/internal/service/document/document_upload.go index a374868e6c..cc7ecb51dd 100644 --- a/internal/service/document/document_upload.go +++ b/internal/service/document/document_upload.go @@ -1,11 +1,13 @@ package document import ( + "context" "fmt" "io" "mime/multipart" "net/http" "path/filepath" + "ragflow/internal/dao" "strings" "ragflow/internal/common" @@ -25,7 +27,7 @@ import ( // // Gaps vs Python (documented, not yet ported): thumbnail generation and // read_potential_broken_pdf repair. -func (s *DocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) { +func (s *DocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) { storageImpl := storage.GetStorageFactory().GetStorage() if storageImpl == nil { return nil, []string{"storage not initialized"} @@ -52,7 +54,7 @@ func (s *DocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantI // Don't silently disable dedupe protection: a transient lookup failure means // the existing-name set is unknown, so fail rather than risk duplicates. - names, err := s.documentDAO.ListNamesByKbID(kb.ID) + names, err := s.documentDAO.ListNamesByKbID(ctx, dao.DB, kb.ID) if err != nil { return nil, []string{err.Error()} } @@ -65,7 +67,8 @@ func (s *DocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantI var errMsgs []string for _, fh := range files { - blob, err := readFileHeaderBytes(fh) + var blob []byte + blob, err = readFileHeaderBytes(fh) if err != nil { errMsgs = append(errMsgs, fh.Filename+": "+err.Error()) continue @@ -86,22 +89,22 @@ func (s *DocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantI for storageImpl.ObjExist(kb.ID, location) { location += "_" } - if err := storageImpl.Put(kb.ID, location, blob); err != nil { + if err = storageImpl.Put(kb.ID, location, blob); err != nil { errMsgs = append(errMsgs, fh.Filename+": "+err.Error()) continue } doc := s.newDatasetDocument(kb, tenantID, filename, location, string(filetype), merged, "local", int64(len(blob)), blob) - if err := s.InsertDocument(doc); err != nil { + 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) errMsgs = append(errMsgs, fh.Filename+": "+err.Error()) continue } - if err := s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil { + if err = s.addFileFromKB(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(doc, kb.ID, err) + err = s.rollbackAddFileFromKBError(ctx, doc, kb.ID, err) _ = storageImpl.Remove(kb.ID, location) errMsgs = append(errMsgs, fh.Filename+": "+err.Error()) continue @@ -115,16 +118,16 @@ func (s *DocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantI } // UploadEmptyDocument inserts a zero-byte "virtual" document into the dataset. -func (s *DocumentService) UploadEmptyDocument(kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) { +func (s *DocumentService) UploadEmptyDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) { // A transient lookup failure means the existing-name set is unknown; fail // rather than write blind and risk a duplicate. - names, err := s.documentDAO.ListNamesByKbID(kb.ID) + names, err := s.documentDAO.ListNamesByKbID(ctx, dao.DB, kb.ID) if err != nil { return nil, common.CodeServerError, err } for _, n := range names { if n == name { - return nil, common.CodeDataError, fmt.Errorf("Duplicated document name in the same dataset.") + return nil, common.CodeDataError, fmt.Errorf("duplicated document name in the same dataset") } } @@ -134,11 +137,11 @@ func (s *DocumentService) UploadEmptyDocument(kb *entity.Knowledgebase, tenantID } doc := s.newDatasetDocument(kb, tenantID, name, "", "virtual", kb.ParserConfig, "local", 0, nil) - if err := s.InsertDocument(doc); err != nil { + if err = s.InsertDocument(doc); err != nil { return nil, common.CodeServerError, err } - if err := s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil { - return nil, common.CodeServerError, s.rollbackAddFileFromKBError(doc, kb.ID, err) + if err = s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil { + return nil, common.CodeServerError, s.rollbackAddFileFromKBError(ctx, doc, kb.ID, err) } return docToRawMap(doc), common.CodeSuccess, nil } @@ -226,7 +229,7 @@ func (s *DocumentService) addFileFromKB(doc *entity.Document, kbFolderID, tenant return nil } -func (s *DocumentService) UploadWebDocument(kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) { +func (s *DocumentService) UploadWebDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) { storageImpl := storage.GetStorageFactory().GetStorage() if storageImpl == nil { return nil, common.CodeServerError, fmt.Errorf("storage not initialized") @@ -237,7 +240,7 @@ func (s *DocumentService) UploadWebDocument(kb *entity.Knowledgebase, tenantID, return nil, common.CodeServerError, err } - names, err := s.documentDAO.ListNamesByKbID(kb.ID) + names, err := s.documentDAO.ListNamesByKbID(ctx, dao.DB, kb.ID) if err != nil { return nil, common.CodeServerError, err } @@ -246,7 +249,7 @@ func (s *DocumentService) UploadWebDocument(kb *entity.Knowledgebase, tenantID, taken[n] = true } - blob, headers, _, err := utility.FetchRemoteFileSafely(url, maxUploadDocSize) + blob, headers, _, err := utility.FetchRemoteFileSafely(ctx, url, maxUploadDocSize) if err != nil { return nil, common.CodeDataError, err } @@ -260,24 +263,24 @@ func (s *DocumentService) UploadWebDocument(kb *entity.Knowledgebase, tenantID, filetype := utility.FilenameType(filename) if filetype == utility.FileTypeOTHER { - return nil, common.CodeDataError, fmt.Errorf("This type of file has not been supported yet!") + return nil, common.CodeDataError, fmt.Errorf("this type of file has not been supported yet") } location := filename for storageImpl.ObjExist(kb.ID, location) { location += "_" } - if err := storageImpl.Put(kb.ID, location, blob); err != nil { + if err = storageImpl.Put(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 { + if err = s.InsertDocument(doc); err != nil { _ = storageImpl.Remove(kb.ID, location) return nil, common.CodeServerError, err } - if err := s.addFileFromKB(doc, kbFolder.ID, kb.TenantID); err != nil { - err = s.rollbackAddFileFromKBError(doc, kb.ID, err) + if err = s.addFileFromKB(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 } diff --git a/internal/service/document/file2document.go b/internal/service/document/file2document.go index 0331d1ef36..ee17da3c9a 100644 --- a/internal/service/document/file2document.go +++ b/internal/service/document/file2document.go @@ -17,6 +17,7 @@ package document import ( + "context" "errors" "path/filepath" "ragflow/internal/service" @@ -81,7 +82,7 @@ type LinkToDatasetsRequest struct { // // On validation failure it returns a sentinel error (see ErrLink* above) so the // handler can map it to a Python-compatible response without leaking internals. -func (s *File2DocumentService) LinkToDatasets(userID string, req *LinkToDatasetsRequest, mode string) error { +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) if err != nil { @@ -149,7 +150,7 @@ func (s *File2DocumentService) LinkToDatasets(userID string, req *LinkToDatasets // ── 6. Run conversion in background (fire-and-forget) ──────────────────── kbIDs := req.KbIDs go func() { - if err := s.convertFiles(allFileIDs, kbIDs, userID, mode); err != nil { + if err = s.convertFiles(ctx, allFileIDs, kbIDs, userID, mode); err != nil { common.Warn("file2document.convertFiles failed", zap.Strings("file_ids", allFileIDs), zap.Strings("kb_ids", kbIDs), @@ -164,7 +165,7 @@ func (s *File2DocumentService) LinkToDatasets(userID string, req *LinkToDatasets // either remove existing documents/mappings (replace) or keep them and skip // already-linked KBs (add), then create a new document in each target KB and // a fresh mapping. -func (s *File2DocumentService) convertFiles(fileIDs, kbIDs []string, userID, mode string) error { +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) @@ -182,7 +183,7 @@ func (s *File2DocumentService) convertFiles(fileIDs, kbIDs []string, userID, mod if m.DocumentID == nil { continue } - if err := s.documentSvc.RemoveDocumentKeepFile(*m.DocumentID); err != nil { + if err = s.documentSvc.RemoveDocumentKeepFile(ctx, *m.DocumentID); err != nil { common.Warn("convertFiles: RemoveDocumentKeepFile failed", zap.String("docID", *m.DocumentID), zap.Error(err)) } @@ -196,11 +197,12 @@ func (s *File2DocumentService) convertFiles(fileIDs, kbIDs []string, userID, mod // "add" mode: collect KB IDs already linked to this file so we // skip them when creating new documents below. Existing links // are preserved (mirrors Python _convert_files add path). + var doc *entity.Document for _, m := range mappings { if m.DocumentID == nil { continue } - doc, err := s.documentDAO.GetByID(*m.DocumentID) + doc, err = s.documentDAO.GetByID(ctx, dao.DB, *m.DocumentID) if err != nil || doc == nil { continue } @@ -227,7 +229,7 @@ func (s *File2DocumentService) convertFiles(fileIDs, kbIDs []string, userID, mod // Mirror Python duplicate_name: generate a non-colliding document // name within the dataset so that linking does not silently create // duplicate-name documents in the same KB. - docName, err := s.documentSvc.UniqueDocumentName(kbID, file.Name) + docName, err := s.documentSvc.UniqueDocumentName(ctx, kbID, file.Name) if err != nil { common.Warn("convertFiles: UniqueDocumentName failed", zap.String("kbID", kbID), zap.String("fileID", fileID), zap.Error(err)) diff --git a/internal/service/file/file.go b/internal/service/file/file.go index 32c9cc8bc1..76d1332f1c 100644 --- a/internal/service/file/file.go +++ b/internal/service/file/file.go @@ -17,6 +17,7 @@ package file import ( + "context" "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/utility" @@ -31,7 +32,7 @@ var ( // DocRemover is the narrow interface FileService needs from the document domain. type DocRemover interface { - RemoveDocumentKeepFile(docID string) error + RemoveDocumentKeepFile(ctx context.Context, docID string) error } // CheckFilePermFunc is the function signature for file-team permission checks, diff --git a/internal/service/file/file_commit.go b/internal/service/file/file_commit.go index 14932b1d43..c62cfcaf75 100644 --- a/internal/service/file/file_commit.go +++ b/internal/service/file/file_commit.go @@ -216,7 +216,7 @@ func (s *FileCommitService) CreateCommit(folderID, authorID, message string, cha if err != nil { return fmt.Errorf("failed to marshal tree state: %w", err) } - if err := tx.Model(&entity.FileCommit{}).Where("id = ?", commitID).Update("tree_state", string(treeJSON)).Error; err != nil { + if err = tx.Model(&entity.FileCommit{}).Where("id = ?", commitID).Update("tree_state", string(treeJSON)).Error; err != nil { return fmt.Errorf("failed to update tree state: %w", err) } treeStr = string(treeJSON) diff --git a/internal/service/file/file_content.go b/internal/service/file/file_content.go index c3fb251143..2385dc115c 100644 --- a/internal/service/file/file_content.go +++ b/internal/service/file/file_content.go @@ -28,7 +28,7 @@ func (s *FileService) GetFileContent(uid, fileID string) (*entity.File, error) { // GetStorageAddress gets storage address for a file (fallback for when direct blob is empty) // Matches Python's File2DocumentService.get_storage_address function -func (s *FileService) GetStorageAddress(fileID string) (*StorageAddress, error) { +func (s *FileService) GetStorageAddress(ctx context.Context, fileID string) (*StorageAddress, error) { // Get file2document mapping f2d, err := s.file2DocumentDAO.GetByFileID(fileID) if err != nil || len(f2d) == 0 { @@ -61,7 +61,7 @@ func (s *FileService) GetStorageAddress(fileID string) (*StorageAddress, error) } documentDAO := dao.NewDocumentDAO() - doc, err := documentDAO.GetByID(*f2d[0].DocumentID) + doc, err := documentDAO.GetByID(ctx, dao.DB, *f2d[0].DocumentID) if err != nil || doc == nil { return nil, fmt.Errorf("document not found") } @@ -77,7 +77,7 @@ func (s *FileService) GetStorageAddress(fileID string) (*StorageAddress, error) } // DownloadAgentFile downloads an agent-generated file directly from MinIO without querying the database. -func (s *FileService) DownloadAgentFile(tenantID, location string) ([]byte, error) { +func (s *FileService) DownloadAgentFile(ctx context.Context, tenantID, location string) ([]byte, error) { storageImpl := storage.GetStorageFactory().GetStorage() if storageImpl == nil { return nil, fmt.Errorf("storage not initialized") diff --git a/internal/service/file/file_delete.go b/internal/service/file/file_delete.go index 5a609a351a..f39a799bbf 100644 --- a/internal/service/file/file_delete.go +++ b/internal/service/file/file_delete.go @@ -2,6 +2,7 @@ package file import ( "context" + "errors" "fmt" "ragflow/internal/common" "ragflow/internal/entity" @@ -78,20 +79,23 @@ func (s *FileService) deleteSingleFile(ctx context.Context, file *entity.File) e } docID := *inform.DocumentID if s.documentService != nil { - if err := s.documentService.RemoveDocumentKeepFile(docID); err != nil { + if err = s.documentService.RemoveDocumentKeepFile(ctx, docID); err != nil { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return fmt.Errorf("context cancelled while removing document %s: %w", docID, err) + } common.Logger.Error(fmt.Sprintf("Fail to remove document: %s, error: %v", docID, err)) } } } // 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(file.ID); err != nil { return fmt.Errorf("failed to delete file2document mapping: %w", err) } } // 3. Delete file record - if err := s.fileDAO.Delete(file.ID); err != nil { + if err = s.fileDAO.Delete(file.ID); err != nil { return err } diff --git a/internal/service/file/file_folder.go b/internal/service/file/file_folder.go index 4f2c0f638f..d1b0815d6f 100644 --- a/internal/service/file/file_folder.go +++ b/internal/service/file/file_folder.go @@ -1,6 +1,7 @@ package file import ( + "context" "fmt" "path/filepath" "ragflow/internal/common" @@ -254,9 +255,9 @@ func (s *FileService) GetAllParentFolders(userID, fileID string) ([]map[string]i } // GetDocCount gets document count for a tenant -func (s *FileService) GetDocCount(tenantID string) (int64, error) { +func (s *FileService) GetDocCount(ctx context.Context, tenantID string) (int64, error) { documentDAO := dao.NewDocumentDAO() - return documentDAO.CountByTenantID(tenantID) + return documentDAO.CountByTenantID(ctx, dao.DB, tenantID) } func (s *FileService) createFolderRecursive(parentFolder *entity.File, names []string, count int, tenantID string) (*entity.File, error) { @@ -335,7 +336,7 @@ func (s *FileService) CreateFolder(tenantID, name, parentID, fileType string) (m // - new_name only: rename in place (no storage operation) // - dest_file_id only: move to new folder (keep names) // - both: move and rename simultaneously -func (s *FileService) MoveFiles(uid string, srcFileIDs []string, destFileID string, newName string) (bool, string) { +func (s *FileService) MoveFiles(ctx context.Context, uid string, srcFileIDs []string, destFileID string, newName string) (bool, string) { // 1. Get all source files files, err := s.fileDAO.GetByIDs(srcFileIDs) if err != nil || len(files) == 0 { @@ -448,7 +449,7 @@ func (s *FileService) MoveFiles(uid string, srcFileIDs []string, destFileID stri if destFolder != nil { // Move to destination folder for _, file := range files { - if err := s.moveEntryRecursive(file, destFolder, newName); err != nil { + if err = s.moveEntryRecursive(ctx, file, destFolder, newName); err != nil { return false, err.Error() } } @@ -461,7 +462,7 @@ func (s *FileService) MoveFiles(uid string, srcFileIDs []string, destFileID stri return false, "Source files not found!" } file := filesMap[srcFileIDs[0]] - if err := s.fileDAO.UpdateByID(file.ID, map[string]interface{}{"name": newName}); err != nil { + if err = s.fileDAO.UpdateByID(file.ID, map[string]interface{}{"name": newName}); err != nil { return false, "Database error (File rename)!" } @@ -470,7 +471,7 @@ func (s *FileService) MoveFiles(uid string, srcFileIDs []string, destFileID stri if err == nil && len(informs) > 0 && informs[0].DocumentID != nil { docID := *informs[0].DocumentID documentDAO := dao.NewDocumentDAO() - if err := documentDAO.UpdateByID(docID, map[string]interface{}{"name": newName}); err != nil { + if err := documentDAO.UpdateByID(ctx, dao.DB, docID, map[string]interface{}{"name": newName}); err != nil { return false, "Database error (Document rename)!" } } @@ -480,7 +481,7 @@ func (s *FileService) MoveFiles(uid string, srcFileIDs []string, destFileID stri } // moveEntryRecursive recursively moves a file or folder entry -func (s *FileService) moveEntryRecursive(sourceFile *entity.File, destFolder *entity.File, overrideName string) error { +func (s *FileService) moveEntryRecursive(ctx context.Context, sourceFile *entity.File, destFolder *entity.File, overrideName string) error { effectiveName := overrideName if effectiveName == "" { effectiveName = sourceFile.Name @@ -511,7 +512,7 @@ func (s *FileService) moveEntryRecursive(sourceFile *entity.File, destFolder *en return err } for _, subFile := range subFiles { - if err := s.moveEntryRecursive(subFile, newFolder, ""); err != nil { + if err = s.moveEntryRecursive(ctx, subFile, newFolder, ""); err != nil { return err } } @@ -566,7 +567,7 @@ func (s *FileService) moveEntryRecursive(sourceFile *entity.File, destFolder *en if err == nil && len(informs) > 0 && informs[0].DocumentID != nil { docID := *informs[0].DocumentID documentDAO := dao.NewDocumentDAO() - if err := documentDAO.UpdateByID(docID, map[string]interface{}{"name": overrideName}); err != nil { + if err = documentDAO.UpdateByID(ctx, dao.DB, docID, map[string]interface{}{"name": overrideName}); err != nil { return fmt.Errorf("database error (Document rename): %w", err) } } diff --git a/internal/service/file/file_test.go b/internal/service/file/file_test.go index 0ea818b78b..29cc90a83f 100644 --- a/internal/service/file/file_test.go +++ b/internal/service/file/file_test.go @@ -236,7 +236,8 @@ func TestFileService_DownloadAgentFile_Success(t *testing.T) { tenantID := "tenant123" location := "file-abc.txt" - blob, err := svc.DownloadAgentFile(tenantID, location) + ctx := t.Context() + blob, err := svc.DownloadAgentFile(ctx, tenantID, location) if err != nil { t.Fatalf("expected no error, got %v", err) } @@ -271,7 +272,8 @@ func TestFileService_DownloadAgentFile_Error(t *testing.T) { tenantID := "tenant123" location := "file-abc.txt" - blob, err := svc.DownloadAgentFile(tenantID, location) + ctx := t.Context() + blob, err := svc.DownloadAgentFile(ctx, tenantID, location) if err == nil { t.Fatalf("expected error, got nil") } @@ -309,8 +311,9 @@ func TestFileService_UploadFromURL_PDFAddsExtensionAndStoresToDownloads(t *testi factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) + ctx := t.Context() svc := testFileService() - resp, err := svc.UploadFromURL("tenant123", server.URL+"/report") + resp, err := svc.UploadFromURL(ctx, "tenant123", server.URL+"/report") if err != nil { t.Fatalf("UploadFromURL failed: %v", err) } @@ -355,8 +358,9 @@ func TestFileService_UploadFromURL_HTMLNormalizesReadableContent(t *testing.T) { factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) + ctx := t.Context() svc := testFileService() - resp, err := svc.UploadFromURL("tenant123", server.URL+"/page") + resp, err := svc.UploadFromURL(ctx, "tenant123", server.URL+"/page") if err != nil { t.Fatalf("UploadFromURL failed: %v", err) } diff --git a/internal/service/file/file_upload.go b/internal/service/file/file_upload.go index 48f9118f89..f7dff13e85 100644 --- a/internal/service/file/file_upload.go +++ b/internal/service/file/file_upload.go @@ -1,6 +1,7 @@ package file import ( + "context" "fmt" "io" "mime/multipart" @@ -14,7 +15,7 @@ import ( ) // UploadFile uploads files to a folder -func (s *FileService) UploadFile(tenantID, parentID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) { +func (s *FileService) UploadFile(ctx context.Context, tenantID, parentID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) { if parentID == "" { rootFolder, err := s.fileDAO.GetRootFolder(tenantID) if err != nil { @@ -33,7 +34,7 @@ func (s *FileService) UploadFile(tenantID, parentID string, files []*multipart.F var maxNum int64 if _, err = fmt.Sscanf(maxFileNumPerUser, "%d", &maxNum); err == nil && maxNum > 0 { var docCount int64 - docCount, err = s.GetDocCount(tenantID) + docCount, err = s.GetDocCount(ctx, tenantID) if err != nil { return nil, fmt.Errorf("failed to get document count: %w", err) } @@ -134,7 +135,7 @@ func (s *FileService) UploadFile(tenantID, parentID string, files []*multipart.F // UploadInfos mirrors Python's upload_info file branch: store raw bytes in the // per-user downloads bucket and return lightweight upload descriptors instead // of creating full File rows in the file-management tree. -func (s *FileService) UploadInfos(userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) { +func (s *FileService) UploadInfos(ctx context.Context, userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) { storageImpl := storage.GetStorageFactory().GetStorage() if storageImpl == nil { return nil, fmt.Errorf("storage not initialized") @@ -143,7 +144,7 @@ func (s *FileService) UploadInfos(userID string, files []*multipart.FileHeader) results := make([]map[string]interface{}, 0, len(files)) for _, fileHeader := range files { filename := fileHeader.Filename - if err := s.checkUploadInfoHealth(userID, filename); err != nil { + if err := s.checkUploadInfoHealth(ctx, userID, filename); err != nil { return nil, err } src, err := fileHeader.Open() @@ -213,7 +214,7 @@ func (s *FileService) toUploadInfoResponse(file *entity.File, mimeType string) m } } -func (s *FileService) checkUploadInfoHealth(userID, filename string) error { +func (s *FileService) checkUploadInfoHealth(ctx context.Context, userID, filename string) error { if filename == "" { return fmt.Errorf("No file selected!") } @@ -222,7 +223,7 @@ func (s *FileService) checkUploadInfoHealth(userID, filename string) error { var maxNum int64 if _, err := fmt.Sscanf(maxFileNumPerUser, "%d", &maxNum); err == nil && maxNum > 0 { var docCount int64 - docCount, err = s.GetDocCount(userID) + docCount, err = s.GetDocCount(ctx, userID) if err != nil { return fmt.Errorf("failed to get document count: %w", err) } @@ -262,8 +263,8 @@ func (s *FileService) storeUploadInfoBlob(storageImpl storage.Storage, userID, f // UploadDocumentInfos is the document-level wrapper that stores uploaded blobs // without creating Document rows, then returns the file metadata including // size/mime-type/extension. -func (s *FileService) UploadDocumentInfos(userID string, files []*multipart.FileHeader) ([]map[string]interface{}, common.ErrorCode, error) { - data, err := s.UploadInfos(userID, files) +func (s *FileService) UploadDocumentInfos(ctx context.Context, userID string, files []*multipart.FileHeader) ([]map[string]interface{}, common.ErrorCode, error) { + data, err := s.UploadInfos(ctx, userID, files) if err != nil { return nil, common.CodeDataError, err } @@ -272,8 +273,8 @@ func (s *FileService) UploadDocumentInfos(userID string, files []*multipart.File // UploadDocumentInfoByURL fetches a remote URL, stores the content without // creating a Document row, then returns file metadata. -func (s *FileService) UploadDocumentInfoByURL(userID, rawURL string) (map[string]interface{}, common.ErrorCode, error) { - data, err := s.UploadFromURL(userID, rawURL) +func (s *FileService) UploadDocumentInfoByURL(ctx context.Context, userID, rawURL string) (map[string]interface{}, common.ErrorCode, error) { + data, err := s.UploadFromURL(ctx, userID, rawURL) if err != nil { return nil, common.CodeDataError, err } diff --git a/internal/service/file/file_url.go b/internal/service/file/file_url.go index 8e19965682..20b2aecacd 100644 --- a/internal/service/file/file_url.go +++ b/internal/service/file/file_url.go @@ -1,6 +1,7 @@ package file import ( + "context" "fmt" "net/url" "path/filepath" @@ -20,7 +21,7 @@ import ( // carries connect and overall timeouts, and the response body is bounded with // truncation detection so an oversized file is rejected rather than silently // clipped. -func (s *FileService) UploadFromURL(tenantID, rawURL string) (map[string]interface{}, error) { +func (s *FileService) UploadFromURL(ctx context.Context, tenantID, rawURL string) (map[string]interface{}, error) { if rawURL == "" { return nil, fmt.Errorf("url is required") } @@ -29,7 +30,7 @@ func (s *FileService) UploadFromURL(tenantID, rawURL string) (map[string]interfa return nil, fmt.Errorf("invalid or unsafe URL") } - data, headers, finalURL, err := utility.FetchRemoteFileSafely(rawURL, maxRemoteFileSize) + data, headers, finalURL, err := utility.FetchRemoteFileSafely(ctx, rawURL, maxRemoteFileSize) if err != nil { return nil, err } @@ -41,7 +42,7 @@ func (s *FileService) UploadFromURL(tenantID, rawURL string) (map[string]interfa contentType := headers.Get("Content-Type") filename := normalizeRemoteUploadFilename(finalURL, contentType, data) - if err := s.checkUploadInfoHealth(tenantID, filename); err != nil { + if err = s.checkUploadInfoHealth(ctx, tenantID, filename); err != nil { return nil, err } filename, contentType, data = utility.NormalizeUploadInfoContent(filename, contentType, data) diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go index 7da77bb64b..6469c0b6ae 100644 --- a/internal/service/ingestion_task_service.go +++ b/internal/service/ingestion_task_service.go @@ -1,6 +1,7 @@ package service import ( + "context" "errors" "fmt" "time" @@ -73,7 +74,7 @@ func (s *IngestionTaskService) ListByUser(userID string, datasetID *string, page return s.ingestionTaskDAO.ListByUserIDAndDatasetID(userID, *datasetID, page, pageSize) } -func (s *IngestionTaskService) CreateForDocuments(datasetID, userID string, docIDs []string) ([]*ParseDocumentResponse, error) { +func (s *IngestionTaskService) CreateForDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*ParseDocumentResponse, error) { uniqueDocIDs := common.Deduplicate(docIDs) if len(uniqueDocIDs) == 0 { return nil, fmt.Errorf("no documents to parse") @@ -81,7 +82,7 @@ func (s *IngestionTaskService) CreateForDocuments(datasetID, userID string, docI responses := make([]*ParseDocumentResponse, 0, len(uniqueDocIDs)) for _, docID := range uniqueDocIDs { - doc, err := s.documentDAO.GetByID(docID) + doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID) if err != nil { responses = append(responses, &ParseDocumentResponse{ DocumentID: docID, @@ -104,7 +105,7 @@ func (s *IngestionTaskService) CreateForDocuments(datasetID, userID string, docI Schema: nil, Status: common.CREATED, } - task, err = s.CreateAndEnqueue(task) + task, err = s.CreateAndEnqueue(ctx, task) if err != nil { responses = append(responses, &ParseDocumentResponse{ DocumentID: docID, @@ -164,7 +165,8 @@ func (s *IngestionTaskService) ListAllForAdmin() ([]map[string]interface{}, erro showTasks := make([]map[string]interface{}, 0, len(ingestionTasks)) for _, task := range ingestionTasks { - user, err := s.userDAO.GetByTenantID(task.UserID) + var user *entity.User + user, err = s.userDAO.GetByTenantID(task.UserID) if err != nil { return nil, err } @@ -198,14 +200,14 @@ func (s *IngestionTaskService) ListAllForAdmin() ([]map[string]interface{}, erro return showTasks, nil } -func (s *IngestionTaskService) StartRunning(taskID string) (*entity.IngestionTask, error) { +func (s *IngestionTaskService) StartRunning(ctx context.Context, taskID string) (*entity.IngestionTask, error) { task, err := s.GetTask(taskID) if err != nil { return nil, err } switch task.Status { case common.CREATED: - task, err := s.transition(taskID, common.RUNNING) + task, err = s.transition(taskID, common.RUNNING) if err != nil { return nil, err } @@ -215,7 +217,7 @@ func (s *IngestionTaskService) StartRunning(taskID string) (*entity.IngestionTas // the task transition and trigger a redelivery loop. run uses the // document's numeric TaskStatus enum ("1"), not the task's string // status label. - if err := s.documentDAO.UpdateByID(task.DocumentID, map[string]interface{}{ + if err = s.documentDAO.UpdateByID(ctx, dao.DB, task.DocumentID, map[string]interface{}{ "run": string(entity.TaskStatusRunning), "progress": float64(0), "chunk_num": int64(0), @@ -244,7 +246,7 @@ func (s *IngestionTaskService) RequestStop(taskID string) (*entity.IngestionTask case common.CREATED: return s.transition(taskID, common.STOPPED) case common.RUNNING: - task, err := s.transition(taskID, common.STOPPING) + task, err = s.transition(taskID, common.STOPPING) if err != nil { return nil, err } @@ -354,7 +356,7 @@ func (s *IngestionTaskService) transition(taskID string, to string) (*entity.Ing if err != nil { return nil, err } - if err := validateTransition(task.Status, to); err != nil { + if err = validateTransition(task.Status, to); err != nil { var transitionErr *InvalidTaskTransitionError if errors.As(err, &transitionErr) { return task, &InvalidTaskTransitionError{TaskID: taskID, From: transitionErr.From, To: transitionErr.To} @@ -372,7 +374,7 @@ func (s *IngestionTaskService) transition(taskID string, to string) (*entity.Ing return task, nil } -func (s *IngestionTaskService) CreateAndEnqueue(task *entity.IngestionTask) (*entity.IngestionTask, error) { +func (s *IngestionTaskService) CreateAndEnqueue(ctx context.Context, task *entity.IngestionTask) (*entity.IngestionTask, error) { existing, err := s.ingestionTaskDAO.GetByDocumentID(task.DocumentID) if err != nil { return nil, err @@ -385,7 +387,7 @@ func (s *IngestionTaskService) CreateAndEnqueue(task *entity.IngestionTask) (*en if err != nil { return nil, err } - if err := s.enqueueTask(existing.ID); err != nil { + if err = s.enqueueTask(existing.ID); err != nil { if rollbackErr := s.rollbackRetriedTask(existing.ID, originalStatus); rollbackErr != nil { return nil, fmt.Errorf("enqueue task %s: %w (rollback failed: %v)", existing.ID, err, rollbackErr) } @@ -400,7 +402,7 @@ func (s *IngestionTaskService) CreateAndEnqueue(task *entity.IngestionTask) (*en if err != nil { return nil, err } - if err := s.enqueueTask(created.ID); err != nil { + if err = s.enqueueTask(created.ID); err != nil { if rollbackErr := s.rollbackCreatedTask(created.ID); rollbackErr != nil { return nil, fmt.Errorf("enqueue task %s: %w (rollback failed: %v)", created.ID, err, rollbackErr) } diff --git a/internal/service/ingestion_task_service_test.go b/internal/service/ingestion_task_service_test.go index 9c65ca63c2..e9a2442060 100644 --- a/internal/service/ingestion_task_service_test.go +++ b/internal/service/ingestion_task_service_test.go @@ -31,7 +31,8 @@ func TestIngestionTaskServiceCreateForDocumentsPublishesTaskMessages(t *testing. svc := NewIngestionTaskService() svc.taskPublisher = publisher - resp, err := svc.CreateForDocuments("kb-1", "user-1", []string{"doc-1"}) + ctx := t.Context() + resp, err := svc.CreateForDocuments(ctx, "kb-1", "user-1", []string{"doc-1"}) if err != nil { t.Fatalf("CreateForDocuments failed: %v", err) } @@ -204,7 +205,8 @@ func TestIngestionTaskServiceStartRunningTransitionsCreatedTask(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := NewIngestionTaskService() - task, err := svc.StartRunning("task-1") + ctx := t.Context() + task, err := svc.StartRunning(ctx, "task-1") if err != nil { t.Fatalf("StartRunning failed: %v", err) } @@ -236,7 +238,8 @@ func TestStartRunningMarksDocumentRunning(t *testing.T) { } svc := NewIngestionTaskService() - if _, err := svc.StartRunning("task-1"); err != nil { + ctx := t.Context() + if _, err := svc.StartRunning(ctx, "task-1"); err != nil { t.Fatalf("StartRunning failed: %v", err) } @@ -289,7 +292,8 @@ func TestStartRunningLeavesTerminalDocumentUntouched(t *testing.T) { } svc := NewIngestionTaskService() - task, err := svc.StartRunning("task-1") + ctx := t.Context() + task, err := svc.StartRunning(ctx, "task-1") if err != nil { t.Fatalf("StartRunning failed: %v", err) } @@ -434,6 +438,7 @@ func TestIngestionTaskServiceCreateAndEnqueueRetriesTerminalTask(t *testing.T) { {name: "stopped", status: common.STOPPED}, } + ctx := t.Context() for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { publisher.subject = "" @@ -446,7 +451,7 @@ func TestIngestionTaskServiceCreateAndEnqueueRetriesTerminalTask(t *testing.T) { t.Fatalf("set terminal status: %v", err) } - task, err := svc.CreateAndEnqueue(&entity.IngestionTask{ + task, err := svc.CreateAndEnqueue(ctx, &entity.IngestionTask{ DocumentID: "doc-1", UserID: "user-1", DatasetID: "kb-1", @@ -483,7 +488,8 @@ func TestIngestionTaskServiceCreateAndEnqueueRejectsActiveExistingTask(t *testin svc := NewIngestionTaskService() svc.taskPublisher = publisher - _, err := svc.CreateAndEnqueue(&entity.IngestionTask{DocumentID: "doc-1", UserID: "user-1", DatasetID: "kb-1", Status: common.CREATED}) + ctx := t.Context() + _, err := svc.CreateAndEnqueue(ctx, &entity.IngestionTask{DocumentID: "doc-1", UserID: "user-1", DatasetID: "kb-1", Status: common.CREATED}) if err == nil { t.Fatal("expected CreateAndEnqueue to reject existing created task") } @@ -499,7 +505,8 @@ func TestIngestionTaskServiceCreateAndEnqueueRollsBackNewTaskOnPublishFailure(t svc := NewIngestionTaskService() svc.taskPublisher = publisher - _, err := svc.CreateAndEnqueue(&entity.IngestionTask{ + ctx := t.Context() + _, err := svc.CreateAndEnqueue(ctx, &entity.IngestionTask{ DocumentID: "doc-1", UserID: "user-1", DatasetID: "kb-1", @@ -529,7 +536,8 @@ func TestIngestionTaskServiceCreateAndEnqueueRollsBackRetriedTaskOnPublishFailur svc := NewIngestionTaskService() svc.taskPublisher = publisher - _, err := svc.CreateAndEnqueue(&entity.IngestionTask{ + ctx := t.Context() + _, err := svc.CreateAndEnqueue(ctx, &entity.IngestionTask{ DocumentID: "doc-1", UserID: "user-1", DatasetID: "kb-1", diff --git a/internal/service/nlp/retrieval.go b/internal/service/nlp/retrieval.go index 6a691af523..addc1dabc3 100644 --- a/internal/service/nlp/retrieval.go +++ b/internal/service/nlp/retrieval.go @@ -152,7 +152,7 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) } // Prune deleted chunks - searchResult, err = s.PruneDeletedChunks(searchResult) + searchResult, err = s.PruneDeletedChunks(ctx, searchResult) if err != nil { return nil, fmt.Errorf("PruneDeletedChunks failed: %w", err) } @@ -970,7 +970,7 @@ func RetrievalByChildren(chunks []map[string]interface{}, tenantIDs []string, do } // PruneDeletedChunks removes chunks whose documents no longer exist -func (s *RetrievalService) PruneDeletedChunks(result *RetrievalSearchResult) (*RetrievalSearchResult, error) { +func (s *RetrievalService) PruneDeletedChunks(ctx context.Context, result *RetrievalSearchResult) (*RetrievalSearchResult, error) { if s.documentDAO == nil { return nil, fmt.Errorf("documentDAO is not initialized") } @@ -997,7 +997,7 @@ func (s *RetrievalService) PruneDeletedChunks(result *RetrievalSearchResult) (*R } // Get existing document IDs - docs, err := s.documentDAO.GetByIDs(uniqueDocIDs) + docs, err := s.documentDAO.GetByIDs(ctx, dao.DB, uniqueDocIDs) if err != nil { return nil, fmt.Errorf("GetByIDs failed: %w", err) } diff --git a/internal/service/user.go b/internal/service/user.go index fbef5d8e7c..8e2b90548e 100644 --- a/internal/service/user.go +++ b/internal/service/user.go @@ -232,26 +232,26 @@ func (s *UserService) Register(req *RegisterRequest) (*entity.User, common.Error } db := dao.GetDB() - if err := db.Transaction(func(tx *gorm.DB) error { - if err := tx.Create(user).Error; err != nil { + if err = db.Transaction(func(tx *gorm.DB) error { + if err = tx.Create(user).Error; err != nil { return fmt.Errorf("failed to create user: %w", err) } - if err := tx.Create(tenant).Error; err != nil { + if err = tx.Create(tenant).Error; err != nil { return fmt.Errorf("failed to create tenant: %w", err) } - if err := tx.Create(userTenant).Error; err != nil { + if err = tx.Create(userTenant).Error; err != nil { return fmt.Errorf("failed to create user tenant relation: %w", err) } if len(tenantLLMs) > 0 { - if err := tx.Create(&tenantLLMs).Error; err != nil { + if err = tx.Create(&tenantLLMs).Error; err != nil { return fmt.Errorf("failed to create tenant llm: %w", err) } } - if err := tx.Create(rootFile).Error; err != nil { + if err = tx.Create(rootFile).Error; err != nil { return fmt.Errorf("failed to create root folder: %w", err) } return nil diff --git a/internal/utility/upload.go b/internal/utility/upload.go index 4d357c121b..030ef1c4e6 100644 --- a/internal/utility/upload.go +++ b/internal/utility/upload.go @@ -17,6 +17,7 @@ package utility import ( + "context" "fmt" "html" "io" @@ -37,7 +38,7 @@ var ( // FetchRemoteFileSafely downloads rawURL with SSRF protection, connect/overall // timeouts, and a hard size cap that rejects (rather than truncates) oversized // bodies. -func FetchRemoteFileSafely(rawURL string, maxSize int64) ([]byte, http.Header, string, error) { +func FetchRemoteFileSafely(ctx context.Context, rawURL string, maxSize int64) ([]byte, http.Header, string, error) { currentURL := rawURL for redirects := 0; redirects < 10; redirects++ { hostname, resolvedIP, err := AssertURLSafe(currentURL) diff --git a/internal/utility/upload_test.go b/internal/utility/upload_test.go index 0016c67c54..fee42df9e3 100644 --- a/internal/utility/upload_test.go +++ b/internal/utility/upload_test.go @@ -55,7 +55,8 @@ func TestFetchRemoteFileSafely_PDFAddsExtension(t *testing.T) { PinnedHTTPClient = origPinned }) - data, headers, _, err := FetchRemoteFileSafely(server.URL+"/report", 100<<20) + ctx := t.Context() + data, headers, _, err := FetchRemoteFileSafely(ctx, server.URL+"/report", 100<<20) if err != nil { t.Fatalf("FetchRemoteFileSafely failed: %v", err) } @@ -87,7 +88,8 @@ func TestFetchRemoteFileSafely_ReturnsContentAndHeaders(t *testing.T) { PinnedHTTPClient = origPinned }) - data, headers, finalURL, err := FetchRemoteFileSafely(server.URL+"/page", 100<<20) + ctx := t.Context() + data, headers, finalURL, err := FetchRemoteFileSafely(ctx, server.URL+"/page", 100<<20) if err != nil { t.Fatalf("FetchRemoteFileSafely failed: %v", err) }