A database/db.go

@@ -0,0 +1,52 @@ 
+package database
+
+import (
+	"github.com/dgraph-io/badger/v3"
+)
+
+// Use this as a global DB handle.
+var DB database
+
+type database struct {
+	*badger.DB
+}
+
+func Open() (*badger.DB, error) {
+	db, err := badger.Open(badger.DefaultOptions("./badger"))
+	if err != nil {
+		return nil, err
+	}
+	return db, nil
+}
+
+// Wrapper function to simplify setting a key/val
+// in badger.
+func (d *database) Set(key, val []byte) error {
+	err := d.Update(func(txn *badger.Txn) error {
+		err := txn.Set(key, val)
+		return err
+	})
+
+	return err
+}
+
+// Wrapper function to simplify getting a key from badger.
+func (d *database) Get(key []byte) ([]byte, error) {
+	var val []byte
+	err := d.View(func(txn *badger.Txn) error {
+		item, err := txn.Get(key)
+		if err != nil {
+			return err
+		}
+
+		err = item.Value(func(v []byte) error {
+			val, err = item.ValueCopy(nil)
+			return err
+		})
+		return err
+	})
+	if err != nil {
+		return nil, err
+	}
+	return val, nil
+}

D db/db.go

@@ -1,7 +0,0 @@ 
-package db
-
-import (
-	"github.com/dgraph-io/badger/v3"
-)
-
-var DB *badger.DB

M main.go → main.go

@@ -4,9 +4,8 @@ import (
 	"log"
 	"net"
 
-	"git.icyphox.sh/paprika/db"
+	"git.icyphox.sh/paprika/database"
 	"git.icyphox.sh/paprika/plugins"
-	"github.com/dgraph-io/badger/v3"
 	"gopkg.in/irc.v3"
 )
 
@@ -45,11 +44,11 @@ Name:    "paprika",
 		Handler: irc.HandlerFunc(ircHandler),
 	}
 
-	db.DB, err = badger.Open(badger.DefaultOptions("./badger"))
+	database.DB.DB, err = database.Open()
 	if err != nil {
 		log.Fatal(err)
 	}
-	defer db.DB.Close()
+	defer database.DB.Close()
 
 	client := irc.NewClient(conn, config)
 	err = client.Run()

M plugins/listenbrainz/db.go → plugins/listenbrainz/db.go

@@ -3,41 +3,24 @@ 
 import (
 	"fmt"
 
-	"git.icyphox.sh/paprika/db"
-	"github.com/dgraph-io/badger/v3"
+	"git.icyphox.sh/paprika/database"
 )
 
 // Store the Listenbrainz username against the nick.
 func Setup(lbzUser, nick string) error {
-	err := db.DB.Update(func(txn *badger.Txn) error {
-		err := txn.Set([]byte(fmt.Sprintf("lbz/%s", nick)), []byte(lbzUser))
-		return err
-	})
-	if err != nil {
-		return err
-	}
-
-	return nil
+	err := database.DB.Set(
+		[]byte(fmt.Sprintf("lbz/%s", nick)),
+		[]byte(lbzUser),
+	)
+	return err
 }
 
 // Gets the Listenbrainz username from the DB.
 func GetUser(nick string) (string, error) {
-	var user string
-	err := db.DB.View(func(txn *badger.Txn) error {
-		item, err := txn.Get([]byte(fmt.Sprintf("lbz/%s", nick)))
-		if err != nil {
-			return err
-		}
-
-		err = item.Value(func(val []byte) error {
-			user = string(val)
-			return nil
-		})
-		return nil
-	})
-
+	nick = fmt.Sprintf("lbz/%s", nick)
+	user, err := database.DB.Get([]byte(nick))
 	if err != nil {
 		return "", err
 	}
-	return user, nil
+	return string(user), nil
 }