Skip to content

Commit

Permalink
cloudflarekv: add support for context management
Browse files Browse the repository at this point in the history
  • Loading branch information
@efectn
efectn committed Feb 11, 2025
1 parent d365f8a commit 9f1466a
Show file tree
Hide file tree
Showing 2 changed files with 110 additions and 68 deletions.
34 changes: 25 additions & 9 deletions cloudflarekv/cloudflarekv.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,8 @@ func New(config ...Config) *Storage {
return storage
}

func (s *Storage) Get(key string) ([]byte, error) {
resp, err := s.api.GetWorkersKV(context.Background(), cloudflare.AccountIdentifier(s.accountID), cloudflare.GetWorkersKVParams{NamespaceID: s.namespaceID, Key: key})
func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
resp, err := s.api.GetWorkersKV(ctx, cloudflare.AccountIdentifier(s.accountID), cloudflare.GetWorkersKVParams{NamespaceID: s.namespaceID, Key: key})

if err != nil {
log.Printf("Error occur in GetWorkersKV: %v", err)
Expand All @@ -66,8 +66,12 @@ func (s *Storage) Get(key string) ([]byte, error) {
return resp, nil
}

func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
_, err := s.api.WriteWorkersKVEntry(context.Background(), cloudflare.AccountIdentifier(s.accountID), cloudflare.WriteWorkersKVEntryParams{
func (s *Storage) Get(key string) ([]byte, error) {
return s.GetWithContext(context.Background(), key)
}

func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
_, err := s.api.WriteWorkersKVEntry(ctx, cloudflare.AccountIdentifier(s.accountID), cloudflare.WriteWorkersKVEntryParams{
NamespaceID: s.namespaceID,
Key: key,
Value: val,
Expand All @@ -81,8 +85,12 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
return nil
}

func (s *Storage) Delete(key string) error {
_, err := s.api.DeleteWorkersKVEntry(context.Background(), cloudflare.AccountIdentifier(s.accountID), cloudflare.DeleteWorkersKVEntryParams{
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
return s.SetWithContext(context.Background(), key, val, exp)
}

func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
_, err := s.api.DeleteWorkersKVEntry(ctx, cloudflare.AccountIdentifier(s.accountID), cloudflare.DeleteWorkersKVEntryParams{
NamespaceID: s.namespaceID,
Key: key,
})
Expand All @@ -95,14 +103,18 @@ func (s *Storage) Delete(key string) error {
return nil
}

func (s *Storage) Reset() error {
func (s *Storage) Delete(key string) error {
return s.DeleteWithContext(context.Background(), key)
}

func (s *Storage) ResetWithContext(ctx context.Context) error {
var (
cursor string
keys []string
)

for {
resp, err := s.api.ListWorkersKVKeys(context.Background(), cloudflare.AccountIdentifier(s.accountID), cloudflare.ListWorkersKVsParams{
resp, err := s.api.ListWorkersKVKeys(ctx, cloudflare.AccountIdentifier(s.accountID), cloudflare.ListWorkersKVsParams{
NamespaceID: s.namespaceID,
Cursor: cursor,
})
Expand All @@ -119,7 +131,7 @@ func (s *Storage) Reset() error {
keys = append(keys, name)
}

_, err = s.api.DeleteWorkersKVEntries(context.Background(), cloudflare.AccountIdentifier(s.accountID), cloudflare.DeleteWorkersKVEntriesParams{
_, err = s.api.DeleteWorkersKVEntries(ctx, cloudflare.AccountIdentifier(s.accountID), cloudflare.DeleteWorkersKVEntriesParams{
NamespaceID: s.namespaceID,
Keys: keys,
})
Expand All @@ -140,6 +152,10 @@ func (s *Storage) Reset() error {
return nil
}

func (s *Storage) Reset() error {
return s.ResetWithContext(context.Background())
}

func (s *Storage) Close() error {
return nil
}
Expand Down
144 changes: 85 additions & 59 deletions cloudflarekv/cloudflarekv_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,21 @@ package cloudflarekv

import (
"bytes"
"context"
"fmt"
"os"
"testing"

"github.com/stretchr/testify/require"
)

func TestMain(m *testing.M) {
var testStore *Storage

var testStore *Storage
func TestMain(m *testing.M) {

testStore = New(Config{
Key: "test",
Key: "test",
Reset: true,
})

code := m.Run()
Expand All @@ -25,12 +28,6 @@ func TestMain(m *testing.M) {
func Test_CloudflareKV_Get(t *testing.T) {
t.Parallel()

var testStore *Storage

testStore = New(Config{
Key: "test",
})

var (
key = "john"
val = []byte("doe")
Expand All @@ -55,14 +52,30 @@ func Test_CloudflareKV_Get(t *testing.T) {
_ = testStore.Close()
}

func Test_CloudflareKV_Set(t *testing.T) {
func Test_CloudflareKV_GetWithContext(t *testing.T) {
t.Parallel()

var testStore *Storage
var (
key = "john"
val = []byte("doe")
)

testStore = New(Config{
Key: "test",
})
err := testStore.Set(key, val, 0)
require.NoError(t, err)

ctx, cancel := context.WithCancel(context.Background())
cancel()

val, err = testStore.GetWithContext(ctx, key)
fmt.Println(err)
require.ErrorContains(t, err, context.Canceled.Error())
require.Nil(t, val)

_ = testStore.Close()
}

func Test_CloudflareKV_Set(t *testing.T) {
t.Parallel()

var (
key = "john"
Expand All @@ -76,14 +89,21 @@ func Test_CloudflareKV_Set(t *testing.T) {
_ = testStore.Close()
}

func Test_CloudflareKV_Delete(t *testing.T) {
t.Parallel()
func Test_CloudflareKV_SetWithContext(t *testing.T) {
var (
key = "john"
val = []byte("doe")
)

var testStore *Storage
ctx, cancel := context.WithCancel(context.Background())
cancel()

testStore = New(Config{
Key: "test",
})
err := testStore.SetWithContext(ctx, key, val, 0)
require.ErrorIs(t, err, context.Canceled)
}

func Test_CloudflareKV_Delete(t *testing.T) {
t.Parallel()

var (
key = "john"
Expand All @@ -99,29 +119,62 @@ func Test_CloudflareKV_Delete(t *testing.T) {
_ = testStore.Close()
}

func Test_CloudflareKV_Reset(t *testing.T) {
t.Parallel()
func Test_CloudflareKV_DeleteWithContext(t *testing.T) {
var (
key = "john"
val = []byte("doe")
)

var testStore *Storage
err := testStore.Set(key, val, 0)
require.NoError(t, err)

testStore = New(Config{
Key: "test",
})
ctx, cancel := context.WithCancel(context.Background())
cancel()

err = testStore.DeleteWithContext(ctx, key)
require.ErrorIs(t, err, context.Canceled)

result, err := testStore.Get(key)
require.NoError(t, err)
require.Equal(t, val, result)
}

func Test_CloudflareKV_Reset(t *testing.T) {
t.Parallel()

err := testStore.Reset()

require.NoError(t, err)

_ = testStore.Close()
}
func Test_CloudflareKV_Close(t *testing.T) {
t.Parallel()

var testStore *Storage
func Test_CloudflareKV_ResetWithContext(t *testing.T) {
val := []byte("doe")

testStore = New(Config{
Key: "test",
})
err := testStore.Set("john1", val, 0)
require.NoError(t, err)

err = testStore.Set("john2", val, 0)
require.NoError(t, err)

ctx, cancel := context.WithCancel(context.Background())
cancel()

err = testStore.ResetWithContext(ctx)
require.ErrorIs(t, err, context.Canceled)

result, err := testStore.Get("john1")
require.NoError(t, err)
require.Equal(t, val, result)

result, err = testStore.Get("john2")
require.NoError(t, err)
require.Equal(t, val, result)
}

func Test_CloudflareKV_Close(t *testing.T) {
t.Parallel()

require.Nil(t, testStore.Close())

Expand All @@ -131,25 +184,12 @@ func Test_CloudflareKV_Close(t *testing.T) {
func Test_CloudflareKV_Conn(t *testing.T) {
t.Parallel()

var testStore *Storage

testStore = New(Config{
Key: "test",
})

require.NotNil(t, testStore.Conn())

_ = testStore.Close()
}

func Benchmark_CloudflareKV_Set(b *testing.B) {

var testStore *Storage

testStore = New(Config{
Key: "test",
})

b.ReportAllocs()
b.ResetTimer()

Expand All @@ -164,13 +204,6 @@ func Benchmark_CloudflareKV_Set(b *testing.B) {
}

func Benchmark_CloudflareKV_Get(b *testing.B) {

var testStore *Storage

testStore = New(Config{
Key: "test",
})

err := testStore.Set("john", []byte("doe"), 0)
require.NoError(b, err)

Expand All @@ -187,13 +220,6 @@ func Benchmark_CloudflareKV_Get(b *testing.B) {
}

func Benchmark_CloudflareKV_SetAndDelete(b *testing.B) {

var testStore *Storage

testStore = New(Config{
Key: "test",
})

b.ReportAllocs()
b.ResetTimer()

Expand Down

0 comments on commit 9f1466a

Please # to comment.