Don't call Load after storing items

connector.go

@@ -63,22 +63,16 @@ func Init(config *redis.Options) error {
 func listenForUpdates(firstMessageSignal chan bool) {
 	receivedFirstMessage := false
 	psc = C.PSubscribe(ctx, "*")
+	ch := psc.Channel()
 
-	for {
-		msg, err := psc.ReceiveMessage(ctx)
-
+	for msg := range ch {
 		if !receivedFirstMessage && msg.Channel == firstMessageChannel {
 			// Send signal on first message to confirm PSubscribe is ready
 			firstMessageSignal <- true
 			receivedFirstMessage = true
 			continue
 		}
 
-		if err != nil {
-			// Panic if we disconnect from Redis
-			panic("Disconnected from redis")
-		}
-
 		handlersMux.RLock()
 		handlers, ok := handlers[msg.Channel]
 		handlersMux.RUnlock()

store.go

@@ -1,6 +1,8 @@
 package grocery
 
 import (
+	"reflect"
+
 	"github.com/google/uuid"
 )
 
@@ -12,6 +14,10 @@ type StoreOptions struct {
 	// this ID, unless Overwrite is set to true.
 	ID string
 
+	// Set to true if you would like to load the object from Redis back into
+	// the pointer after storing it.
+	Load bool
+
 	// Set to true if you would like to overwrite existing objects stored with
 	// this ID.
 	Overwrite bool
@@ -26,7 +32,7 @@ type StoreOptions struct {
 // prefix:id. If you would like to set a specific ID, use StoreWithOptions.
 func Store(ptr interface{}) (string, error) {
 	id := uuid.NewString()
-	return id, StoreWithOptions(ptr, &StoreOptions{id, false, nil})
+	return id, StoreWithOptions(ptr, &StoreOptions{id, false, false, nil})
 }
 
 // StoreWithOptions saves an object in Redis, like Store, but with options.
@@ -38,5 +44,25 @@ func StoreWithOptions(ptr interface{}, opts *StoreOptions) error {
 	opts.UpdateOptions.isStore = true
 	opts.UpdateOptions.storeOverwrite = opts.Overwrite
 
-	return updateInternal(opts.ID, ptr, opts.UpdateOptions)
+	if err := updateInternal(opts.ID, ptr, opts.UpdateOptions); err != nil {
+		return err
+	}
+
+	if reflect.TypeOf(ptr).Kind() == reflect.Ptr {
+		if opts.Load {
+			// Load object back into the pointer
+			if err := Load(opts.ID, ptr); err != nil {
+				return err
+			}
+		} else {
+			// Set ID
+			fi := reflect.Indirect(reflect.ValueOf(ptr).Elem()).FieldByName("ID")
+
+			if fi.String() != opts.ID {
+				fi.SetString(opts.ID)
+			}
+		}
+	}
+
+	return nil
 }

store_ref_test.go

@@ -32,6 +32,7 @@ func TestStoreReference(t *testing.T) {
 
 	if err != nil {
 		t.Error(err)
+		return
 	}
 
 	loadedB := new(B)
@@ -59,6 +60,7 @@ func TestStoreListReference(t *testing.T) {
 
 	if err != nil {
 		t.Error(err)
+		return
 	}
 
 	loadedC := new(Ctest)

update.go

@@ -58,16 +58,30 @@ func UpdateWithOptions(id string, ptr interface{}, opts *UpdateOptions) error {
 }
 
 func updateInternal(id string, ptr interface{}, opts *UpdateOptions) error {
-	if reflect.TypeOf(ptr).Kind() != reflect.Ptr || reflect.TypeOf(ptr).Elem().Kind() != reflect.Struct {
-		return errors.New("ptr must be a struct pointer")
-	}
-
 	if id == "" {
 		return errors.New("ID must not be empty")
 	}
 
+	var val reflect.Value
+	var typ reflect.Type
+
+	switch reflect.TypeOf(ptr).Kind() {
+	case reflect.Ptr:
+		if reflect.TypeOf(ptr).Elem().Kind() != reflect.Struct {
+			return errors.New("ptr must be a struct pointer")
+		}
+
+		val = reflect.ValueOf(ptr).Elem()
+		typ = reflect.TypeOf(ptr).Elem()
+	case reflect.Struct:
+		val = reflect.ValueOf(ptr)
+		typ = reflect.TypeOf(ptr)
+	default:
+		return errors.New("ptr must be a struct pointer")
+	}
+
 	// Get prefix for the struct (e.g. 'answer:' from Answer)
-	prefix := strings.ToLower(reflect.TypeOf(ptr).Elem().Name())
+	prefix := strings.ToLower(typ.Name())
 
 	// Make sure the object exists on an update, or not on a store
 	exists, _ := C.Exists(ctx, prefix+":"+id).Result()
@@ -78,8 +92,6 @@ func updateInternal(id string, ptr interface{}, opts *UpdateOptions) error {
 		return fmt.Errorf("%s:%s does not exist", prefix, id)
 	}
 
-	val := reflect.ValueOf(ptr).Elem()
-	typ := val.Type()
 	pip := opts.Pipeline
 
 	if opts.Pipeline == nil {
@@ -202,13 +214,6 @@ func updateInternal(id string, ptr interface{}, opts *UpdateOptions) error {
 	pip.HSet(ctx, prefix+":"+id, "updatedAt", time.Now().Unix())
 
 	if opts.isStore {
-		// Set ID
-		fi := reflect.Indirect(val).FieldByName("ID")
-
-		if fi.String() != id {
-			fi.SetString(id)
-		}
-
 		// Set createdAt timestamp
 		pip.HSet(ctx, prefix+":"+id, "createdAt", time.Now().Unix())
 
@@ -228,10 +233,6 @@ func updateInternal(id string, ptr interface{}, opts *UpdateOptions) error {
 		if _, err := pip.Exec(ctx); err != nil {
 			return err
 		}
-
-		if err := Load(id, ptr); err != nil {
-			return err
-		}
 	}
 
 	return nil