From 7002ab90ae00631de5bb623b6f3c66fa12f5de84 Mon Sep 17 00:00:00 2001 From: Georges Chaudy Date: Fri, 11 Jul 2025 18:25:48 +0200 Subject: [PATCH] unistore: save returns a writecloser (#107955) * unistore: save returns a writecloser * go-lint * address comments --- pkg/storage/unified/resource/datastore.go | 12 +++- pkg/storage/unified/resource/eventstore.go | 14 +++-- pkg/storage/unified/resource/kv.go | 66 ++++++++++++++----- pkg/storage/unified/resource/kv_test.go | 35 ++++++----- pkg/storage/unified/resource/metadata.go | 15 +++-- pkg/storage/unified/testing/kv.go | 73 +++++++++++++++------- 6 files changed, 149 insertions(+), 66 deletions(-) diff --git a/pkg/storage/unified/resource/datastore.go b/pkg/storage/unified/resource/datastore.go index 5838b585b8f..a3301062048 100644 --- a/pkg/storage/unified/resource/datastore.go +++ b/pkg/storage/unified/resource/datastore.go @@ -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 { diff --git a/pkg/storage/unified/resource/eventstore.go b/pkg/storage/unified/resource/eventstore.go index e1f4b565952..93f1922c5cc 100644 --- a/pkg/storage/unified/resource/eventstore.go +++ b/pkg/storage/unified/resource/eventstore.go @@ -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) { diff --git a/pkg/storage/unified/resource/kv.go b/pkg/storage/unified/resource/kv.go index 0e9c7a13d0e..613841c5994 100644 --- a/pkg/storage/unified/resource/kv.go +++ b/pkg/storage/unified/resource/kv.go @@ -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") diff --git a/pkg/storage/unified/resource/kv_test.go b/pkg/storage/unified/resource/kv_test.go index 6b7ead40879..27f24e19122 100644 --- a/pkg/storage/unified/resource/kv_test.go +++ b/pkg/storage/unified/resource/kv_test.go @@ -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) +} diff --git a/pkg/storage/unified/resource/metadata.go b/pkg/storage/unified/resource/metadata.go index ba0ea429faf..83240f97727 100644 --- a/pkg/storage/unified/resource/metadata.go +++ b/pkg/storage/unified/resource/metadata.go @@ -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 diff --git a/pkg/storage/unified/testing/kv.go b/pkg/storage/unified/testing/kv.go index 60b69d38497..79704b48b55 100644 --- a/pkg/storage/unified/testing/kv.go +++ b/pkg/storage/unified/testing/kv.go @@ -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) +}