Merge pull request #40 from w-h-a/store-seed

feat: add seed option for store
w-h-a · Jul 8, 2024 · c408541 · c408541
2 parents cd315fc + a3c98f3
commit c408541
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/store/memory/memory.go b/store/memory/memory.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/patrickmn/go-cache"
 	"github.com/w-h-a/pkg/store"
+	"github.com/w-h-a/pkg/telemetry/log"
 )
 
 type memoryStore struct {
@@ -193,5 +194,13 @@ func NewStore(opts ...store.StoreOption) store.Store {
 		store:   cache.New(cache.NoExpiration, 5*time.Minute),
 	}
 
+	if len(options.Seed) != 0 {
+		for _, rec := range options.Seed {
+			if err := s.Write(rec); err != nil {
+				log.Fatalf("failed to seed database: %v", err)
+			}
+		}
+	}
+
 	return s
 }
diff --git a/store/options.go b/store/options.go
@@ -10,6 +10,7 @@ type StoreOptions struct {
 	Nodes    []string
 	Database string
 	Table    string
+	Seed     []*Record
 	Context  context.Context
 }
 
@@ -31,6 +32,12 @@ func StoreWithTable(tbl string) StoreOption {
 	}
 }
 
+func StoreWithSeed(rs ...*Record) StoreOption {
+	return func(o *StoreOptions) {
+		o.Seed = append(o.Seed, rs...)
+	}
+}
+
 func NewStoreOptions(opts ...StoreOption) StoreOptions {
 	options := StoreOptions{
 		Context: context.Background(),