Merge pull request #9 from chenyanchen/feat/batchkv

feat: define BatchKV and implement the cacheBatchKV with cacheKV

README.md

@@ -1,3 +1,3 @@
 # db
 
-Basic database abstract and implement.
+Basic database abstract and implementations.

batchkv.go

@@ -0,0 +1,10 @@
+package db
+
+import "context"
+
+// BatchKV is a batch key-value storage.
+type BatchKV[K comparable, V any] interface {
+	Get(ctx context.Context, keys []K) (map[K]V, error)
+	Set(ctx context.Context, kvs map[K]V) error
+	Del(ctx context.Context, keys []K) error
+}

cachebatchkv/cache.go

@@ -0,0 +1,105 @@
+package cachebatchkv
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/chenyanchen/db"
+	"github.com/chenyanchen/db/cachekv"
+)
+
+// cacheBatchKV is a struct that contains a cache and a source BatchKV.
+// It is used to cache batch operations.
+//
+// Important: cacheBatchKV are not guaranteed to get all the values of keys, it guarantees no error.
+type cacheBatchKV[K comparable, V any] struct {
+	cache db.KV[K, V]
+
+	source db.BatchKV[K, V]
+}
+
+// New creates a new cacheBatchKV instance with the given source and options.
+func New[K comparable, V any](
+	source db.BatchKV[K, V],
+	options ...cachekv.CacheOption[K, V],
+) *cacheBatchKV[K, V] {
+	return &cacheBatchKV[K, V]{
+		cache:  cachekv.New(options...),
+		source: source,
+	}
+}
+
+// Get retrieves the values for the given keys from the cache.
+// If a key is not found in the cache, it is retrieved from the source BatchKV.
+func (c *cacheBatchKV[K, V]) Get(ctx context.Context, keys []K) (map[K]V, error) {
+	result := make(map[K]V, len(keys))
+
+	// misses is a slice that contains the keys that are not found in the cache.
+	var misses []K
+
+	for _, key := range keys {
+		v, err := c.cache.Get(ctx, key)
+		if err == nil {
+			result[key] = v
+			continue
+		}
+
+		if errors.Is(err, db.ErrNotFound) {
+			misses = append(misses, key)
+			continue
+		}
+
+		return nil, err
+	}
+
+	if c.source == nil {
+		return result, nil
+	}
+
+	get, err := c.source.Get(ctx, misses)
+	if err != nil {
+		return result, err
+	}
+
+	for k, v := range get {
+		result[k] = v
+		if err = c.cache.Set(ctx, k, v); err != nil {
+			return nil, fmt.Errorf("set cache: %w", err)
+		}
+	}
+
+	return result, nil
+}
+
+// Set sets the values for the given keys in the cache.
+// It also sets the values in the source BatchKV if it exists.
+func (c *cacheBatchKV[K, V]) Set(ctx context.Context, m map[K]V) error {
+	for k, v := range m {
+		if err := c.cache.Set(ctx, k, v); err != nil {
+			return fmt.Errorf("set cache: %w", err)
+		}
+	}
+
+	if c.source == nil {
+		return nil
+	}
+
+	return c.source.Set(ctx, m)
+}
+
+// Del deletes the values for the given keys from the cache.
+// It also deletes the values from the source BatchKV if it exists.
+func (c *cacheBatchKV[K, V]) Del(ctx context.Context, keys []K) error {
+	for _, key := range keys {
+		if err := c.cache.Del(ctx, key); err != nil {
+			return fmt.Errorf("del cache: %w", err)
+		}
+	}
+
+	if c.source == nil {
+		return nil
+	}
+
+	return c.source.Del(ctx, keys)
+}

cachekv/cache.go

@@ -104,7 +104,7 @@ func (c *cacheKV[K, V]) Get(ctx context.Context, k K) (V, error) {
 	}
 	if c.source == nil {
 		c.telemetry(k, "miss_mem")
-		return v, fmt.Errorf("not found: %+v", k)
+		return v, db.ErrNotFound
 	}
 	got, err := c.source.Get(ctx, k)
 	if err != nil {

examples/batch/main.go

@@ -3,67 +3,49 @@ package main
 import (
 	"context"
 	"fmt"
+	"strconv"
+	"time"
 
-	"github.com/chenyanchen/db"
+	"github.com/chenyanchen/db/cachebatchkv"
+	"github.com/chenyanchen/db/cachekv"
 )
 
 func main() {
 	ctx := context.Background()
 
-	// You can define K as all comparable types, like this:
-	// array: [3]int64
-	// string (SQL): "id in (1, 2, 3)"
-	// string (JSON): `{"ids": [1, 2, 3]}`
-	var mapKV db.KV[[3]int64, []Content] = NewMapKV()
+	batchKV := cachebatchkv.New[int64, Content](
+		&fakeContentKV{},
+		cachekv.WithExpires[int64, Content](time.Hour),
+	)
 
-	if err := mapKV.Set(ctx, [3]int64{}, []Content{{ID: 1, Title: "A"}, {ID: 2, Title: "B"}}); err != nil {
-		panic(err)
-	}
-
-	contents, err := mapKV.Get(ctx, [3]int64{1, 3})
+	contents, err := batchKV.Get(ctx, []int64{1, 3, 5})
 	if err != nil {
 		panic(err)
 	}
 
 	fmt.Println("contents:", contents)
-	// output: contents: [{1 A}]
+	// output: contents: map[1:{1 Title: 1} 3:{3 Title: 3} 5:{5 Title: 5}]
 }
 
 type Content struct {
 	ID    int64
 	Title string
 }
 
-type mapKV struct {
-	m map[int64]Content
-}
-
-func NewMapKV() *mapKV {
-	return &mapKV{m: make(map[int64]Content)}
-}
+type fakeContentKV struct{}
 
-func (s *mapKV) Get(ctx context.Context, ids [3]int64) ([]Content, error) {
-	contents := make([]Content, 0, len(ids))
-	for _, id := range ids {
-		content, ok := s.m[id]
-		if !ok {
-			continue
-		}
-		contents = append(contents, content)
+func (s *fakeContentKV) Get(ctx context.Context, keys []int64) (map[int64]Content, error) {
+	result := make(map[int64]Content, len(keys))
+	for _, key := range keys {
+		result[key] = Content{ID: key, Title: "Title: " + strconv.FormatInt(key, 10)}
 	}
-	return contents, nil
+	return result, nil
 }
 
-func (s *mapKV) Set(ctx context.Context, ids [3]int64, contents []Content) error {
-	for _, content := range contents {
-		s.m[content.ID] = content
-	}
+func (s *fakeContentKV) Set(ctx context.Context, kvs map[int64]Content) error {
 	return nil
 }
 
-func (s *mapKV) Del(ctx context.Context, ids [3]int64) error {
-	for _, id := range ids {
-		delete(s.m, id)
-	}
+func (s *fakeContentKV) Del(ctx context.Context, keys []int64) error {
 	return nil
 }