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

Go: add context to storage (#17690)

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

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

@@ -32,6 +32,7 @@ import (
func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) {
requireTokenizerPool(t)
ctx := t.Context()
mustLoadTaskTestConfig(t)
origDB := dao.DB
@@ -55,7 +56,7 @@ func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) {
}
templateBytes = disableTokenizerEmbeddingForTaskTemplate(t, templateBytes)
var templateDSL entity.JSONMap
if err := json.Unmarshal(templateBytes, &templateDSL); err != nil {
if err = json.Unmarshal(templateBytes, &templateDSL); err != nil {
t.Fatalf("unmarshal template dsl: %v", err)
}
@@ -71,10 +72,10 @@ func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) {
content := "Alpha paragraph\n\nBeta paragraph."
mustSeedTaskRealPipelineDocument(t, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath, docName, content)
if err := realDB.Model(&entity.Document{}).Where("id = ?", docID).Update("pipeline_id", canvasID).Error; err != nil {
if err = realDB.Model(&entity.Document{}).Where("id = ?", docID).Update("pipeline_id", canvasID).Error; err != nil {
t.Fatalf("set document pipeline_id: %v", err)
}
if err := realDB.Create(&entity.UserCanvas{
if err = realDB.Create(&entity.UserCanvas{
ID: canvasID,
UserID: tenantID,
Permission: "me",
@@ -84,8 +85,9 @@ func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) {
t.Fatalf("create user canvas: %v", err)
}
t.Cleanup(func() {
_ = realDB.Where("id = ?", canvasID).Delete(&entity.UserCanvas{}).Error
cleanupTaskRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath)
cleanUpCtx := context.Background()
_ = realDB.WithContext(cleanUpCtx).Where("id = ?", canvasID).Delete(&entity.UserCanvas{}).Error
cleanupTaskRealPipelineDocument(cleanUpCtx, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath)
})
taskCtx := &TaskContext{
@@ -115,7 +117,7 @@ func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) {
return nil, nil
})
if _, err := svc.Execute(context.Background()); err != nil {
if _, err = svc.Execute(ctx); err != nil {
t.Fatalf("Run: %v", err)
}
if len(inserted) != 1 {
@@ -136,6 +138,7 @@ func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) {
func TestPipelineExecutor_Run_RealPDF_ProducesIndexedChunks(t *testing.T) {
requireTokenizerPool(t)
ctx := t.Context()
// Loads service config (server.Init side effect) without requiring any
// external MySQL/MinIO/ES. The pipeline runs against an in-memory sqlite
@@ -198,7 +201,7 @@ func TestPipelineExecutor_Run_RealPDF_ProducesIndexedChunks(t *testing.T) {
}
t.Cleanup(func() {
_ = realDB.Where("id = ?", canvasID).Delete(&entity.UserCanvas{}).Error
cleanupTaskRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath)
cleanupTaskRealPipelineDocument(ctx, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath)
})
taskCtx := &TaskContext{
@@ -229,7 +232,6 @@ func TestPipelineExecutor_Run_RealPDF_ProducesIndexedChunks(t *testing.T) {
}).
WithLogCreateFunc(func(ctx context.Context, db *gorm.DB, log *entity.PipelineOperationLog) error { return nil })
ctx := t.Context()
if _, err = svc.Execute(ctx); err != nil {
t.Fatalf("Run: %v", err)
}
@@ -269,6 +271,7 @@ func TestPipelineExecutor_Run_RealPDF_ProducesIndexedChunks(t *testing.T) {
func TestRunPipeline_RealPipelineOutput_ProducesIndexFields(t *testing.T) {
requireTokenizerPool(t)
ctx := t.Context()
mustLoadTaskTestConfig(t)
origDB := dao.DB
@@ -309,10 +312,9 @@ func TestRunPipeline_RealPipelineOutput_ProducesIndexFields(t *testing.T) {
mustSeedTaskRealPipelineDocument(t, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath, docName, content)
t.Cleanup(func() {
cleanupTaskRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath)
cleanupTaskRealPipelineDocument(ctx, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath)
})
ctx := t.Context()
pipelineOut, err := pipe.Run(ctx, map[string]any{
"doc_id": docID,
}, nil)
@@ -618,7 +620,8 @@ func mustSeedTaskRealPipelineDocumentBytes(
}).Error; err != nil {
t.Fatalf("create kb: %v", err)
}
if err := stg.Put(bucket, objectPath, content); err != nil {
ctx := t.Context()
if err := stg.Put(ctx, bucket, objectPath, content); err != nil {
t.Fatalf("put real minio object: %v", err)
}
if err := db.Create(&entity.File{
@@ -657,13 +660,13 @@ func mustSeedTaskRealPipelineDocumentBytes(
}
}
func cleanupTaskRealPipelineDocument(db *gorm.DB, stg storage.Storage, tenantID, kbID, docID, fileID, bucket, objectPath string) {
_ = db.Where("document_id = ?", docID).Delete(&entity.File2Document{}).Error
_ = db.Where("id = ?", docID).Delete(&entity.Document{}).Error
_ = db.Where("id = ?", fileID).Delete(&entity.File{}).Error
_ = db.Where("id = ?", kbID).Delete(&entity.Knowledgebase{}).Error
_ = db.Where("id = ?", tenantID).Delete(&entity.Tenant{}).Error
_ = stg.Remove(bucket, objectPath)
func cleanupTaskRealPipelineDocument(ctx context.Context, db *gorm.DB, storageImpl storage.Storage, tenantID, kbID, docID, fileID, bucket, objectPath string) {
_ = db.WithContext(ctx).Where("document_id = ?", docID).Delete(&entity.File2Document{}).Error
_ = db.WithContext(ctx).Where("id = ?", docID).Delete(&entity.Document{}).Error
_ = db.WithContext(ctx).Where("id = ?", fileID).Delete(&entity.File{}).Error
_ = db.WithContext(ctx).Where("id = ?", kbID).Delete(&entity.Knowledgebase{}).Error
_ = db.WithContext(ctx).Where("id = ?", tenantID).Delete(&entity.Tenant{}).Error
_ = storageImpl.Remove(ctx, bucket, objectPath)
}
func deepCopyTaskChunks(in []map[string]any) []map[string]any {