set storage in context

vinceanalytics · Feb 9, 2024 · 4477cc9 · 4477cc9
1 parent 5f06c00
commit 4477cc9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/db/persistance.go b/db/persistance.go
@@ -2,6 +2,7 @@ package db
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -200,6 +201,16 @@ type KV struct {
 	db *badger.DB
 }
 
+type dbKey struct{}
+
+func With(ctx context.Context, kv Storage) context.Context {
+	return context.WithValue(ctx, dbKey{}, kv)
+}
+
+func Get(ctx context.Context) Storage {
+	return ctx.Value(dbKey{}).(Storage)
+}
+
 func NewKV(path string) (*KV, error) {
 	db, err := badger.Open(badger.DefaultOptions(filepath.Join(path, "db")).
 		WithLogger(&badgerLogger{

diff --git a/main.go b/main.go
@@ -192,7 +192,7 @@ func app() *cli.Command {
 			if err != nil {
 				return err
 			}
-
+			ctx = db.With(ctx, store)
 			log.Info("Setup session")
 			alloc := memory.NewGoAllocator()
 			log.Info("Loading primary index")