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

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 {