public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

unistore: save returns a writecloser (#107955)

Browse Source 
* unistore: save returns a writecloser

* go-lint

* address comments
This commit is contained in:
Georges Chaudy
2025-07-11 18:25:48 +02:00
committed by GitHub
parent 519e0449e7
commit 7002ab90ae
6 changed files with 149 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/storage/unified/resource/datastore.go
+11 -1
View File
@@ -247,7 +247,17 @@ func (d *dataStore) Save(ctx context.Context, key DataKey, value io.Reader) erro
return fmt.Errorf("invalid data key: %w", err)
}
return d.kv.Save(ctx, dataSection, key.String(), value)
writer, err := d.kv.Save(ctx, dataSection, key.String())
if err != nil {
return err
}
_, err = io.Copy(writer, value)
if err != nil {
_ = writer.Close()
return err
}
return writer.Close()
}
func (d *dataStore) Delete(ctx context.Context, key DataKey) error {
pkg/storage/unified/resource/eventstore.go
+9 -5
View File
@@ -1,7 +1,6 @@
package resource
import (
"bytes"
"context"
"encoding/json"
"fmt"
@@ -134,12 +133,17 @@ func (n *eventStore) Save(ctx context.Context, event Event) error {
return fmt.Errorf("invalid event key: %w", err)
}
var buf bytes.Buffer
encoder := json.NewEncoder(&buf)
if err := encoder.Encode(event); err != nil {
writer, err := n.kv.Save(ctx, eventsSection, eventKey.String())
if err != nil {
return err
}
return n.kv.Save(ctx, eventsSection, eventKey.String(), &buf)
encoder := json.NewEncoder(writer)
if err := encoder.Encode(event); err != nil {
_ = writer.Close()
return err
}
return writer.Close()
}
func (n *eventStore) Get(ctx context.Context, key EventKey) (Event, error) {
pkg/storage/unified/resource/kv.go
+50 -16
View File
@@ -36,8 +36,8 @@ type KV interface {
// Get retrieves the value for a key from the store
Get(ctx context.Context, section string, key string) (io.ReadCloser, error)
// Save a new value
Save(ctx context.Context, section string, key string, value io.Reader) error
// Save a new value - returns a WriteCloser to write the value to
Save(ctx context.Context, section string, key string) (io.WriteCloser, error)
// Delete a value
Delete(ctx context.Context, section string, key string) error
@@ -92,32 +92,66 @@ func (k *badgerKV) Get(ctx context.Context, section string, key string) (io.Read
return io.NopCloser(bytes.NewReader(value)), nil
}
func (k *badgerKV) Save(ctx context.Context, section string, key string, value io.Reader) error {
if k.db.IsClosed() {
// badgerWriteCloser implements io.WriteCloser for badgerKV
type badgerWriteCloser struct {
db *badger.DB
keyWithSection string
buf *bytes.Buffer
closed bool
}
// Write implements io.Writer
func (w *badgerWriteCloser) Write(p []byte) (int, error) {
if w.closed {
return 0, fmt.Errorf("write to closed writer")
}
return w.buf.Write(p)
}
// Close implements io.Closer - stores the buffered data in BadgerDB
func (w *badgerWriteCloser) Close() error {
if w.closed {
return nil
}
w.closed = true
if w.db.IsClosed() {
return fmt.Errorf("database is closed")
}
if section == "" {
return fmt.Errorf("section is required")
}
data := w.buf.Bytes()
key = section + "/" + key
data, err := io.ReadAll(value)
if err != nil {
return fmt.Errorf("failed to read value: %w", err)
}
txn := k.db.NewTransaction(true)
txn := w.db.NewTransaction(true)
defer txn.Discard()
err = txn.Set([]byte(key), data)
err := txn.Set([]byte(w.keyWithSection), data)
if err != nil {
return err
}
return txn.Commit()
}
func (k *badgerKV) Save(ctx context.Context, section string, key string) (io.WriteCloser, error) {
if k.db.IsClosed() {
return nil, fmt.Errorf("database is closed")
}
if section == "" {
return nil, fmt.Errorf("section is required")
}
if key == "" {
return nil, fmt.Errorf("key is required")
}
return &badgerWriteCloser{
db: k.db,
keyWithSection: section + "/" + key,
buf: &bytes.Buffer{},
closed: false,
}, nil
}
func (k *badgerKV) Delete(ctx context.Context, section string, key string) error {
if k.db.IsClosed() {
return fmt.Errorf("database is closed")
pkg/storage/unified/resource/kv_test.go
+20 -15
View File
@@ -64,11 +64,10 @@ func TestBadgerKV_UnderlyingStorage(t *testing.T) {
expectedInternalKey := section + "/" + key
// Save through KV interface
err := kv.Save(ctx, section, key, strings.NewReader(value))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, key, strings.NewReader(value))
// Verify the raw key exists in badger with correct format
err = db.View(func(txn *badger.Txn) error {
err := db.View(func(txn *badger.Txn) error {
item, err := txn.Get([]byte(expectedInternalKey))
require.NoError(t, err)
@@ -90,13 +89,11 @@ func TestBadgerKV_UnderlyingStorage(t *testing.T) {
value2 := "value-from-section2"
// Save same key in different sections
err := kv.Save(ctx, section1, key, strings.NewReader(value1))
require.NoError(t, err)
err = kv.Save(ctx, section2, key, strings.NewReader(value2))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section1, key, strings.NewReader(value1))
saveKVHelper(t, kv, ctx, section2, key, strings.NewReader(value2))
// Verify both keys exist in badger with different internal keys
err = db.View(func(txn *badger.Txn) error {
err := db.View(func(txn *badger.Txn) error {
// Check section1 key
item1, err := txn.Get([]byte(section1 + "/" + key))
require.NoError(t, err)
@@ -140,11 +137,10 @@ func TestBadgerKV_UnderlyingStorage(t *testing.T) {
internalKey := section + "/" + key
// Save and verify it exists
err := kv.Save(ctx, section, key, strings.NewReader(value))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, key, strings.NewReader(value))
// Verify it exists in badger
err = db.View(func(txn *badger.Txn) error {
err := db.View(func(txn *badger.Txn) error {
_, err := txn.Get([]byte(internalKey))
return err
})
@@ -172,12 +168,10 @@ func TestBadgerKV_UnderlyingStorage(t *testing.T) {
keys2 := []string{"b1", "b2", "b3"}
for _, k := range keys1 {
err := kv.Save(ctx, section1, k, strings.NewReader("value"+k))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section1, k, strings.NewReader("value"+k))
}
for _, k := range keys2 {
err := kv.Save(ctx, section2, k, strings.NewReader("value"+k))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section2, k, strings.NewReader("value"+k))
}
// List keys from section1 only
@@ -262,3 +256,14 @@ func TestIsValidKey(t *testing.T) {
})
}
}
// saveKVHelper is a helper function to save data to KV store using the new WriteCloser interface
func saveKVHelper(t *testing.T, kv KV, ctx context.Context, section, key string, value io.Reader) {
t.Helper()
writer, err := kv.Save(ctx, section, key)
require.NoError(t, err)
_, err = io.Copy(writer, value)
require.NoError(t, err)
err = writer.Close()
require.NoError(t, err)
}
pkg/storage/unified/resource/metadata.go
+10 -5
View File
@@ -1,7 +1,6 @@
package resource
import (
"bytes"
"context"
"encoding/json"
"fmt"
@@ -340,12 +339,18 @@ func (d *metadataStore) Save(ctx context.Context, obj MetaDataObj) error {
if err := obj.Key.Validate(); err != nil {
return fmt.Errorf("invalid metadata key: %w", err)
}
var buf bytes.Buffer
encoder := json.NewEncoder(&buf)
if err := encoder.Encode(obj.Value); err != nil {
writer, err := d.kv.Save(ctx, metaSection, obj.Key.String())
if err != nil {
return err
}
return d.kv.Save(ctx, metaSection, obj.Key.String(), &buf)
encoder := json.NewEncoder(writer)
if err := encoder.Encode(obj.Value); err != nil {
_ = writer.Close()
return err
}
return writer.Close()
}
// parseMetaDataKey parses a string key into a MetaDataKey struct
pkg/storage/unified/testing/kv.go
+49 -24
View File
@@ -81,8 +81,7 @@ func runTestKVGet(t *testing.T, kv resource.KV, nsPrefix string) {
t.Run("get existing key", func(t *testing.T) {
// First save a key
testValue := "test value for get"
err := kv.Save(ctx, section, "existing-key", strings.NewReader(testValue))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, "existing-key", strings.NewReader(testValue))
// Now get it
reader, err := kv.Get(ctx, section, "existing-key")
@@ -117,8 +116,7 @@ func runTestKVSave(t *testing.T, kv resource.KV, nsPrefix string) {
t.Run("save new key", func(t *testing.T) {
testValue := "new test value"
err := kv.Save(ctx, section, "new-key", strings.NewReader(testValue))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, "new-key", strings.NewReader(testValue))
// Verify it was saved
reader, err := kv.Get(ctx, section, "new-key")
@@ -133,13 +131,11 @@ func runTestKVSave(t *testing.T, kv resource.KV, nsPrefix string) {
t.Run("save overwrite existing key", func(t *testing.T) {
// First save
err := kv.Save(ctx, section, "overwrite-key", strings.NewReader("old value"))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, "overwrite-key", strings.NewReader("old value"))
// Overwrite
newValue := "new value"
err = kv.Save(ctx, section, "overwrite-key", strings.NewReader(newValue))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, "overwrite-key", strings.NewReader(newValue))
// Verify it was updated
reader, err := kv.Get(ctx, section, "overwrite-key")
@@ -153,15 +149,14 @@ func runTestKVSave(t *testing.T, kv resource.KV, nsPrefix string) {
})
t.Run("save with empty section", func(t *testing.T) {
err := kv.Save(ctx, "", "some-key", strings.NewReader("some value"))
_, err := kv.Save(ctx, "", "some-key")
assert.Error(t, err)
assert.Contains(t, err.Error(), "section is required")
})
t.Run("save binary data", func(t *testing.T) {
binaryData := []byte{0x00, 0x01, 0x02, 0x03, 0xFF, 0xFE, 0xFD}
err := kv.Save(ctx, section, "binary-key", bytes.NewReader(binaryData))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, "binary-key", bytes.NewReader(binaryData))
// Verify binary data
reader, err := kv.Get(ctx, section, "binary-key")
@@ -176,8 +171,7 @@ func runTestKVSave(t *testing.T, kv resource.KV, nsPrefix string) {
t.Run("save key with no data", func(t *testing.T) {
// Save a key with empty data
err := kv.Save(ctx, section, "empty-key", strings.NewReader(""))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, "empty-key", strings.NewReader(""))
// Verify it was saved with empty data
reader, err := kv.Get(ctx, section, "empty-key")
@@ -198,11 +192,10 @@ func runTestKVDelete(t *testing.T, kv resource.KV, nsPrefix string) {
t.Run("delete existing key", func(t *testing.T) {
// First create a key
err := kv.Save(ctx, section, "delete-key", strings.NewReader("delete me"))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, "delete-key", strings.NewReader("delete me"))
// Verify it exists
_, err = kv.Get(ctx, section, "delete-key")
_, err := kv.Get(ctx, section, "delete-key")
require.NoError(t, err)
// Delete it
@@ -235,8 +228,7 @@ func runTestKVKeys(t *testing.T, kv resource.KV, nsPrefix string) {
// Setup test data
testKeys := []string{"a1", "a2", "b1", "b2", "c1"}
for _, key := range testKeys {
err := kv.Save(ctx, section, key, strings.NewReader("value"+key))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, key, strings.NewReader("value"+key))
}
t.Run("list all keys", func(t *testing.T) {
@@ -284,8 +276,7 @@ func runTestKVKeysWithLimits(t *testing.T, kv resource.KV, nsPrefix string) {
// Setup test data
testKeys := []string{"a1", "a2", "b1", "b2", "c1", "c2", "d1", "d2"}
for _, key := range testKeys {
err := kv.Save(ctx, section, key, strings.NewReader("value"+key))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, key, strings.NewReader("value"+key))
}
t.Run("keys with limit", func(t *testing.T) {
@@ -339,8 +330,7 @@ func runTestKVKeysWithSort(t *testing.T, kv resource.KV, nsPrefix string) {
// Setup test data
testKeys := []string{"a1", "a2", "b1", "b2", "c1"}
for _, key := range testKeys {
err := kv.Save(ctx, section, key, strings.NewReader("value"+key))
require.NoError(t, err)
saveKVHelper(t, kv, ctx, section, key, strings.NewReader("value"+key))
}
t.Run("keys in ascending order (default)", func(t *testing.T) {
@@ -407,7 +397,19 @@ func runTestKVConcurrent(t *testing.T, kv resource.KV, nsPrefix string) {
value := fmt.Sprintf("concurrent-value-%d-%d", goroutineID, j)
// Save
err = kv.Save(ctx, section, key, strings.NewReader(value))
writer, err := kv.Save(ctx, section, key)
if err != nil {
return
}
defer func() {
err := writer.Close()
require.NoError(t, err)
}()
_, err = io.Copy(writer, strings.NewReader(value))
if err != nil {
return
}
err = writer.Close()
if err != nil {
return
}
@@ -447,7 +449,19 @@ func runTestKVConcurrent(t *testing.T, kv resource.KV, nsPrefix string) {
value := fmt.Sprintf("concurrent-ops-value-%d", goroutineID)
// Save
err = kv.Save(ctx, section, key, strings.NewReader(value))
writer, err := kv.Save(ctx, section, key)
if err != nil {
return
}
defer func() {
err := writer.Close()
require.NoError(t, err)
}()
_, err = io.Copy(writer, strings.NewReader(value))
if err != nil {
return
}
err = writer.Close()
if err != nil {
return
}
@@ -512,3 +526,14 @@ func runTestKVUnixTimestamp(t *testing.T, kv resource.KV, nsPrefix string) {
require.InDelta(t, timestamp1, timestamp2, 1)
})
}
// saveKVHelper is a helper function to save data to KV store using the new WriteCloser interface
func saveKVHelper(t *testing.T, kv resource.KV, ctx context.Context, section, key string, value io.Reader) {
t.Helper()
writer, err := kv.Save(ctx, section, key)
require.NoError(t, err)
_, err = io.Copy(writer, value)
require.NoError(t, err)
err = writer.Close()
require.NoError(t, err)
}