datastore: Load key even if Load returns error

Load as much as possible of both the key and properties when
a custom KeyLoader is used. If both load methods return an
error, the error of the LoadKey method will be returned.

Fixes #1575.

Change-Id: Ib063224a9ff7f7c3cdf77dfa882c2d13dc6d2ac8
Reviewed-on: https://code-review.googlesource.com/c/gocloud/+/45730
Reviewed-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Jean de Klerk <deklerk@google.com>
diff --git a/datastore/load.go b/datastore/load.go
index ffed38b..8981044 100644
--- a/datastore/load.go
+++ b/datastore/load.go
@@ -431,14 +431,19 @@
 
 func loadEntity(dst interface{}, ent *Entity) error {
 	if pls, ok := dst.(PropertyLoadSaver); ok {
-		err := pls.Load(ent.Properties)
-		if err != nil {
-			return err
-		}
+		// Load both key and properties. Try to load as much as possible, even
+		// if an error occurs during loading loading either the key or the
+		// properties.
+		var keyLoadErr error
 		if e, ok := dst.(KeyLoader); ok {
-			err = e.LoadKey(ent.Key)
+			keyLoadErr = e.LoadKey(ent.Key)
 		}
-		return err
+		loadErr := pls.Load(ent.Properties)
+		// Let any error returned by LoadKey prevail above any error from Load.
+		if keyLoadErr != nil {
+			return keyLoadErr
+		}
+		return loadErr
 	}
 	return loadEntityToStruct(dst, ent)
 }
diff --git a/datastore/load_test.go b/datastore/load_test.go
index 24167ba..75ccab9 100644
--- a/datastore/load_test.go
+++ b/datastore/load_test.go
@@ -688,6 +688,110 @@
 	return nil
 }
 
+type PLS1 struct {
+	A string
+}
+
+func (p *PLS1) Load(props []Property) error {
+	for _, pp := range props {
+		if pp.Name == "A" {
+			p.A = pp.Value.(string)
+		}
+	}
+	return nil
+}
+
+func (p *PLS1) Save() (props []Property, err error) {
+	return []Property{{Name: "A", Value: p.A}}, nil
+}
+
+type KeyLoader6 struct {
+	A string
+	B string
+	K *Key
+}
+
+func (kl *KeyLoader6) Load(props []Property) error {
+	for _, pp := range props {
+		if pp.Name == "A" {
+			kl.A = pp.Value.(string)
+		}
+	}
+	return &ErrFieldMismatch{
+		StructType: reflect.TypeOf(kl),
+		FieldName:  "B",
+		Reason:     "no value found",
+	}
+}
+
+func (kl *KeyLoader6) LoadKey(k *Key) error {
+	kl.K = k
+	return nil
+}
+
+func (kl *KeyLoader6) Save() (props []Property, err error) {
+	return []Property{{}}, nil
+}
+
+type KeyLoader7 struct {
+	A string
+	K *Key
+}
+
+func (kl *KeyLoader7) Load(props []Property) error {
+	for _, pp := range props {
+		if pp.Name == "A" {
+			kl.A = pp.Value.(string)
+		}
+	}
+	return nil
+}
+
+func (kl *KeyLoader7) LoadKey(k *Key) error {
+	return &ErrFieldMismatch{
+		StructType: reflect.TypeOf(kl),
+		FieldName:  "key",
+		Reason:     "no value found",
+	}
+}
+
+func (kl *KeyLoader7) Save() (props []Property, err error) {
+	return []Property{{}}, nil
+}
+
+type KeyLoader8 struct {
+	A string
+	B string
+	K *Key
+}
+
+type customLoadError struct{}
+
+func (e *customLoadError) Error() string {
+	return "custom load error"
+}
+
+func (kl *KeyLoader8) Load(props []Property) error {
+	for _, pp := range props {
+		if pp.Name == "A" {
+			kl.A = pp.Value.(string)
+		}
+	}
+	return &customLoadError{}
+}
+
+func (kl *KeyLoader8) LoadKey(k *Key) error {
+	return &ErrFieldMismatch{
+		StructType: reflect.TypeOf(kl),
+		FieldName:  "key",
+		Reason:     "no value found",
+	}
+}
+
+func (kl *KeyLoader8) Save() (props []Property, err error) {
+	return []Property{{}}, nil
+}
+
 type NotKeyLoader struct {
 	A string
 	K *Key
@@ -842,6 +946,50 @@
 				PLS: &NotKeyLoader{A: "something"},
 			},
 		},
+		{
+			desc: "simple key loader with ErrFieldMismatch error",
+			src: &pb.Entity{
+				Key: keyToProto(testKey0),
+				Properties: map[string]*pb.Value{
+					"A": {ValueType: &pb.Value_StringValue{StringValue: "hello"}},
+				},
+			},
+			dst: &KeyLoader6{},
+			want: &KeyLoader6{
+				A: "hello",
+				B: "",
+				K: testKey0,
+			},
+		},
+		{
+			desc: "simple key loader with ErrFieldMismatch during key load",
+			src: &pb.Entity{
+				Key: keyToProto(testKey0),
+				Properties: map[string]*pb.Value{
+					"A": {ValueType: &pb.Value_StringValue{StringValue: "hello"}},
+				},
+			},
+			dst: &KeyLoader7{},
+			want: &KeyLoader7{
+				A: "hello",
+				K: nil,
+			},
+		},
+		{
+			desc: "simple key loader with other error during Load and ErrFieldMismatch during KeyLoad",
+			src: &pb.Entity{
+				Key: keyToProto(testKey0),
+				Properties: map[string]*pb.Value{
+					"A": {ValueType: &pb.Value_StringValue{StringValue: "hello"}},
+				},
+			},
+			dst: &KeyLoader8{},
+			want: &KeyLoader8{
+				A: "hello",
+				B: "",
+				K: nil,
+			},
+		},
 	}
 
 	for _, tc := range testCases {