Go: add context to storage (#17690)

Signed-off-by: Jin Hai <haijin.chn@gmail.com>

internal/ingestion/component/document_storage.go

@@ -45,8 +45,8 @@ var ResolveDocumentStorageOverride func(docID string) (*DocumentStorageRef, erro
 // the source PDF to crop section images on demand) can reuse the same
 // storage resolution the Parser uses.
 func FetchBinary(ctx context.Context, bucket, path string) ([]byte, error) {
-	stg := resolveStorage()
-	if stg == nil {
+	storageImpl := resolveStorage()
+	if storageImpl == nil {
 		return nil, fmt.Errorf("no storage backend registered")
 	}
 
@@ -56,7 +56,7 @@ func FetchBinary(ctx context.Context, bucket, path string) ([]byte, error) {
 	}
 	done := make(chan result, 1)
 	go func() {
-		data, err := stg.Get(bucket, path)
+		data, err := storageImpl.Get(ctx, bucket, path)
 		done <- result{data: data, err: err}
 	}()
 	select {

internal/ingestion/component/extractor_tag.go

@@ -245,7 +245,7 @@ func (c *ExtractorComponent) loadTagFileIndexed(ctx context.Context) (*indexedTa
 		return nil, false
 	}
 	tenantID := globals.GlobalOrInput(ctx, nil, "tenant_id", "")
-	data, err := stg.Get(f.ParentID, *f.Location, tenantID)
+	data, err := stg.Get(ctx, f.ParentID, *f.Location, tenantID)
 	if err != nil {
 		common.Warn(fmt.Sprintf("extractor tags: load tag source %q/%q: %v", f.ParentID, *f.Location, err))
 		return nil, false

internal/ingestion/component/file_test.go

@@ -83,7 +83,8 @@ func TestFileComponent_Registered(t *testing.T) {
 // even when explicit storage overrides are supplied.
 func TestFileComponent_Invoke_HappyPath(t *testing.T) {
 	ms := withMemoryStorage(t)
-	if err := ms.Put("bucketA", "path/to/file.txt", []byte("hello, ragflow")); err != nil {
+	ctx := t.Context()
+	if err := ms.Put(ctx, "bucketA", "path/to/file.txt", []byte("hello, ragflow")); err != nil {
 		t.Fatalf("seed: %v", err)
 	}
 
@@ -113,9 +114,10 @@ func TestFileComponent_Invoke_HappyPath(t *testing.T) {
 
 func TestFileComponent_Invoke_ResolvesDocIDViaDocumentLocation(t *testing.T) {
 	ms := withMemoryStorage(t)
+	ctx := t.Context()
 	db := withFileComponentTestDB(t)
 	location := "docs/from-document.bin"
-	if err := ms.Put("kb-doc", location, []byte("doc-location")); err != nil {
+	if err := ms.Put(ctx, "kb-doc", location, []byte("doc-location")); err != nil {
 		t.Fatalf("seed storage: %v", err)
 	}
 	docName := "report.pdf"
@@ -151,9 +153,10 @@ func TestFileComponent_Invoke_ResolvesDocIDViaDocumentLocation(t *testing.T) {
 
 func TestFileComponent_Invoke_ResolvesDocIDViaFileMapping(t *testing.T) {
 	ms := withMemoryStorage(t)
+	ctx := t.Context()
 	db := withFileComponentTestDB(t)
 	location := "tenant-root/from-file.bin"
-	if err := ms.Put("folder-1", location, []byte("file-mapping")); err != nil {
+	if err := ms.Put(ctx, "folder-1", location, []byte("file-mapping")); err != nil {
 		t.Fatalf("seed storage: %v", err)
 	}
 	docName := "deck.pptx"
@@ -249,8 +252,9 @@ func TestFileComponent_Invoke_MissingDoc(t *testing.T) {
 // bucket/path overrides still flow through for downstream Parser use.
 func TestFileComponent_Invoke_IncludesCheckpointPath(t *testing.T) {
 	ms := withMemoryStorage(t)
+	ctx := t.Context()
 	const wantPath = "checkpoint/expected/path.bin"
-	if err := ms.Put("b", wantPath, []byte("x")); err != nil {
+	if err := ms.Put(ctx, "b", wantPath, []byte("x")); err != nil {
 		t.Fatalf("seed: %v", err)
 	}
 	c := &FileComponent{}

internal/ingestion/component/image_uploader.go

@@ -32,12 +32,12 @@ type ImageUploader func(ctx context.Context, kbID, chunkID string, data []byte)
 // returns img_id "<kb_id>-<chunk_id>". Mirrors Python image2id's storage_put +
 // f"{bucket}-{objname}" (rag/utils/base64_image.py:80-82), minus re-encoding:
 // the bytes are stored in whatever format the chunker produced.
-func DefaultImageUploader(_ context.Context, kbID, chunkID string, data []byte) (string, error) {
-	stg := resolveStorage()
-	if stg == nil {
+func DefaultImageUploader(ctx context.Context, kbID, chunkID string, data []byte) (string, error) {
+	storageImpl := resolveStorage()
+	if storageImpl == nil {
 		return "", fmt.Errorf("no storage backend registered")
 	}
-	if err := stg.Put(kbID, chunkID, data); err != nil {
+	if err := storageImpl.Put(ctx, kbID, chunkID, data); err != nil {
 		return "", fmt.Errorf("store chunk image (%q,%q): %w", kbID, chunkID, err)
 	}
 	return kbID + "-" + chunkID, nil

internal/ingestion/component/parser_test.go

@@ -244,7 +244,8 @@ func TestParserComponent_Invoke_ResolvesBinaryFromDocID(t *testing.T) {
 	ms := withMemoryStorage(t)
 	db := withFileComponentTestDB(t)
 	location := "docs/from-parser.txt"
-	if err := ms.Put("kb-parser", location, []byte("alpha\fbeta")); err != nil {
+	ctx := t.Context()
+	if err := ms.Put(ctx, "kb-parser", location, []byte("alpha\fbeta")); err != nil {
 		t.Fatalf("seed storage: %v", err)
 	}
 	docName := "parser.txt"
@@ -263,7 +264,7 @@ func TestParserComponent_Invoke_ResolvesBinaryFromDocID(t *testing.T) {
 	}
 
 	c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
-	out, err := c.Invoke(context.Background(), db, map[string]any{"doc_id": "doc-parser"})
+	out, err := c.Invoke(ctx, db, map[string]any{"doc_id": "doc-parser"})
 	if err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
@@ -281,12 +282,13 @@ func TestParserComponent_Invoke_ResolvesBinaryFromDocID(t *testing.T) {
 
 func TestParserComponent_Invoke_ResolvesBinaryFromBucketPath(t *testing.T) {
 	ms := withMemoryStorage(t)
-	if err := ms.Put("bucket-1", "docs/explicit.txt", []byte("bucket content")); err != nil {
+	ctx := t.Context()
+	if err := ms.Put(ctx, "bucket-1", "docs/explicit.txt", []byte("bucket content")); err != nil {
 		t.Fatalf("seed storage: %v", err)
 	}
 
 	c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
-	out, err := c.Invoke(context.Background(), nil, map[string]any{
+	out, err := c.Invoke(ctx, nil, map[string]any{
 		"bucket": "bucket-1",
 		"path":   "docs/explicit.txt",
 	})