zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-05 07:10:29 +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

57
internal/storage/gcs.go
View File

@@ -37,21 +37,19 @@ type GCSStorage struct {
}
// NewGCSStorage creates a new GCS storage instance
func NewGCSStorage(config config.GCSConfig) (*GCSStorage, error) {
func NewGCSStorage(ctx context.Context, config config.GCSConfig) (*GCSStorage, error) {
gcsStorage := &GCSStorage{
config: config,
}
if err := gcsStorage.connect(); err != nil {
if err := gcsStorage.connect(ctx); err != nil {
return nil, err
}
return gcsStorage, nil
}
func (m *GCSStorage) connect() error {
ctx := context.Background()
func (m *GCSStorage) connect(ctx context.Context) error {
client, err := storage.NewClient(ctx)
if err != nil {
@@ -62,8 +60,8 @@ func (m *GCSStorage) connect() error {
return nil
}
func (m *GCSStorage) reconnect() {
if err := m.connect(); err != nil {
func (m *GCSStorage) reconnect(ctx context.Context) {
if err := m.connect(ctx); err != nil {
common.Fatal(fmt.Sprintf("Failed to reconnect to GCS, %s", err.Error()))
}
}
@@ -71,13 +69,12 @@ func (m *GCSStorage) reconnect() {
func (m *GCSStorage) Type() string { return "gcs" }
// Health checks GCS service availability
func (m *GCSStorage) Health() bool {
return m.BucketExists(m.config.Bucket)
func (m *GCSStorage) Health(ctx context.Context) bool {
return m.BucketExists(ctx, m.config.Bucket)
}
// Put uploads an object to GCS
func (m *GCSStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
ctx := context.Background()
func (m *GCSStorage) Put(ctx context.Context, bucket, fnm string, binary []byte, tenantID ...string) error {
obj := m.client.Bucket(bucket).Object(fnm)
w := obj.NewWriter(ctx)
@@ -93,8 +90,7 @@ func (m *GCSStorage) Put(bucket, fnm string, binary []byte, tenantID ...string)
}
// Get retrieves an object from GCS
func (m *GCSStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
ctx := context.Background()
func (m *GCSStorage) Get(ctx context.Context, bucket, fnm string, tenantID ...string) ([]byte, error) {
r, err := m.client.Bucket(bucket).Object(fnm).NewReader(ctx)
if err != nil {
@@ -111,8 +107,7 @@ func (m *GCSStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error)
}
// Remove removes an object from GCS
func (m *GCSStorage) Remove(bucketName, objectName string, tenantID ...string) error {
ctx := context.Background()
func (m *GCSStorage) Remove(ctx context.Context, bucketName, objectName string, tenantID ...string) error {
obj := m.client.Bucket(bucketName).Object(objectName)
if err := obj.Delete(ctx); err != nil {
@@ -123,8 +118,7 @@ func (m *GCSStorage) Remove(bucketName, objectName string, tenantID ...string) e
}
// ObjExist checks if an object exists in GCS
func (m *GCSStorage) ObjExist(bucketName, objectName string, tenantID ...string) bool {
ctx := context.Background()
func (m *GCSStorage) ObjExist(ctx context.Context, bucketName, objectName string, tenantID ...string) bool {
obj := m.client.Bucket(bucketName).Object(objectName)
@@ -136,8 +130,7 @@ func (m *GCSStorage) ObjExist(bucketName, objectName string, tenantID ...string)
return true
}
func (m *GCSStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
ctx := context.Background()
func (m *GCSStorage) ListObjects(ctx context.Context, bucket string, tenantID ...string) ([]string, error) {
bucketObject := m.client.Bucket(bucket)
it := bucketObject.Objects(ctx, nil)
@@ -159,7 +152,7 @@ func (m *GCSStorage) ListObjects(bucket string, tenantID ...string) ([]string, e
}
// GetPresignedURL generates a presigned URL for accessing an object
func (m *GCSStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
func (m *GCSStorage) GetPresignedURL(ctx context.Context, bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
bucketObject := m.client.Bucket(bucket)
objectPath := fmt.Sprintf("%s/%s", bucket, fnm)
@@ -175,14 +168,12 @@ func (m *GCSStorage) GetPresignedURL(bucket, fnm string, expires time.Duration,
}
// BucketExists checks if a bucket exists
func (m *GCSStorage) BucketExists(bucket string) bool {
func (m *GCSStorage) BucketExists(ctx context.Context, bucket string) bool {
actualBucket := bucket
if m.config.Bucket != "" {
actualBucket = m.config.Bucket
}
ctx := context.Background()
_, err := m.client.Bucket(actualBucket).Attrs(ctx)
if err != nil {
return false
@@ -192,13 +183,11 @@ func (m *GCSStorage) BucketExists(bucket string) bool {
}
// RemoveBucket removes a bucket and all its objects
func (m *GCSStorage) RemoveBucket(bucketName string) error {
func (m *GCSStorage) RemoveBucket(ctx context.Context, bucketName string) error {
if bucketName == "" {
return fmt.Errorf("attempt to delete bucket without name")
}
ctx := context.Background()
bucket := m.client.Bucket(bucketName)
it := bucket.Objects(ctx, nil)
@@ -224,8 +213,8 @@ func (m *GCSStorage) RemoveBucket(bucketName string) error {
}
// Copy copies an object from source to destination
func (m *GCSStorage) Copy(srcBucket, srcObject, destBucket, destObject string) bool {
ctx := context.Background()
func (m *GCSStorage) Copy(ctx context.Context, srcBucket, srcObject, destBucket, destObject string) bool {
src := m.client.Bucket(srcBucket).Object(srcObject)
dst := m.client.Bucket(destBucket).Object(destObject)
copier := dst.CopierFrom(src)
@@ -238,10 +227,16 @@ func (m *GCSStorage) Copy(srcBucket, srcObject, destBucket, destObject string) b
}
// Move moves an object from source to destination
func (m *GCSStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
if m.Copy(srcBucket, srcPath, destBucket, destPath) {
if err := m.Remove(srcBucket, srcPath); err != nil {
func (m *GCSStorage) Move(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
if m.Copy(ctx, srcBucket, srcPath, destBucket, destPath) {
if err := m.Remove(ctx, srcBucket, srcPath); err != nil {
common.Warn("Failed to remove source object after copy", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
rollbackCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
defer cancel()
err = m.Remove(rollbackCtx, destBucket, destPath)
if err != nil {
common.Warn("Failed to roll back copied destination object", zap.String("bucket", destBucket), zap.String("key", destPath), zap.Error(err))
}
return false
}
return true

57
internal/storage/memory.go
View File

@@ -17,10 +17,14 @@
package storage
import (
"context"
"errors"
"fmt"
"ragflow/internal/common"
"sync"
"time"
"go.uber.org/zap"
)
// ErrMemoryNotFound is returned when a key does not exist in the in-memory backend.
@@ -49,13 +53,13 @@ func NewMemoryStorage() Storage {
func (m *MemoryStorage) Type() string { return "memory_storage" }
// Health always reports healthy for the in-memory backend.
func (m *MemoryStorage) Health() bool {
func (m *MemoryStorage) Health(ctx context.Context) bool {
return true
}
// Put uploads an object to the in-memory backend, creating the bucket
// on demand if it does not yet exist. The stored bytes are a defensive copy.
func (m *MemoryStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
func (m *MemoryStorage) Put(ctx context.Context, bucket, fnm string, binary []byte, tenantID ...string) error {
if bucket == "" {
return fmt.Errorf("memory storage: bucket is required")
}
@@ -63,6 +67,10 @@ func (m *MemoryStorage) Put(bucket, fnm string, binary []byte, tenantID ...strin
return fmt.Errorf("memory storage: key is required")
}
if err := ctx.Err(); err != nil {
return err
}
m.mu.Lock()
defer m.mu.Unlock()
@@ -80,7 +88,7 @@ func (m *MemoryStorage) Put(bucket, fnm string, binary []byte, tenantID ...strin
// Get retrieves an object from the in-memory backend. Returns
// ErrMemoryNotFound when the bucket or key is missing.
func (m *MemoryStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
func (m *MemoryStorage) Get(ctx context.Context, bucket, fnm string, tenantID ...string) ([]byte, error) {
m.mu.RLock()
defer m.mu.RUnlock()
@@ -100,7 +108,11 @@ func (m *MemoryStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, err
// Remove deletes an object from the in-memory backend. Removing a
// non-existent key is a no-op and returns nil.
func (m *MemoryStorage) Remove(bucket, fnm string, tenantID ...string) error {
func (m *MemoryStorage) Remove(ctx context.Context, bucket, fnm string, tenantID ...string) error {
if err := ctx.Err(); err != nil {
return err
}
m.mu.Lock()
defer m.mu.Unlock()
@@ -113,7 +125,7 @@ func (m *MemoryStorage) Remove(bucket, fnm string, tenantID ...string) error {
}
// ObjExist reports whether the given bucket and key are present.
func (m *MemoryStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
func (m *MemoryStorage) ObjExist(ctx context.Context, bucket, fnm string, tenantID ...string) bool {
m.mu.RLock()
defer m.mu.RUnlock()
@@ -125,7 +137,7 @@ func (m *MemoryStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
return ok
}
func (m *MemoryStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
func (m *MemoryStorage) ListObjects(ctx context.Context, bucket string, tenantID ...string) ([]string, error) {
m.mu.RLock()
defer m.mu.RUnlock()
@@ -146,7 +158,7 @@ func (m *MemoryStorage) ListObjects(bucket string, tenantID ...string) ([]string
// GetPresignedURL returns a deterministic, non-network URL string for tests.
// Format: memory://<bucket>/<key>?exp=<unix-seconds>
func (m *MemoryStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
func (m *MemoryStorage) GetPresignedURL(ctx context.Context, bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
m.mu.RLock()
defer m.mu.RUnlock()
@@ -163,7 +175,7 @@ func (m *MemoryStorage) GetPresignedURL(bucket, fnm string, expires time.Duratio
}
// BucketExists reports whether the named bucket has been created.
func (m *MemoryStorage) BucketExists(bucket string) bool {
func (m *MemoryStorage) BucketExists(ctx context.Context, bucket string) bool {
m.mu.RLock()
defer m.mu.RUnlock()
@@ -173,7 +185,11 @@ func (m *MemoryStorage) BucketExists(bucket string) bool {
// RemoveBucket deletes a bucket and all of its keys. Removing a
// non-existent bucket is a no-op and returns nil.
func (m *MemoryStorage) RemoveBucket(bucket string) error {
func (m *MemoryStorage) RemoveBucket(ctx context.Context, bucket string) error {
if err := ctx.Err(); err != nil {
return err
}
m.mu.Lock()
defer m.mu.Unlock()
@@ -184,7 +200,11 @@ func (m *MemoryStorage) RemoveBucket(bucket string) error {
// Copy duplicates an object from srcBucket/srcKey to destBucket/destKey.
// The source is left untouched. Returns false if the source does not exist
// or if the destination bucket creation fails.
func (m *MemoryStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
func (m *MemoryStorage) Copy(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
if err := ctx.Err(); err != nil {
return false
}
m.mu.RLock()
srcBucketMap, ok := m.objects[srcBucket]
if !ok {
@@ -213,11 +233,22 @@ func (m *MemoryStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bo
// Move transfers an object to a new location, deleting the source on success.
// Returns false if the source does not exist or the copy step fails.
func (m *MemoryStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
if !m.Copy(srcBucket, srcPath, destBucket, destPath) {
func (m *MemoryStorage) Move(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
if err := ctx.Err(); err != nil {
return false
}
if err := m.Remove(srcBucket, srcPath); err != nil {
if !m.Copy(ctx, srcBucket, srcPath, destBucket, destPath) {
return false
}
if err := m.Remove(ctx, srcBucket, srcPath); err != nil {
rollbackCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
defer cancel()
err = m.Remove(rollbackCtx, destBucket, destPath)
if err != nil {
common.Warn("Failed to roll back copied destination object", zap.String("bucket", destBucket), zap.String("key", destPath), zap.Error(err))
}
return false
}
return true

108
internal/storage/memory_test.go
View File

@@ -38,13 +38,14 @@ func newTestMemory(t *testing.T) *MemoryStorage {
func TestMemoryStorage_PutGet(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
payload := []byte("hello, world")
if err := ms.Put("b1", "k1", payload); err != nil {
if err := ms.Put(ctx, "b1", "k1", payload); err != nil {
t.Fatalf("Put returned error: %v", err)
}
got, err := ms.Get("b1", "k1")
got, err := ms.Get(ctx, "b1", "k1")
if err != nil {
t.Fatalf("Get returned error: %v", err)
}
@@ -54,7 +55,7 @@ func TestMemoryStorage_PutGet(t *testing.T) {
// Mutating the caller's slice after Put must not affect stored data.
payload[0] = 'X'
got2, err := ms.Get("b1", "k1")
got2, err := ms.Get(ctx, "b1", "k1")
if err != nil {
t.Fatalf("Get returned error: %v", err)
}
@@ -65,110 +66,115 @@ func TestMemoryStorage_PutGet(t *testing.T) {
func TestMemoryStorage_GetMissing(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
if _, err := ms.Get("missing-bucket", "k"); !errors.Is(err, ErrMemoryNotFound) {
if _, err := ms.Get(ctx, "missing-bucket", "k"); !errors.Is(err, ErrMemoryNotFound) {
t.Fatalf("Get on missing bucket: expected ErrMemoryNotFound, got %v", err)
}
if err := ms.Put("b1", "exists", []byte("data")); err != nil {
if err := ms.Put(ctx, "b1", "exists", []byte("data")); err != nil {
t.Fatalf("Put failed: %v", err)
}
if _, err := ms.Get("b1", "missing-key"); !errors.Is(err, ErrMemoryNotFound) {
if _, err := ms.Get(ctx, "b1", "missing-key"); !errors.Is(err, ErrMemoryNotFound) {
t.Fatalf("Get on missing key: expected ErrMemoryNotFound, got %v", err)
}
}
func TestMemoryStorage_ObjExist(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
if ms.ObjExist("b1", "k1") {
if ms.ObjExist(ctx, "b1", "k1") {
t.Fatalf("ObjExist on empty bucket returned true")
}
if err := ms.Put("b1", "k1", []byte("v")); err != nil {
if err := ms.Put(ctx, "b1", "k1", []byte("v")); err != nil {
t.Fatalf("Put failed: %v", err)
}
if !ms.ObjExist("b1", "k1") {
if !ms.ObjExist(ctx, "b1", "k1") {
t.Fatalf("ObjExist after Put returned false")
}
if ms.ObjExist("b1", "other") {
if ms.ObjExist(ctx, "b1", "other") {
t.Fatalf("ObjExist for sibling key returned true")
}
if ms.ObjExist("other-bucket", "k1") {
if ms.ObjExist(ctx, "other-bucket", "k1") {
t.Fatalf("ObjExist for sibling bucket returned true")
}
}
func TestMemoryStorage_Remove(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
// Idempotent: removing a key from a missing bucket is a no-op.
if err := ms.Remove("ghost", "k"); err != nil {
if err := ms.Remove(ctx, "ghost", "k"); err != nil {
t.Fatalf("Remove on missing bucket returned error: %v", err)
}
if err := ms.Put("b1", "k1", []byte("v")); err != nil {
if err := ms.Put(ctx, "b1", "k1", []byte("v")); err != nil {
t.Fatalf("Put failed: %v", err)
}
if err := ms.Remove("b1", "k1"); err != nil {
if err := ms.Remove(ctx, "b1", "k1"); err != nil {
t.Fatalf("Remove failed: %v", err)
}
if ms.ObjExist("b1", "k1") {
if ms.ObjExist(ctx, "b1", "k1") {
t.Fatalf("ObjExist after Remove returned true")
}
// Removing the same key again must not error.
if err := ms.Remove("b1", "k1"); err != nil {
if err := ms.Remove(ctx, "b1", "k1"); err != nil {
t.Fatalf("Remove on already-removed key returned error: %v", err)
}
}
func TestMemoryStorage_RemoveBucket(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
for _, k := range []string{"a", "b", "c"} {
if err := ms.Put("b1", k, []byte(k)); err != nil {
if err := ms.Put(ctx, "b1", k, []byte(k)); err != nil {
t.Fatalf("Put failed: %v", err)
}
}
if err := ms.Put("b2", "x", []byte("x")); err != nil {
if err := ms.Put(ctx, "b2", "x", []byte("x")); err != nil {
t.Fatalf("Put failed: %v", err)
}
if err := ms.RemoveBucket("b1"); err != nil {
if err := ms.RemoveBucket(ctx, "b1"); err != nil {
t.Fatalf("RemoveBucket failed: %v", err)
}
if ms.BucketExists("b1") {
if ms.BucketExists(ctx, "b1") {
t.Fatalf("BucketExists returned true after RemoveBucket")
}
if !ms.BucketExists("b2") {
if !ms.BucketExists(ctx, "b2") {
t.Fatalf("sibling bucket was removed unexpectedly")
}
// Idempotent: removing a missing bucket is a no-op.
if err := ms.RemoveBucket("b1"); err != nil {
if err := ms.RemoveBucket(ctx, "b1"); err != nil {
t.Fatalf("RemoveBucket on missing bucket returned error: %v", err)
}
}
func TestMemoryStorage_CopyMove(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
if err := ms.Put("src", "k", []byte("payload")); err != nil {
if err := ms.Put(ctx, "src", "k", []byte("payload")); err != nil {
t.Fatalf("Put failed: %v", err)
}
// Copy preserves source.
if !ms.Copy("src", "k", "dst", "k2") {
if !ms.Copy(ctx, "src", "k", "dst", "k2") {
t.Fatalf("Copy returned false on existing source")
}
if !ms.ObjExist("src", "k") {
if !ms.ObjExist(ctx, "src", "k") {
t.Fatalf("source missing after Copy")
}
if !ms.ObjExist("dst", "k2") {
if !ms.ObjExist(ctx, "dst", "k2") {
t.Fatalf("destination missing after Copy")
}
got, err := ms.Get("dst", "k2")
got, err := ms.Get(ctx, "dst", "k2")
if err != nil {
t.Fatalf("Get copy failed: %v", err)
}
@@ -177,53 +183,55 @@ func TestMemoryStorage_CopyMove(t *testing.T) {
}
// Move deletes the source.
if !ms.Move("src", "k", "dst2", "k3") {
if !ms.Move(ctx, "src", "k", "dst2", "k3") {
t.Fatalf("Move returned false on existing source")
}
if ms.ObjExist("src", "k") {
if ms.ObjExist(ctx, "src", "k") {
t.Fatalf("source still exists after Move")
}
if !ms.ObjExist("dst2", "k3") {
if !ms.ObjExist(ctx, "dst2", "k3") {
t.Fatalf("destination missing after Move")
}
// Copy/Move on missing source returns false.
if ms.Copy("src", "k", "dst", "k4") {
if ms.Copy(ctx, "src", "k", "dst", "k4") {
t.Fatalf("Copy on missing source returned true")
}
if ms.Move("src", "k", "dst", "k4") {
if ms.Move(ctx, "src", "k", "dst", "k4") {
t.Fatalf("Move on missing source returned true")
}
}
func TestMemoryStorage_BucketExists(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
if ms.BucketExists("b1") {
if ms.BucketExists(ctx, "b1") {
t.Fatalf("BucketExists returned true for empty backend")
}
if err := ms.Put("b1", "k", []byte("v")); err != nil {
if err := ms.Put(ctx, "b1", "k", []byte("v")); err != nil {
t.Fatalf("Put failed: %v", err)
}
if !ms.BucketExists("b1") {
if !ms.BucketExists(ctx, "b1") {
t.Fatalf("BucketExists returned false after Put")
}
if err := ms.RemoveBucket("b1"); err != nil {
if err := ms.RemoveBucket(ctx, "b1"); err != nil {
t.Fatalf("RemoveBucket failed: %v", err)
}
if ms.BucketExists("b1") {
if ms.BucketExists(ctx, "b1") {
t.Fatalf("BucketExists returned true after RemoveBucket")
}
}
func TestMemoryStorage_PresignedURL(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
if err := ms.Put("b1", "k1", []byte("v")); err != nil {
if err := ms.Put(ctx, "b1", "k1", []byte("v")); err != nil {
t.Fatalf("Put failed: %v", err)
}
url, err := ms.GetPresignedURL("b1", "k1", time.Minute)
url, err := ms.GetPresignedURL(ctx, "b1", "k1", time.Minute)
if err != nil {
t.Fatalf("GetPresignedURL failed: %v", err)
}
@@ -237,20 +245,23 @@ func TestMemoryStorage_PresignedURL(t *testing.T) {
t.Fatalf("presigned URL has unexpected scheme: %s", url)
}
if _, err := ms.GetPresignedURL("b1", "missing", time.Minute); !errors.Is(err, ErrMemoryNotFound) {
if _, err := ms.GetPresignedURL(ctx, "b1", "missing", time.Minute); !errors.Is(err, ErrMemoryNotFound) {
t.Fatalf("GetPresignedURL on missing key: expected ErrMemoryNotFound, got %v", err)
}
}
func TestMemoryStorage_Health(t *testing.T) {
ms := newTestMemory(t)
if !ms.Health() {
ctx := t.Context()
if !ms.Health(ctx) {
t.Fatalf("Health returned false for in-memory backend")
}
}
func TestMemoryStorage_Concurrent(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
const writers = 100
var wg sync.WaitGroup
@@ -261,7 +272,7 @@ func TestMemoryStorage_Concurrent(t *testing.T) {
defer wg.Done()
key := fmt.Sprintf("k-%d", i)
payload := []byte(fmt.Sprintf("payload-%d", i))
if err := ms.Put("race", key, payload); err != nil {
if err := ms.Put(ctx, "race", key, payload); err != nil {
t.Errorf("Put failed for %s: %v", key, err)
return
}
@@ -272,7 +283,7 @@ func TestMemoryStorage_Concurrent(t *testing.T) {
for i := 0; i < writers; i++ {
key := fmt.Sprintf("k-%d", i)
want := fmt.Sprintf("payload-%d", i)
got, err := ms.Get("race", key)
got, err := ms.Get(ctx, "race", key)
if err != nil {
t.Fatalf("Get %s failed: %v", key, err)
}
@@ -284,18 +295,19 @@ func TestMemoryStorage_Concurrent(t *testing.T) {
func TestMemoryStorage_Inspect(t *testing.T) {
ms := newTestMemory(t)
ctx := t.Context()
if got := ms.Inspect(); len(got) != 0 {
t.Fatalf("Inspect on empty backend returned %d entries", len(got))
}
if err := ms.Put("b1", "k1", []byte("12345")); err != nil {
if err := ms.Put(ctx, "b1", "k1", []byte("12345")); err != nil {
t.Fatalf("Put failed: %v", err)
}
if err := ms.Put("b1", "k2", []byte("hello")); err != nil {
if err := ms.Put(ctx, "b1", "k2", []byte("hello")); err != nil {
t.Fatalf("Put failed: %v", err)
}
if err := ms.Put("b2", "only", []byte("x")); err != nil {
if err := ms.Put(ctx, "b2", "only", []byte("x")); err != nil {
t.Fatalf("Put failed: %v", err)
}
@@ -325,10 +337,10 @@ func TestMemoryStorage_Inspect(t *testing.T) {
}
// After cleanup, Inspect should be empty again.
if err := ms.RemoveBucket("b1"); err != nil {
if err := ms.RemoveBucket(ctx, "b1"); err != nil {
t.Fatalf("RemoveBucket failed: %v", err)
}
if err := ms.RemoveBucket("b2"); err != nil {
if err := ms.RemoveBucket(ctx, "b2"); err != nil {
t.Fatalf("RemoveBucket failed: %v", err)
}
if got := ms.Inspect(); len(got) != 0 {

91
internal/storage/minio.go
View File

@@ -110,7 +110,7 @@ func (m *MinioStorage) resolveBucketAndPath(bucket, fnm string) (string, string)
func (m *MinioStorage) Type() string { return "minio" }
// Health checks MinIO service availability
func (m *MinioStorage) Health() bool {
func (m *MinioStorage) Health(ctx context.Context) bool {
cancelFunction, err := m.client.HealthCheck(time.Second * 5)
if cancelFunction != nil {
defer cancelFunction()
@@ -125,11 +125,9 @@ func (m *MinioStorage) Health() bool {
}
// Put uploads an object to MinIO
func (m *MinioStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
func (m *MinioStorage) Put(ctx context.Context, bucket, fnm string, binary []byte, tenantID ...string) error {
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
var err error
for i := 0; i < 3; i++ {
@@ -138,16 +136,26 @@ func (m *MinioStorage) Put(bucket, fnm string, binary []byte, tenantID ...string
if m.bucket == "" {
exists, err = m.client.BucketExists(ctx, bucket)
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return ctxErr
}
common.Warn("Failed to check bucket existence", zap.String("bucket", bucket), zap.Error(err))
m.reconnect()
time.Sleep(time.Second)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return err
}
continue
}
if !exists {
if err = m.client.MakeBucket(ctx, bucket, minio.MakeBucketOptions{}); err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return ctxErr
}
common.Warn("Failed to create bucket", zap.String("bucket", bucket), zap.Error(err))
m.reconnect()
time.Sleep(time.Second)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return err
}
continue
}
}
@@ -156,9 +164,14 @@ func (m *MinioStorage) Put(bucket, fnm string, binary []byte, tenantID ...string
reader := bytes.NewReader(binary)
_, err = m.client.PutObject(ctx, bucket, fnm, reader, int64(len(binary)), minio.PutObjectOptions{})
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return ctxErr
}
common.Warn("Failed to put object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
m.reconnect()
time.Sleep(time.Second)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return err
}
continue
}
@@ -169,26 +182,34 @@ func (m *MinioStorage) Put(bucket, fnm string, binary []byte, tenantID ...string
}
// Get retrieves an object from MinIO
func (m *MinioStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
func (m *MinioStorage) Get(ctx context.Context, bucket, fnm string, tenantID ...string) ([]byte, error) {
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
for i := 0; i < 2; i++ {
obj, err := m.client.GetObject(ctx, bucket, fnm, minio.GetObjectOptions{})
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, ctxErr
}
common.Warn("failed to get object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
m.reconnect()
time.Sleep(time.Second)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return nil, err
}
continue
}
defer obj.Close()
buf := new(bytes.Buffer)
if _, err = buf.ReadFrom(obj); err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, ctxErr
}
common.Warn("failed to read object data", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
m.reconnect()
time.Sleep(time.Second)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return nil, err
}
continue
}
@@ -199,11 +220,9 @@ func (m *MinioStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, erro
}
// Remove removes an object from MinIO
func (m *MinioStorage) Remove(bucket, fnm string, tenantID ...string) error {
func (m *MinioStorage) Remove(ctx context.Context, bucket, fnm string, tenantID ...string) error {
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
if err := m.client.RemoveObject(ctx, bucket, fnm, minio.RemoveObjectOptions{}); err != nil {
common.Warn("Failed to remove object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
return err
@@ -213,11 +232,9 @@ func (m *MinioStorage) Remove(bucket, fnm string, tenantID ...string) error {
}
// ObjExist checks if an object exists in MinIO
func (m *MinioStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
func (m *MinioStorage) ObjExist(ctx context.Context, bucket, fnm string, tenantID ...string) bool {
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
exists, err := m.client.BucketExists(ctx, bucket)
if err != nil || !exists {
return false
@@ -237,17 +254,20 @@ func (m *MinioStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
}
// GetPresignedURL generates a presigned URL for accessing an object
func (m *MinioStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
func (m *MinioStorage) GetPresignedURL(ctx context.Context, bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
for i := 0; i < 10; i++ {
url, err := m.client.PresignedGetObject(ctx, bucket, fnm, expires, nil)
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return "", ctxErr
}
common.Warn("Failed to get presigned URL", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
m.reconnect()
time.Sleep(time.Second)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return "", err
}
continue
}
@@ -258,14 +278,12 @@ func (m *MinioStorage) GetPresignedURL(bucket, fnm string, expires time.Duration
}
// BucketExists checks if a bucket exists
func (m *MinioStorage) BucketExists(bucket string) bool {
func (m *MinioStorage) BucketExists(ctx context.Context, bucket string) bool {
actualBucket := bucket
if m.bucket != "" {
actualBucket = m.bucket
}
ctx := context.Background()
exists, err := m.client.BucketExists(ctx, actualBucket)
if err != nil {
common.Warn("Failed to check bucket existence", zap.String("bucket", actualBucket), zap.Error(err))
@@ -275,8 +293,7 @@ func (m *MinioStorage) BucketExists(bucket string) bool {
return exists
}
func (m *MinioStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
ctx := context.Background()
func (m *MinioStorage) ListObjects(ctx context.Context, bucket string, tenantID ...string) ([]string, error) {
var objects []string
for obj := range m.client.ListObjects(ctx, bucket, minio.ListObjectsOptions{
@@ -294,7 +311,7 @@ func (m *MinioStorage) ListObjects(bucket string, tenantID ...string) ([]string,
}
// RemoveBucket removes a bucket and all its objects
func (m *MinioStorage) RemoveBucket(bucket string) error {
func (m *MinioStorage) RemoveBucket(ctx context.Context, bucket string) error {
actualBucket := bucket
origBucket := bucket
@@ -302,8 +319,6 @@ func (m *MinioStorage) RemoveBucket(bucket string) error {
actualBucket = m.bucket
}
ctx := context.Background()
// Build prefix for single-bucket mode
prefix := ""
if m.bucket != "" {
@@ -346,12 +361,10 @@ func (m *MinioStorage) RemoveBucket(bucket string) error {
}
// Copy copies an object from source to destination
func (m *MinioStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
func (m *MinioStorage) Copy(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
srcBucket, srcPath = m.resolveBucketAndPath(srcBucket, srcPath)
destBucket, destPath = m.resolveBucketAndPath(destBucket, destPath)
ctx := context.Background()
// Ensure destination bucket exists
if m.bucket == "" {
exists, err := m.client.BucketExists(ctx, destBucket)
@@ -394,10 +407,16 @@ func (m *MinioStorage) Copy(srcBucket, srcPath, destBucket, destPath string) boo
}
// Move moves an object from source to destination
func (m *MinioStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
if m.Copy(srcBucket, srcPath, destBucket, destPath) {
if err := m.Remove(srcBucket, srcPath); err != nil {
func (m *MinioStorage) Move(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
if m.Copy(ctx, srcBucket, srcPath, destBucket, destPath) {
if err := m.Remove(ctx, srcBucket, srcPath); err != nil {
common.Warn("Failed to remove source object after copy", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
rollbackCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
defer cancel()
err = m.Remove(rollbackCtx, destBucket, destPath)
if err != nil {
common.Warn("Failed to roll back copied destination object", zap.String("bucket", destBucket), zap.String("key", destPath), zap.Error(err))
}
return false
}
return true

143
internal/storage/oss.go
View File

@@ -21,6 +21,7 @@ import (
"context"
"errors"
"fmt"
"ragflow/internal/common"
"ragflow/internal/server/config"
"time"
@@ -42,22 +43,21 @@ type OSSStorage struct {
}
// NewOSSStorage creates a new OSS storage instance
func NewOSSStorage(config config.OSSConfig) (*OSSStorage, error) {
func NewOSSStorage(ctx context.Context, config config.OSSConfig) (*OSSStorage, error) {
storage := &OSSStorage{
bucket: config.Bucket,
prefixPath: config.PrefixPath,
config: config,
}
if err := storage.connect(); err != nil {
if err := storage.connect(ctx); err != nil {
return nil, err
}
return storage, nil
}
func (o *OSSStorage) connect() error {
ctx := context.Background()
func (o *OSSStorage) connect(ctx context.Context) error {
// Create static credentials
creds := credentials.NewStaticCredentialsProvider(
@@ -83,9 +83,9 @@ func (o *OSSStorage) connect() error {
return nil
}
func (o *OSSStorage) reconnect() {
if err := o.connect(); err != nil {
zap.L().Error("Failed to reconnect to OSS", zap.Error(err))
func (o *OSSStorage) reconnect(ctx context.Context) {
if err := o.connect(ctx); err != nil {
common.Error("Failed to reconnect to OSS", err)
}
}
@@ -106,7 +106,7 @@ func (o *OSSStorage) resolveBucketAndPath(bucket, fnm string) (string, string) {
func (o *OSSStorage) Type() string { return "oss" }
// Health checks OSS service availability
func (o *OSSStorage) Health() bool {
func (o *OSSStorage) Health(ctx context.Context) bool {
bucket := o.bucket
if bucket == "" {
bucket = "health-check-bucket"
@@ -118,15 +118,13 @@ func (o *OSSStorage) Health() bool {
}
binary := []byte("_t@@@1")
ctx := context.Background()
// Ensure bucket exists
if !o.BucketExists(bucket) {
if !o.BucketExists(ctx, bucket) {
_, err := o.client.CreateBucket(ctx, &s3.CreateBucketInput{
Bucket: aws.String(bucket),
})
if err != nil {
zap.L().Error("Failed to create bucket for health check", zap.String("bucket", bucket), zap.Error(err))
common.Error("Failed to create bucket for health check", err, zap.String("bucket", bucket))
return false
}
}
@@ -140,7 +138,7 @@ func (o *OSSStorage) Health() bool {
})
if err != nil {
zap.L().Error("Health check failed", zap.Error(err))
common.Error("Health check failed", err)
return false
}
@@ -148,24 +146,27 @@ func (o *OSSStorage) Health() bool {
}
// Put uploads an object to OSS
func (o *OSSStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
func (o *OSSStorage) Put(ctx context.Context, bucket, fnm string, binary []byte, tenantID ...string) error {
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
for i := 0; i < 2; i++ {
// Ensure bucket exists
if !o.BucketExists(bucket) {
if !o.BucketExists(ctx, bucket) {
_, err := o.client.CreateBucket(ctx, &s3.CreateBucketInput{
Bucket: aws.String(bucket),
})
if err != nil {
zap.L().Error("Failed to create bucket", zap.String("bucket", bucket), zap.Error(err))
o.reconnect()
time.Sleep(time.Second)
if ctxErr := ctx.Err(); ctxErr != nil {
return ctxErr
}
common.Error("Failed to create bucket", err, zap.String("bucket", bucket))
o.reconnect(ctx)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return err
}
continue
}
zap.L().Info("Created bucket", zap.String("bucket", bucket))
common.Info("Created bucket", zap.String("bucket", bucket))
}
reader := bytes.NewReader(binary)
@@ -175,9 +176,14 @@ func (o *OSSStorage) Put(bucket, fnm string, binary []byte, tenantID ...string)
Body: reader,
})
if err != nil {
zap.L().Error("Failed to put object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
o.reconnect()
time.Sleep(time.Second)
if ctxErr := ctx.Err(); ctxErr != nil {
return ctxErr
}
common.Error("Failed to put object", err, zap.String("bucket", bucket), zap.String("key", fnm))
o.reconnect(ctx)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return err
}
continue
}
@@ -188,29 +194,37 @@ func (o *OSSStorage) Put(bucket, fnm string, binary []byte, tenantID ...string)
}
// Get retrieves an object from OSS
func (o *OSSStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
func (o *OSSStorage) Get(ctx context.Context, bucket, fnm string, tenantID ...string) ([]byte, error) {
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
for i := 0; i < 2; i++ {
result, err := o.client.GetObject(ctx, &s3.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(fnm),
})
if err != nil {
zap.L().Error("Failed to get object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
o.reconnect()
time.Sleep(time.Second)
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, ctxErr
}
common.Error("Failed to get object", err, zap.String("bucket", bucket), zap.String("key", fnm))
o.reconnect(ctx)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return nil, err
}
continue
}
defer result.Body.Close()
buf := new(bytes.Buffer)
if _, err := buf.ReadFrom(result.Body); err != nil {
zap.L().Error("Failed to read object data", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
o.reconnect()
time.Sleep(time.Second)
if _, err = buf.ReadFrom(result.Body); err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, ctxErr
}
common.Error("Failed to read object data", err, zap.String("bucket", bucket), zap.String("key", fnm))
o.reconnect(ctx)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return nil, err
}
continue
}
@@ -221,17 +235,15 @@ func (o *OSSStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error)
}
// Remove removes an object from OSS
func (o *OSSStorage) Remove(bucket, fnm string, tenantID ...string) error {
func (o *OSSStorage) Remove(ctx context.Context, bucket, fnm string, tenantID ...string) error {
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
_, err := o.client.DeleteObject(ctx, &s3.DeleteObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(fnm),
})
if err != nil {
zap.L().Error("Failed to remove object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
common.Error("Failed to remove object", err, zap.String("bucket", bucket), zap.String("key", fnm))
return err
}
@@ -239,11 +251,9 @@ func (o *OSSStorage) Remove(bucket, fnm string, tenantID ...string) error {
}
// ObjExist checks if an object exists in OSS
func (o *OSSStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
func (o *OSSStorage) ObjExist(ctx context.Context, bucket, fnm string, tenantID ...string) bool {
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
_, err := o.client.HeadObject(ctx, &s3.HeadObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(fnm),
@@ -258,8 +268,7 @@ func (o *OSSStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
return true
}
func (o *OSSStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
ctx := context.Background()
func (o *OSSStorage) ListObjects(ctx context.Context, bucket string, tenantID ...string) ([]string, error) {
listInput := &s3.ListObjectsV2Input{
Bucket: aws.String(bucket),
@@ -279,11 +288,9 @@ func (o *OSSStorage) ListObjects(bucket string, tenantID ...string) ([]string, e
return objects, nil
}
func (o *OSSStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
func (o *OSSStorage) GetPresignedURL(ctx context.Context, bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
presignClient := s3.NewPresignClient(o.client)
for i := 0; i < 10; i++ {
@@ -292,8 +299,8 @@ func (o *OSSStorage) GetPresignedURL(bucket, fnm string, expires time.Duration,
Key: aws.String(fnm),
}, s3.WithPresignExpires(expires))
if err != nil {
zap.L().Error("Failed to generate presigned URL", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
o.reconnect()
common.Error("Failed to generate presigned URL", err, zap.String("bucket", bucket), zap.String("key", fnm))
o.reconnect(ctx)
time.Sleep(time.Second)
continue
}
@@ -305,19 +312,17 @@ func (o *OSSStorage) GetPresignedURL(bucket, fnm string, expires time.Duration,
}
// BucketExists checks if a bucket exists
func (o *OSSStorage) BucketExists(bucket string) bool {
func (o *OSSStorage) BucketExists(ctx context.Context, bucket string) bool {
actualBucket := bucket
if o.bucket != "" {
actualBucket = o.bucket
}
ctx := context.Background()
_, err := o.client.HeadBucket(ctx, &s3.HeadBucketInput{
Bucket: aws.String(actualBucket),
})
if err != nil {
zap.L().Debug("Bucket does not exist or error", zap.String("bucket", actualBucket), zap.Error(err))
common.Error("Bucket does not exist or error", err, zap.String("bucket", actualBucket))
return false
}
@@ -325,16 +330,14 @@ func (o *OSSStorage) BucketExists(bucket string) bool {
}
// RemoveBucket removes a bucket and all its objects
func (o *OSSStorage) RemoveBucket(bucket string) error {
func (o *OSSStorage) RemoveBucket(ctx context.Context, bucket string) error {
actualBucket := bucket
if o.bucket != "" {
actualBucket = o.bucket
}
ctx := context.Background()
// Check if bucket exists
if !o.BucketExists(actualBucket) {
if !o.BucketExists(ctx, actualBucket) {
return nil
}
@@ -346,17 +349,17 @@ func (o *OSSStorage) RemoveBucket(bucket string) error {
for {
result, err := o.client.ListObjectsV2(ctx, listInput)
if err != nil {
zap.L().Error("Failed to list objects", zap.String("bucket", actualBucket), zap.Error(err))
common.Error("Failed to list objects", err, zap.String("bucket", actualBucket))
return err
}
for _, obj := range result.Contents {
_, err := o.client.DeleteObject(ctx, &s3.DeleteObjectInput{
_, err = o.client.DeleteObject(ctx, &s3.DeleteObjectInput{
Bucket: aws.String(actualBucket),
Key: obj.Key,
})
if err != nil {
zap.L().Error("Failed to delete object", zap.String("bucket", actualBucket), zap.Error(err))
common.Error("Failed to delete object", err, zap.String("bucket", actualBucket))
}
}
@@ -371,7 +374,7 @@ func (o *OSSStorage) RemoveBucket(bucket string) error {
Bucket: aws.String(actualBucket),
})
if err != nil {
zap.L().Error("Failed to delete bucket", zap.String("bucket", actualBucket), zap.Error(err))
common.Error("Failed to delete bucket", err, zap.String("bucket", actualBucket))
return err
}
@@ -379,12 +382,10 @@ func (o *OSSStorage) RemoveBucket(bucket string) error {
}
// Copy copies an object from source to destination
func (o *OSSStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
func (o *OSSStorage) Copy(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
srcBucket, srcPath = o.resolveBucketAndPath(srcBucket, srcPath)
destBucket, destPath = o.resolveBucketAndPath(destBucket, destPath)
ctx := context.Background()
copySource := fmt.Sprintf("%s/%s", srcBucket, srcPath)
_, err := o.client.CopyObject(ctx, &s3.CopyObjectInput{
@@ -393,7 +394,7 @@ func (o *OSSStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool
CopySource: aws.String(copySource),
})
if err != nil {
zap.L().Error("Failed to copy object", zap.String("src", copySource), zap.String("dest", fmt.Sprintf("%s/%s", destBucket, destPath)), zap.Error(err))
common.Error("Failed to copy object", err, zap.String("src", copySource), zap.String("dest", fmt.Sprintf("%s/%s", destBucket, destPath)))
return false
}
@@ -401,10 +402,16 @@ func (o *OSSStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool
}
// Move moves an object from source to destination
func (o *OSSStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
if o.Copy(srcBucket, srcPath, destBucket, destPath) {
if err := o.Remove(srcBucket, srcPath); err != nil {
zap.L().Error("Failed to remove source object after copy", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
func (o *OSSStorage) Move(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
if o.Copy(ctx, srcBucket, srcPath, destBucket, destPath) {
if err := o.Remove(ctx, srcBucket, srcPath); err != nil {
common.Error("Failed to remove source object after copy", err, zap.String("bucket", srcBucket), zap.String("key", srcPath))
rollbackCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
defer cancel()
err = o.Remove(rollbackCtx, destBucket, destPath)
if err != nil {
common.Warn("Failed to roll back copied destination object", zap.String("bucket", destBucket), zap.String("key", destPath), zap.Error(err))
}
return false
}
return true

129
internal/storage/s3.go
View File

@@ -21,6 +21,7 @@ import (
"context"
"errors"
"fmt"
"ragflow/internal/common"
"ragflow/internal/server/config"
"time"
@@ -41,20 +42,19 @@ type S3Storage struct {
}
// NewS3Storage creates a new S3 storage instance
func NewS3Storage(config config.S3Config) (*S3Storage, error) {
func NewS3Storage(ctx context.Context, config config.S3Config) (*S3Storage, error) {
storage := &S3Storage{
config: config,
}
if err := storage.connect(); err != nil {
if err := storage.connect(ctx); err != nil {
return nil, err
}
return storage, nil
}
func (s *S3Storage) connect() error {
ctx := context.Background()
func (s *S3Storage) connect(ctx context.Context) error {
var opts []func(*s3Config.LoadOptions) error
@@ -91,9 +91,9 @@ func (s *S3Storage) connect() error {
return nil
}
func (s *S3Storage) reconnect() {
if err := s.connect(); err != nil {
zap.L().Error("Failed to reconnect to S3", zap.Error(err))
func (s *S3Storage) reconnect(ctx context.Context) {
if err := s.connect(ctx); err != nil {
common.Error("Failed to reconnect to S3", err, zap.Error(err))
}
}
@@ -114,7 +114,7 @@ func (s *S3Storage) resolveBucketAndPath(bucket, fnm string) (string, string) {
func (s *S3Storage) Type() string { return "s3" }
// Health checks S3 service availability
func (s *S3Storage) Health() bool {
func (s *S3Storage) Health(ctx context.Context) bool {
bucket := s.bucket
if bucket == "" {
bucket = "health-check-bucket"
@@ -126,15 +126,13 @@ func (s *S3Storage) Health() bool {
}
binary := []byte("_t@@@1")
ctx := context.Background()
// Ensure bucket exists
if !s.BucketExists(bucket) {
if !s.BucketExists(ctx, bucket) {
_, err := s.client.CreateBucket(ctx, &s3.CreateBucketInput{
Bucket: aws.String(bucket),
})
if err != nil {
zap.L().Error("Failed to create bucket for health check", zap.String("bucket", bucket), zap.Error(err))
common.Error("Failed to create bucket for health check", err, zap.String("bucket", bucket), zap.Error(err))
return false
}
}
@@ -148,7 +146,7 @@ func (s *S3Storage) Health() bool {
})
if err != nil {
zap.L().Error("Health check failed", zap.Error(err))
common.Error("Health check failed", err, zap.Error(err))
return false
}
@@ -156,24 +154,27 @@ func (s *S3Storage) Health() bool {
}
// Put uploads an object to S3
func (s *S3Storage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
func (s *S3Storage) Put(ctx context.Context, bucket, fnm string, binary []byte, tenantID ...string) error {
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
for i := 0; i < 2; i++ {
// Ensure bucket exists
if !s.BucketExists(bucket) {
if !s.BucketExists(ctx, bucket) {
_, err := s.client.CreateBucket(ctx, &s3.CreateBucketInput{
Bucket: aws.String(bucket),
})
if err != nil {
zap.L().Error("Failed to create bucket", zap.String("bucket", bucket), zap.Error(err))
s.reconnect()
time.Sleep(time.Second)
if ctxErr := ctx.Err(); ctxErr != nil {
return ctxErr
}
common.Error("Failed to create bucket", err, zap.String("bucket", bucket), zap.Error(err))
s.reconnect(ctx)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return err
}
continue
}
zap.L().Info("Created bucket", zap.String("bucket", bucket))
common.Info("Created bucket", zap.String("bucket", bucket))
}
reader := bytes.NewReader(binary)
@@ -183,9 +184,14 @@ func (s *S3Storage) Put(bucket, fnm string, binary []byte, tenantID ...string) e
Body: reader,
})
if err != nil {
zap.L().Error("Failed to put object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
s.reconnect()
time.Sleep(time.Second)
if ctxErr := ctx.Err(); ctxErr != nil {
return ctxErr
}
common.Error("Failed to put object", err, zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
s.reconnect(ctx)
if err = sleepOrAbort(ctx, time.Second); err != nil {
return err
}
continue
}
@@ -196,28 +202,26 @@ func (s *S3Storage) Put(bucket, fnm string, binary []byte, tenantID ...string) e
}
// Get retrieves an object from S3
func (s *S3Storage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
func (s *S3Storage) Get(ctx context.Context, bucket, fnm string, tenantID ...string) ([]byte, error) {
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
for i := 0; i < 2; i++ {
result, err := s.client.GetObject(ctx, &s3.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(fnm),
})
if err != nil {
zap.L().Error("Failed to get object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
s.reconnect()
common.Error("Failed to get object", err, zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
s.reconnect(ctx)
time.Sleep(time.Second)
continue
}
defer result.Body.Close()
buf := new(bytes.Buffer)
if _, err := buf.ReadFrom(result.Body); err != nil {
zap.L().Error("Failed to read object data", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
s.reconnect()
if _, err = buf.ReadFrom(result.Body); err != nil {
common.Error("Failed to read object data", err, zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
s.reconnect(ctx)
time.Sleep(time.Second)
continue
}
@@ -229,17 +233,15 @@ func (s *S3Storage) Get(bucket, fnm string, tenantID ...string) ([]byte, error)
}
// Remove removes an object from S3
func (s *S3Storage) Remove(bucket, fnm string, tenantID ...string) error {
func (s *S3Storage) Remove(ctx context.Context, bucket, fnm string, tenantID ...string) error {
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(fnm),
})
if err != nil {
zap.L().Error("Failed to remove object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
common.Error("Failed to remove object", err, zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
return err
}
@@ -247,11 +249,9 @@ func (s *S3Storage) Remove(bucket, fnm string, tenantID ...string) error {
}
// ObjExist checks if an object exists in S3
func (s *S3Storage) ObjExist(bucket, fnm string, tenantID ...string) bool {
func (s *S3Storage) ObjExist(ctx context.Context, bucket, fnm string, tenantID ...string) bool {
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
_, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(fnm),
@@ -266,8 +266,7 @@ func (s *S3Storage) ObjExist(bucket, fnm string, tenantID ...string) bool {
return true
}
func (s *S3Storage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
ctx := context.Background()
func (s *S3Storage) ListObjects(ctx context.Context, bucket string, tenantID ...string) ([]string, error) {
listInput := &s3.ListObjectsV2Input{
Bucket: aws.String(bucket),
@@ -288,11 +287,9 @@ func (s *S3Storage) ListObjects(bucket string, tenantID ...string) ([]string, er
}
// GetPresignedURL generates a presigned URL for accessing an object
func (s *S3Storage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
func (s *S3Storage) GetPresignedURL(ctx context.Context, bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
ctx := context.Background()
presignClient := s3.NewPresignClient(s.client)
for i := 0; i < 10; i++ {
@@ -301,8 +298,8 @@ func (s *S3Storage) GetPresignedURL(bucket, fnm string, expires time.Duration, t
Key: aws.String(fnm),
}, s3.WithPresignExpires(expires))
if err != nil {
zap.L().Error("Failed to generate presigned URL", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
s.reconnect()
common.Error("Failed to generate presigned URL", err, zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
s.reconnect(ctx)
time.Sleep(time.Second)
continue
}
@@ -314,19 +311,17 @@ func (s *S3Storage) GetPresignedURL(bucket, fnm string, expires time.Duration, t
}
// BucketExists checks if a bucket exists
func (s *S3Storage) BucketExists(bucket string) bool {
func (s *S3Storage) BucketExists(ctx context.Context, bucket string) bool {
actualBucket := bucket
if s.bucket != "" {
actualBucket = s.bucket
}
ctx := context.Background()
_, err := s.client.HeadBucket(ctx, &s3.HeadBucketInput{
Bucket: aws.String(actualBucket),
})
if err != nil {
zap.L().Debug("Bucket does not exist or error", zap.String("bucket", actualBucket), zap.Error(err))
common.Debug("Bucket does not exist or error", zap.String("bucket", actualBucket), zap.Error(err))
return false
}
@@ -334,16 +329,14 @@ func (s *S3Storage) BucketExists(bucket string) bool {
}
// RemoveBucket removes a bucket and all its objects
func (s *S3Storage) RemoveBucket(bucket string) error {
func (s *S3Storage) RemoveBucket(ctx context.Context, bucket string) error {
actualBucket := bucket
if s.bucket != "" {
actualBucket = s.bucket
}
ctx := context.Background()
// Check if bucket exists
if !s.BucketExists(actualBucket) {
if !s.BucketExists(ctx, actualBucket) {
return nil
}
@@ -355,17 +348,17 @@ func (s *S3Storage) RemoveBucket(bucket string) error {
for {
result, err := s.client.ListObjectsV2(ctx, listInput)
if err != nil {
zap.L().Error("Failed to list objects", zap.String("bucket", actualBucket), zap.Error(err))
common.Error("Failed to list objects", err, zap.String("bucket", actualBucket), zap.Error(err))
return err
}
for _, obj := range result.Contents {
_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
_, err = s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
Bucket: aws.String(actualBucket),
Key: obj.Key,
})
if err != nil {
zap.L().Error("Failed to delete object", zap.String("bucket", actualBucket), zap.Error(err))
common.Error("Failed to delete object", err, zap.String("bucket", actualBucket), zap.Error(err))
}
}
@@ -380,7 +373,7 @@ func (s *S3Storage) RemoveBucket(bucket string) error {
Bucket: aws.String(actualBucket),
})
if err != nil {
zap.L().Error("Failed to delete bucket", zap.String("bucket", actualBucket), zap.Error(err))
common.Error("Failed to delete bucket", err, zap.String("bucket", actualBucket), zap.Error(err))
return err
}
@@ -388,12 +381,10 @@ func (s *S3Storage) RemoveBucket(bucket string) error {
}
// Copy copies an object from source to destination
func (s *S3Storage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
func (s *S3Storage) Copy(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
srcBucket, srcPath = s.resolveBucketAndPath(srcBucket, srcPath)
destBucket, destPath = s.resolveBucketAndPath(destBucket, destPath)
ctx := context.Background()
copySource := fmt.Sprintf("%s/%s", srcBucket, srcPath)
_, err := s.client.CopyObject(ctx, &s3.CopyObjectInput{
@@ -402,7 +393,7 @@ func (s *S3Storage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
CopySource: aws.String(copySource),
})
if err != nil {
zap.L().Error("Failed to copy object", zap.String("src", copySource), zap.String("dest", fmt.Sprintf("%s/%s", destBucket, destPath)), zap.Error(err))
common.Error("Failed to copy object", err, zap.String("src", copySource), zap.String("dest", fmt.Sprintf("%s/%s", destBucket, destPath)), zap.Error(err))
return false
}
@@ -410,10 +401,16 @@ func (s *S3Storage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
}
// Move moves an object from source to destination
func (s *S3Storage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
if s.Copy(srcBucket, srcPath, destBucket, destPath) {
if err := s.Remove(srcBucket, srcPath); err != nil {
zap.L().Error("Failed to remove source object after copy", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
func (s *S3Storage) Move(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool {
if s.Copy(ctx, srcBucket, srcPath, destBucket, destPath) {
if err := s.Remove(ctx, srcBucket, srcPath); err != nil {
common.Error("Failed to remove source object after copy", err, zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
rollbackCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
defer cancel()
err = s.Remove(rollbackCtx, destBucket, destPath)
if err != nil {
common.Warn("Failed to roll back copied destination object", zap.String("bucket", destBucket), zap.String("key", destPath), zap.Error(err))
}
return false
}
return true

69
internal/storage/storage_factory.go
View File

@@ -17,10 +17,12 @@
package storage
import (
"context"
"fmt"
"ragflow/internal/common"
"ragflow/internal/server"
"sync"
"time"
)
var (
@@ -43,12 +45,12 @@ func GetStorageFactory() *StorageFactory {
}
// Init initializes the storage factory with configuration
func Init() error {
func Init(ctx context.Context) error {
factory := GetStorageFactory()
globalConfig := server.GetConfig()
// Initialize storage based on type
if err := factory.initStorage(); err != nil {
if err := factory.initStorage(ctx); err != nil {
return err
}
@@ -65,17 +67,17 @@ func CloseStorage() error {
return factory.storage.Close()
}
func (f *StorageFactory) initStorage() error {
func (f *StorageFactory) initStorage(ctx context.Context) error {
globalConfig := server.GetConfig()
switch globalConfig.StorageEngineType() {
case "minio":
return f.initMinio()
case "s3":
return f.initS3()
return f.initS3(ctx)
case "oss":
return f.initOSS()
return f.initOSS(ctx)
case "gcs":
return f.initGCS()
return f.initGCS(ctx)
default:
return fmt.Errorf("unsupported storage type: %s", globalConfig.StorageEngineType())
}
@@ -95,9 +97,9 @@ func (f *StorageFactory) initMinio() error {
return nil
}
func (f *StorageFactory) initS3() error {
func (f *StorageFactory) initS3(ctx context.Context) error {
globalConfig := server.GetConfig()
storage, err := NewS3Storage(globalConfig.GetS3Config())
storage, err := NewS3Storage(ctx, globalConfig.GetS3Config())
if err != nil {
return fmt.Errorf("failed to create S3 storage: %w", err)
}
@@ -109,9 +111,9 @@ func (f *StorageFactory) initS3() error {
return nil
}
func (f *StorageFactory) initOSS() error {
func (f *StorageFactory) initOSS(ctx context.Context) error {
globalConfig := server.GetConfig()
storage, err := NewOSSStorage(globalConfig.GetOSSConfig())
storage, err := NewOSSStorage(ctx, globalConfig.GetOSSConfig())
if err != nil {
return fmt.Errorf("failed to create OSS storage: %w", err)
}
@@ -123,9 +125,9 @@ func (f *StorageFactory) initOSS() error {
return nil
}
func (f *StorageFactory) initGCS() error {
func (f *StorageFactory) initGCS(ctx context.Context) error {
globalConfig := server.GetConfig()
storage, err := NewGCSStorage(globalConfig.GetGCSConfig())
storage, err := NewGCSStorage(ctx, globalConfig.GetGCSConfig())
if err != nil {
return fmt.Errorf("failed to create GCS storage: %w", err)
}
@@ -144,38 +146,23 @@ func (f *StorageFactory) GetStorage() Storage {
return f.storage
}
// Create creates a new storage instance based on the storage type
// This is the factory method equivalent to Python's StorageFactory.create()
//func (f *StorageFactory) Create(storageType StorageType) (Storage, error) {
// var storage Storage
// var err error
//
// switch storageType {
// case StorageMinio:
// storage, err = NewMinioStorage(f.config.Minio)
// if err != nil {
// return nil, fmt.Errorf("MinIO config not available: %w, %v", err, f.config.Minio)
// }
// case StorageAWSS3:
// storage, err = NewS3Storage(f.config.S3)
// if err != nil {
// return nil, fmt.Errorf("S3 config not available: %w, %v", err, f.config.S3)
// }
// case StorageOSS:
// storage, err = NewOSSStorage(f.config.OSS)
// if err != nil {
// return nil, fmt.Errorf("OSS config not available: %w, %v", err, f.config.OSS)
// }
// default:
// return nil, fmt.Errorf("unsupported storage type: %v", storageType)
// }
//
// return storage, nil
//}
// SetStorage sets the storage instance (useful for testing)
func (f *StorageFactory) SetStorage(storage Storage) {
f.mu.Lock()
defer f.mu.Unlock()
f.storage = storage
}
func sleepOrAbort(ctx context.Context, d time.Duration) error {
if err := ctx.Err(); err != nil {
return err
}
timer := time.NewTimer(d)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
}

23
internal/storage/types.go
View File

@@ -17,6 +17,7 @@
package storage
import (
"context"
"time"
)
@@ -59,43 +60,43 @@ type Storage interface {
Type() string
// Health checks the storage service availability
Health() bool
Health(ctx context.Context) bool
// Put uploads an object to storage
// bucket: the bucket/container name
// fnm: the file/object name (key)
// binary: the data to upload
// tenantID: optional tenant identifier
Put(bucket, fnm string, binary []byte, tenantID ...string) error
Put(ctx context.Context, bucket, fnm string, binary []byte, tenantID ...string) error
// Get retrieves an object from storage
// Returns the data or nil if not found
Get(bucket, fnm string, tenantID ...string) ([]byte, error)
Get(ctx context.Context, bucket, fnm string, tenantID ...string) ([]byte, error)
// Remove removes an object from storage
Remove(bucket, fnm string, tenantID ...string) error
Remove(ctx context.Context, bucket, fnm string, tenantID ...string) error
// ObjExist checks if an object exists
ObjExist(bucket, fnm string, tenantID ...string) bool
ObjExist(ctx context.Context, bucket, fnm string, tenantID ...string) bool
// ListObjects list all objects of the bucket
ListObjects(bucket string, tenantID ...string) ([]string, error)
ListObjects(ctx context.Context, bucket string, tenantID ...string) ([]string, error)
// GetPresignedURL generates a presigned URL for accessing an object
// expires: duration until the URL expires
GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error)
GetPresignedURL(ctx context.Context, bucket, fnm string, expires time.Duration, tenantID ...string) (string, error)
// BucketExists checks if a bucket exists
BucketExists(bucket string) bool
BucketExists(ctx context.Context, bucket string) bool
// RemoveBucket removes a bucket and all its objects
RemoveBucket(bucket string) error
RemoveBucket(ctx context.Context, bucket string) error
// Copy copies an object from source to destination
Copy(srcBucket, srcPath, destBucket, destPath string) bool
Copy(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool
// Move moves an object from source to destination
Move(srcBucket, srcPath, destBucket, destPath string) bool
Move(ctx context.Context, srcBucket, srcPath, destBucket, destPath string) bool
// Close closes the storage connection
Close() error