From 559377ac03a27738515143a4966d7d3efc6ad2ac Mon Sep 17 00:00:00 2001
From: tauraamui <adampstringer@protonmail.com>
Date: Sat, 8 Jul 2023 22:58:06 +0100
Subject: [PATCH] feat: ability to query specific records by their row ID

---
 example/hiarchy.go    |  2 +-
 example/simple.go     |  2 +-
 storage/store.go      | 33 ++++++++++++++++++++++++++++++++-
 storage/store_test.go | 33 ++++++++++++++++++++++++++++++++-
 4 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/example/hiarchy.go b/example/hiarchy.go
index 45b2f45..87960a8 100644
--- a/example/hiarchy.go
+++ b/example/hiarchy.go
@@ -47,7 +47,7 @@ func hiarchy() {
 	store.Save(child.UUID, &healthyishCarrotCake)
 	store.Save(child.UUID, &redVelvetCake)
 
-	bs, err := storage.LoadAllByOwner(store, Cake{}, child.UUID)
+	bs, err := storage.LoadAll(store, Cake{}, child.UUID)
 	for _, cake := range bs {
 		fmt.Printf("ROWID: %d, %+v\n", cake.ID, cake)
 	}
diff --git a/example/simple.go b/example/simple.go
index 149ef1e..c909eab 100644
--- a/example/simple.go
+++ b/example/simple.go
@@ -30,7 +30,7 @@ func main() {
 	store.Save(kvs.RootOwner{}, &Balloon{Color: "RED", Size: 695})
 	store.Save(kvs.RootOwner{}, &Balloon{Color: "WHITE", Size: 366})
 
-	bs, err := storage.LoadAllByOwner(store, Balloon{}, kvs.RootOwner{})
+	bs, err := storage.LoadAll(store, Balloon{}, kvs.RootOwner{})
 	for _, balloon := range bs {
 		fmt.Printf("ROWID: %d, %+v\n", balloon.ID, balloon)
 	}
diff --git a/storage/store.go b/storage/store.go
index bbd6b3c..6d8a5ae 100644
--- a/storage/store.go
+++ b/storage/store.go
@@ -49,7 +49,38 @@ type TableNamer interface {
 	TableName() string
 }
 
-func LoadAllByOwner[T TableNamer](s Store, v T, owner kvs.UUID) ([]T, error) {
+func Load[T TableNamer](s Store, dest T, owner kvs.UUID, rowID uint32) error {
+	db := s.db
+
+	blankEntries := kvs.ConvertToBlankEntries(dest.TableName(), owner, rowID, dest)
+	for _, ent := range blankEntries {
+		db.View(func(txn *badger.Txn) error {
+			item, err := txn.Get(ent.Key())
+			if err != nil {
+				return err
+			}
+
+			if err := item.Value(func(val []byte) error {
+				ent.Data = val
+				return nil
+			}); err != nil {
+				return err
+			}
+
+			return nil
+		})
+
+		ent.RowID = rowID
+
+		if err := kvs.LoadEntry(dest, ent); err != nil {
+			return err
+		}
+	}
+
+	return kvs.LoadID(dest, rowID)
+}
+
+func LoadAll[T TableNamer](s Store, v T, owner kvs.UUID) ([]T, error) {
 	db := s.db
 	dest := []T{}
 
diff --git a/storage/store_test.go b/storage/store_test.go
index dcf6d60..3eda596 100644
--- a/storage/store_test.go
+++ b/storage/store_test.go
@@ -45,7 +45,7 @@ func TestStoreAndLoadMultipleBalloonsSuccess(t *testing.T) {
 	is.NoErr(store.Save(kvs.RootOwner{}, &smallYellowBalloon))
 	is.NoErr(store.Save(kvs.RootOwner{}, &mediumWhiteBalloon))
 
-	bs, err := storage.LoadAllByOwner(store, Balloon{}, kvs.RootOwner{})
+	bs, err := storage.LoadAll(store, Balloon{}, kvs.RootOwner{})
 	is.NoErr(err)
 
 	is.True(len(bs) == 3)
@@ -55,6 +55,37 @@ func TestStoreAndLoadMultipleBalloonsSuccess(t *testing.T) {
 	is.Equal(bs[2], Balloon{ID: 2, Color: "WHITE", Size: 366})
 }
 
+func TestStoreLoadMultipleLoadIndividualBalloonsSuccess(t *testing.T) {
+	is := is.New(t)
+
+	db, err := kvs.NewMemKVDB()
+	is.NoErr(err)
+	defer db.Close()
+
+	store := storage.New(db)
+	defer store.Close()
+
+	bigRedBalloon := Balloon{Color: "RED", Size: 695}
+	smallYellowBalloon := Balloon{Color: "YELLOW", Size: 112}
+	mediumWhiteBalloon := Balloon{Color: "WHITE", Size: 366}
+	is.NoErr(store.Save(kvs.RootOwner{}, &bigRedBalloon))
+	is.NoErr(store.Save(kvs.RootOwner{}, &smallYellowBalloon))
+	is.NoErr(store.Save(kvs.RootOwner{}, &mediumWhiteBalloon))
+
+	bs0 := Balloon{}
+	is.NoErr(storage.Load(store, &bs0, kvs.RootOwner{}, 0))
+
+	bs1 := Balloon{}
+	is.NoErr(storage.Load(store, &bs1, kvs.RootOwner{}, 1))
+
+	bs2 := Balloon{}
+	is.NoErr(storage.Load(store, &bs2, kvs.RootOwner{}, 2))
+
+	is.Equal(bs0, Balloon{ID: 0, Color: "RED", Size: 695})
+	is.Equal(bs1, Balloon{ID: 1, Color: "YELLOW", Size: 112})
+	is.Equal(bs2, Balloon{ID: 2, Color: "WHITE", Size: 366})
+}
+
 func TestStoreMultipleBalloonsSuccess(t *testing.T) {
 	is := is.New(t)