fix(core): fix datastore.Key or []datastore.Key Save & Load handling

mercari · Nov 22, 2017 · 29f465d · 29f465d
1 parent cc8c004
commit 29f465d
Show file tree

Hide file tree

Showing 4 changed files with 162 additions and 61 deletions.
diff --git a/load.go b/load.go
@@ -356,6 +356,23 @@ func setVal(ctx context.Context, v reflect.Value, p Property) error {
 		default:
 			return typeMismatchReason(p, v)
 		}
+	case reflect.Interface:
+		switch v.Type() {
+		case typeOfKey:
+			x, ok := pValue.(Key)
+			if !ok && pValue != nil {
+				return typeMismatchReason(p, v)
+			}
+			if x == nil {
+				if !v.IsNil() {
+					v.Set(reflect.ValueOf(nil))
+				}
+			} else {
+				v.Set(reflect.ValueOf(x))
+			}
+		default:
+			return typeMismatchReason(p, v)
+		}
 	case reflect.Struct:
 		switch v.Type() {
 		case typeOfTime:

diff --git a/save.go b/save.go
@@ -68,74 +68,80 @@ func saveStructProperty(ctx context.Context, props *[]Property, name string, opt
 		return nil
 	}
 
-	switch x := v.Interface().(type) {
-	case Key, time.Time, GeoPoint:
-		p.Value = x
-	case PropertyTranslator:
-		v, err := x.ToPropertyValue(ctx)
-		if err != nil {
-			return err
-		}
-		p.Value = v
-	default:
-		switch v.Kind() {
-		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
-			p.Value = v.Int()
-		case reflect.Bool:
-			p.Value = v.Bool()
-		case reflect.String:
-			p.Value = v.String()
-		case reflect.Float32, reflect.Float64:
-			p.Value = v.Float()
-		case reflect.Slice:
-			if v.Type().Elem().Kind() == reflect.Uint8 {
-				p.Value = v.Bytes()
-			} else {
-				return saveSliceProperty(ctx, props, name, opts, v)
-			}
-		case reflect.Ptr:
-			if v.Type().Elem().Kind() != reflect.Struct {
-				return fmt.Errorf("datastore: unsupported struct field type: %s", v.Type())
-			}
-			if v.IsNil() {
-				return nil
-			}
-			v = v.Elem()
-			fallthrough
-		case reflect.Struct:
-			if !v.CanAddr() {
-				return fmt.Errorf("datastore: unsupported struct field: value is unaddressable")
-			}
-			vi := v.Addr().Interface()
-
-			sub, err := newStructPLS(vi)
-			if err != nil {
-				return fmt.Errorf("datastore: unsupported struct field: %v", err)
-			}
-
-			if opts.flatten {
-				return sub.save(ctx, props, opts, name+".")
-			}
+	if v.Type().AssignableTo(typeOfKey) {
+		p.Value = v.Interface()
 
-			var subProps []Property
-			err = sub.save(ctx, &subProps, opts, "")
-			if err != nil {
-				return err
-			}
-			subKey, err := sub.key(v)
+	} else {
+		switch x := v.Interface().(type) {
+		case time.Time, GeoPoint:
+			p.Value = x
+		case PropertyTranslator:
+			v, err := x.ToPropertyValue(ctx)
 			if err != nil {
 				return err
 			}
-
-			p.Value = &Entity{
-				Key:        subKey,
-				Properties: subProps,
+			p.Value = v
+		default:
+			switch v.Kind() {
+			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+				p.Value = v.Int()
+			case reflect.Bool:
+				p.Value = v.Bool()
+			case reflect.String:
+				p.Value = v.String()
+			case reflect.Float32, reflect.Float64:
+				p.Value = v.Float()
+			case reflect.Slice:
+				if v.Type().Elem().Kind() == reflect.Uint8 {
+					p.Value = v.Bytes()
+				} else {
+					return saveSliceProperty(ctx, props, name, opts, v)
+				}
+			case reflect.Ptr:
+				if v.Type().Elem().Kind() != reflect.Struct {
+					return fmt.Errorf("datastore: unsupported struct field type: %s", v.Type())
+				}
+				if v.IsNil() {
+					return nil
+				}
+				v = v.Elem()
+				fallthrough
+			case reflect.Struct:
+				if !v.CanAddr() {
+					return fmt.Errorf("datastore: unsupported struct field: value is unaddressable")
+				}
+				vi := v.Addr().Interface()
+
+				sub, err := newStructPLS(vi)
+				if err != nil {
+					return fmt.Errorf("datastore: unsupported struct field: %v", err)
+				}
+
+				if opts.flatten {
+					return sub.save(ctx, props, opts, name+".")
+				}
+
+				var subProps []Property
+				err = sub.save(ctx, &subProps, opts, "")
+				if err != nil {
+					return err
+				}
+				subKey, err := sub.key(v)
+				if err != nil {
+					return err
+				}
+
+				p.Value = &Entity{
+					Key:        subKey,
+					Properties: subProps,
+				}
 			}
 		}
+		if p.Value == nil {
+			return fmt.Errorf("datastore: unsupported struct field type: %v", v.Type())
+		}
 	}
-	if p.Value == nil {
-		return fmt.Errorf("datastore: unsupported struct field type: %v", v.Type())
-	}
+
 	*props = append(*props, p)
 	return nil
 }

diff --git a/testsuite/entity_ops.go b/testsuite/entity_ops.go
@@ -225,6 +225,82 @@ func PutEntityType(t *testing.T, ctx context.Context, client datastore.Client) {
 	}
 }
 
+func PutAndGetNilKey(t *testing.T, ctx context.Context, client datastore.Client) {
+	defer func() {
+		err := client.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	type Data struct {
+		KeyA datastore.Key
+		KeyB datastore.Key
+	}
+
+	key := client.IncompleteKey("Test", nil)
+	key, err := client.Put(ctx, key, &Data{
+		KeyA: client.NameKey("Test", "a", nil),
+		KeyB: nil,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	obj := &Data{}
+	err = client.Get(ctx, key, obj)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if v := obj.KeyA; v == nil {
+		t.Errorf("unexpected: %v", v)
+	}
+	if v := obj.KeyB; v != nil {
+		t.Errorf("unexpected: %v", v)
+	}
+}
+
+func PutAndGetNilKeySlice(t *testing.T, ctx context.Context, client datastore.Client) {
+	defer func() {
+		err := client.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	type Data struct {
+		Keys []datastore.Key
+	}
+
+	key := client.IncompleteKey("Test", nil)
+	key, err := client.Put(ctx, key, &Data{
+		Keys: []datastore.Key{
+			client.NameKey("Test", "a", nil),
+			nil,
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	obj := &Data{}
+	err = client.Get(ctx, key, obj)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if v := len(obj.Keys); v != 2 {
+		t.Fatalf("unexpected: %v", v)
+	}
+	if v := obj.Keys[0]; v == nil {
+		t.Errorf("unexpected: %v", v)
+	}
+	if v := obj.Keys[1]; v != nil {
+		t.Errorf("unexpected: %v", v)
+	}
+}
+
 type EntityInterface interface {
 	Kind() string
 	ID() string

diff --git a/testsuite/suite.go b/testsuite/suite.go
@@ -18,6 +18,8 @@ var TestSuite = map[string]Test{
 	"PutAndGet_ObjectHasObjectSlice":            PutAndGet_ObjectHasObjectSlice,
 	"PutAndGet_ObjectHasObjectSliceWithFlatten": PutAndGet_ObjectHasObjectSliceWithFlatten,
 	"PutEntityType":                             PutEntityType,
+	"PutAndGetNilKey":                           PutAndGetNilKey,
+	"PutAndGetNilKeySlice":                      PutAndGetNilKeySlice,
 	"PutInterface":                              PutInterface,
 	"GeoPoint_PutAndGet":                        GeoPoint_PutAndGet,
 	"Key_Equal":                                 Key_Equal,