feat(boom): support PropertyTranslator return Key with ParentKey

mercari · Aug 30, 2020 · c8a9ccb · c8a9ccb
1 parent 0931919
commit c8a9ccb
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 12 deletions.
diff --git a/boom/boom.go b/boom/boom.go
@@ -226,6 +226,7 @@ func (bm *Boom) KeyError(src interface{}) (datastore.Key, error) {
 	}
 
 	var parent datastore.Key
+	var key datastore.Key
 	var keyName string
 	var keyID int64
 	var kind string
@@ -250,35 +251,30 @@ func (bm *Boom) KeyError(src interface{}) (datastore.Key, error) {
 						return nil, err
 					}
 					if id, ok := pv.(int64); ok {
-						if keyID != 0 || keyName != "" {
+						if key != nil || keyID != 0 || keyName != "" {
 							return nil, fmt.Errorf("boom: Only one field may be marked id")
 						}
 						keyID = id
 					} else if name, ok := pv.(string); ok {
-						if keyID != 0 || keyName != "" {
+						if key != nil || keyID != 0 || keyName != "" {
 							return nil, fmt.Errorf("boom: Only one field may be marked id")
 						}
 						keyName = name
-					} else if key, ok := pv.(datastore.Key); ok {
-						if keyID != 0 || keyName != "" {
+					} else if propertyKey, ok := pv.(datastore.Key); ok {
+						if key != nil || keyID != 0 || keyName != "" {
 							return nil, fmt.Errorf("boom: Only one field may be marked id")
 						}
-						if key.ID() != 0 {
-							keyID = key.ID()
-						} else if key.Name() != "" {
-							keyName = key.Name()
-						}
-						// TODO check kind value
+						key = propertyKey
 					}
 				} else {
 					switch vf.Kind() {
 					case reflect.Int64:
-						if keyID != 0 || keyName != "" {
+						if key != nil || keyID != 0 || keyName != "" {
 							return nil, fmt.Errorf("boom: Only one field may be marked id")
 						}
 						keyID = vf.Int()
 					case reflect.String:
-						if keyID != 0 || keyName != "" {
+						if key != nil || keyID != 0 || keyName != "" {
 							return nil, fmt.Errorf("boom: Only one field may be marked id")
 						}
 						keyName = vf.String()
@@ -328,6 +324,25 @@ func (bm *Boom) KeyError(src interface{}) (datastore.Key, error) {
 		kind = t.Name()
 	}
 
+	if key != nil {
+		if key.ParentKey() != nil && parent != nil {
+			return nil, fmt.Errorf("boom: ID field returns key. don't use parent annotated field at same time")
+		}
+		if key.Kind() != kind {
+			return nil, fmt.Errorf("boom: ID field returns key that contains unexpected kind")
+		}
+
+		if key.ParentKey() != nil {
+			return key, nil
+		}
+
+		if keyName := key.Name(); keyName != "" {
+			return bm.Client.NameKey(kind, keyName, parent), nil
+		}
+
+		return bm.Client.IDKey(kind, key.ID(), parent), nil
+	}
+
 	if keyName != "" {
 		return bm.Client.NameKey(kind, keyName, parent), nil
 	}

diff --git a/boom/boom_test.go b/boom/boom_test.go
@@ -3,6 +3,8 @@ package boom
 import (
 	"context"
 	"fmt"
+	"strconv"
+	"strings"
 	"testing"
 
 	"go.mercari.io/datastore"
@@ -11,13 +13,15 @@ import (
 
 var _ datastore.PropertyTranslator = UserID(0)
 var _ datastore.PropertyTranslator = DataID(0)
+var _ datastore.PropertyTranslator = WithAncestorID("")
 var _ datastore.PropertyTranslator = IntID(0)
 var _ datastore.PropertyTranslator = StringID("")
 
 type contextClient struct{}
 
 type UserID int64
 type DataID int64
+type WithAncestorID string
 
 type IntID int64
 type StringID string
@@ -50,6 +54,35 @@ func (id DataID) FromPropertyValue(ctx context.Context, p datastore.Property) (d
 	return DataID(key.ID()), nil
 }
 
+func (id WithAncestorID) ToPropertyValue(ctx context.Context) (interface{}, error) {
+	client := ctx.Value(contextClient{}).(datastore.Client)
+	ss := strings.SplitN(string(id), "-", 2)
+	if len(ss) != 2 {
+		return nil, fmt.Errorf("unexpected id format: %s", id)
+	}
+	userID, err := strconv.ParseInt(ss[0], 10, 64)
+	if err != nil {
+		return nil, err
+	}
+	dataID, err := strconv.ParseInt(ss[1], 10, 64)
+	if err != nil {
+		return nil, err
+	}
+	parentKey := client.IDKey("User", userID, nil)
+	key := client.IDKey("Data", dataID, parentKey)
+	return key, nil
+}
+
+func (id WithAncestorID) FromPropertyValue(ctx context.Context, p datastore.Property) (dst interface{}, err error) {
+	key, ok := p.Value.(datastore.Key)
+	if !ok {
+		return nil, datastore.ErrInvalidEntityType
+	}
+	userID := key.ParentKey().ID()
+	dataID := key.ID()
+	return WithAncestorID(fmt.Sprintf("%d-%d", userID, dataID)), nil
+}
+
 func (id IntID) ToPropertyValue(ctx context.Context) (interface{}, error) {
 	// for boom.KeyError
 	return int64(id), nil
@@ -606,6 +639,34 @@ func TestBoom_TagIDWithPropertyTranslator(t *testing.T) {
 			t.Fatal(err)
 		}
 	}
+	{ // Put & Get	 with boom:"id" that has ParentKey
+		type Data struct {
+			ID WithAncestorID `datastore:"-" boom:"id"`
+		}
+
+		key, err := bm.Put(&Data{ID: "20-100"})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if v := key.Kind(); v != "Data" {
+			t.Errorf("unexpected: %v", v)
+		}
+		if v := key.ID(); v != 100 {
+			t.Errorf("unexpected: %v", v)
+		}
+		if v := key.ParentKey().Kind(); v != "User" {
+			t.Errorf("unexpected: %v", v)
+		}
+		if v := key.ParentKey().ID(); v != 20 {
+			t.Errorf("unexpected: %v", v)
+		}
+
+		err = bm.Get(&Data{ID: "20-100"})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
 }
 
 func TestBoom_TagParent(t *testing.T) {