Refactor keys and queries in datastore service and implementation.

service/datastore/key.go

Index: service/datastore/key.go
diff --git a/service/datastore/key.go b/service/datastore/key.go
new file mode 100644
index 0000000000000000000000000000000000000000..e4e3927d5db02484e9a66e26d3633fa56c4bf316
--- /dev/null
+++ b/service/datastore/key.go
@@ -0,0 +1,384 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package datastore
+
+import (
+	"bytes"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/golang/protobuf/proto"
+	pb "github.com/luci/gae/service/datastore/internal/protos/datastore"
+)
+
+// KeyTok is a single token from a multi-part Key.
+type KeyTok struct {
+	Kind     string
+	IntID    int64
+	StringID string
+}
+
+func (k KeyTok) Incomplete() bool {
+	return k.StringID == "" && k.IntID == 0
+}
+
+func (k KeyTok) Special() bool {
+	return len(k.Kind) >= 2 && k.Kind[:2] == "__"
+}
+
+func (k KeyTok) ID() Property {
+	if k.StringID != "" {
+		return Property{value: k.StringID, propType: PTString}
+	}
+	return Property{value: k.IntID, propType: PTInt}
+}
+
+func (k KeyTok) Less(other KeyTok) bool {
+	if k.Kind < other.Kind {

[dnj, 2015/09/18 16:47:58]

(see note about double-compare below)

+		return true
+	} else if k.Kind > other.Kind {
+		return false
+	}
+	a, b := k.ID(), other.ID()
+	return a.Less(&b)
+}
+
+// Key is the type used for all datastore operations.
+type Key struct {
+	appID     string
+	namespace string
+	toks      []KeyTok

[dnj, 2015/09/18 16:47:58]

Note that by construction, this array is expected to have nonzero length if the
key is not nil.

[iannucci, 2015/09/18 22:25:48]

done

+}
+
+var _ interface {
+	json.Marshaler
+	json.Unmarshaler
+} = (*Key)(nil)
+
+// NewToks creates a new Key. It is the Key implementation
+// returned from the various PropertyMap serialization routines, as well as
+// the native key implementation for the in-memory implementation of gae.
+//
+// See Interface.NewKeyToks for a version of this function which automatically
+// provides aid and ns.
+func NewKeyToks(aid, ns string, toks []KeyTok) *Key {
+	if len(toks) == 0 {
+		return nil
+	}
+	newToks := make([]KeyTok, len(toks))
+	copy(newToks, toks)

[dnj, 2015/09/18 16:47:58]

Since a KeyTok's fields are mutable, should this be a deep copy?

[iannucci, 2015/09/18 22:25:48]

It is a deep copy, since there are no pointers involved :)

+	return &Key{aid, ns, newToks}
+}
+
+// NewKey is a wrapper around NewToks which has an interface similar
+// to NewKey in the SDK.
+//
+// See Interface.NewKey for a version of this function which automatically
+// provides aid and ns.
+func NewKey(aid, ns, kind, stringID string, intID int64, parent *Key) *Key {
+	if parent == nil {
+		return &Key{aid, ns, []KeyTok{{kind, intID, stringID}}}
+	}
+
+	toks := parent.toks
+	newToks := make([]KeyTok, len(toks), len(toks)+1)
+	copy(newToks, toks)
+	newToks = append(newToks, KeyTok{kind, intID, stringID})
+	return &Key{aid, ns, newToks}
+}
+
+// MakeKey is a convenience function for manufacturing a *Key. It should only
+// be used when elems... is known statically (e.g. in the code) to be correct.
+//
+// elems is pairs of (string, string|int|int32|int64) pairs, which correspond to
+// Kind/id pairs. Example:
+//   MakeKey("aid", "namespace", "Parent", 1, "Child", "id")
+//
+// Would create the key:
+//   aid:namespace:/Parent,1/Child,id
+//
+// If elems is not parsable (e.g. wrong length, wrong types, etc.) this method
+// will panic.
+//
+// See Interface.MakeKey for a version of this function which automatically
+// provides aid and ns.
+func MakeKey(aid, ns string, elems ...interface{}) *Key {
+	if len(elems) == 0 {
+		return nil
+	}
+
+	if len(elems)%2 != 0 {
+		panic(fmt.Errorf("datastore.MakeKey: odd number of tokens: %v", elems))
+	}
+	toks := make([]KeyTok, len(elems)/2)
+	for i := 0; i < len(elems); i += 2 {

[dnj, 2015/09/18 16:47:58]

Consider rewriting using slice-based traversal. I've always found that a bit
cleaner.

for len(elems) > 0 {
  knd, ok := elems[0].(string)
  ...
  tok := &toks[tokIdx++]
  // operate on tok, avoids all the repeated division/indexing.

  elems = elems[2:]
}

[iannucci, 2015/09/18 22:25:48]

Done

+		knd, ok := elems[i].(string)
+		if !ok {
+			panic(fmt.Errorf("datastore.MakeKey: bad kind: %v", elems[i]))
+		}
+		toks[i/2].Kind = knd
+		switch x := elems[i+1].(type) {
+		case string:
+			toks[i/2].StringID = x
+		case int:
+			toks[i/2].IntID = int64(x)
+		case int32:
+			toks[i/2].IntID = int64(x)
+		case int64:
+			toks[i/2].IntID = int64(x)
+		default:
+			panic(fmt.Errorf("datastore.MakeKey: bad id: %v", x))
+		}
+	}
+
+	return NewKeyToks(aid, ns, toks)
+}
+
+// NewKeyEncoded decodes and returns a *Key
+func NewKeyEncoded(encoded string) (ret *Key, err error) {
+	ret = &Key{}
+	// Re-add padding
+	if m := len(encoded) % 4; m != 0 {
+		encoded += strings.Repeat("=", 4-m)
+	}
+	b, err := base64.URLEncoding.DecodeString(encoded)
+	if err != nil {
+		return
+	}
+
+	r := &pb.Reference{}
+	if err = proto.Unmarshal(b, r); err != nil {
+		return
+	}
+
+	ret.appID = r.GetApp()
+	ret.namespace = r.GetNameSpace()
+	ret.toks = make([]KeyTok, len(r.Path.Element))
+	for i, e := range r.Path.Element {
+		ret.toks[i] = KeyTok{
+			Kind:     e.GetType(),
+			IntID:    e.GetId(),
+			StringID: e.GetName(),
+		}
+	}
+	return
+}
+
+func (k *Key) Last() KeyTok {
+	if len(k.toks) > 0 {
+		return k.toks[len(k.toks)-1]
+	}
+	return KeyTok{}
+}
+
+// AppID returns the application ID that this Key is for.
+func (k *Key) AppID() string { return k.appID }

[dnj, 2015/09/18 16:47:58]

(*Key)(nil) is treated inconsistently. For some methods it returns a default
value, for others it would clearly panic. I generally don't like the pattern of
calling methods on nil objects, but I see the appeal WRT allocating empty keys.
This seems like an uncommon case, though, so WDYT about
s/nil-key/key-with-no-tokens/?

[iannucci, 2015/09/18 22:25:48]

In the SDK it's impossible to have a non-nil key without values. I'd rather
retain that. I just let everything panic when k is nil now.

+
+// Namespace returns the namespace that this Key is for.
+func (k *Key) Namespace() string { return k.namespace }
+
+// String returns a human-readable representation of the key in the form of
+//   AID:NS:/Kind,id/Kind,id/...
+func (k *Key) String() string {
+	if k == nil {
+		return ""
+	}
+	b := bytes.NewBuffer(make([]byte, 0, 512))

[dnj, 2015/09/18 16:47:57]

I prefer:
bytes.Buffer{}
b.Grow(512)

But really why preallocate at all? String() is not a necessary function, so IMO
let the Buffer grow on demand.

[iannucci, 2015/09/18 22:25:48]

This is a carry over from
https://github.com/golang/appengine/blob/master/datastore/key.go#L129 . I think
I'll just leave it in.

+	fmt.Fprintf(b, "%s:%s:", k.appID, k.namespace)
+	for _, t := range k.toks {
+		if t.StringID != "" {
+			fmt.Fprintf(b, "/%s,%q", t.Kind, t.StringID)
+		} else {
+			fmt.Fprintf(b, "/%s,%d", t.Kind, t.IntID)
+		}
+	}
+	return b.String()
+}
+
+// Incomplete returns true iff k doesn't have an id yet.
+func (k *Key) Incomplete() bool {
+	return k != nil && k.Last().Incomplete()
+}
+
+// Valid determines if a key is valid, according to a couple rules:
+//   - k is not nil
+//   - every token of k:
+//     - (if !allowSpecial) token's kind doesn't start with '__'
+//     - token's kind and appid are non-blank
+//     - token is not incomplete
+//   - all tokens have the same namespace and appid
+func (k *Key) Valid(allowSpecial bool, aid, ns string) bool {
+	if k == nil {
+		return false
+	}
+	if aid != k.AppID() || ns != k.Namespace() {
+		return false
+	}
+	for _, t := range k.toks {
+		if t.Incomplete() {
+			return false
+		}
+		if !allowSpecial && t.Special() {
+			return false
+		}
+		if t.Kind == "" {
+			return false
+		}
+		if t.StringID != "" && t.IntID != 0 {
+			return false
+		}
+	}
+	return true
+}
+
+// PartialValid returns true iff this key is suitable for use in a Put
+// operation. This is the same as Valid(k, false, ...), but also allowing k to
+// be Incomplete().
+func (k *Key) PartialValid(aid, ns string) bool {
+	if k.Incomplete() {
+		k = NewKey(k.AppID(), k.Namespace(), k.Last().Kind, "", 1, k.Parent())
+	}
+	return k.Valid(false, aid, ns)
+}
+
+// Parent returns the parent Key of this *Key, or nil. The parent
+// will always have the concrete type of *Key.
+func (k *Key) Parent() *Key {
+	if len(k.toks) <= 1 {
+		return nil
+	}
+	return &Key{k.appID, k.namespace, k.toks[:len(k.toks)-1]}
+}
+
+// MarshalJSON allows this key to be automatically marshaled by encoding/json.
+func (k *Key) MarshalJSON() ([]byte, error) {
+	return []byte(`"` + k.Encode() + `"`), nil
+}
+
+// Encode encodes the provided key as a base64-encoded protobuf.
+//
+// This encoding is compatible with the SDK-provided encoding and is agnostic
+// to the underlying implementation of the Key.
+//
+// It's encoded with the urlsafe base64 table.

[dnj, 2015/09/18 16:47:58]

... with padding stripped off.

[iannucci, 2015/09/18 22:25:48]

Done

+func (k *Key) Encode() string {
+	e := make([]*pb.Path_Element, len(k.toks))
+	for i, t := range k.toks {
+		t := t
+		e[i] = &pb.Path_Element{
+			Type: &t.Kind,
+		}
+		if t.StringID != "" {
+			e[i].Name = &t.StringID
+		} else {
+			e[i].Id = &t.IntID
+		}
+	}
+	var namespace *string
+	if k.namespace != "" {
+		namespace = &k.namespace
+	}
+	r, err := proto.Marshal(&pb.Reference{
+		App:       &k.appID,
+		NameSpace: namespace,
+		Path: &pb.Path{
+			Element: e,
+		},
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	// trim padding
+	return strings.TrimRight(base64.URLEncoding.EncodeToString(r), "=")
+}
+
+// UnmarshalJSON allows this key to be automatically unmarshaled by encoding/json.
+func (k *Key) UnmarshalJSON(buf []byte) error {
+	if len(buf) < 2 || buf[0] != '"' || buf[len(buf)-1] != '"' {
+		return errors.New("datastore: bad JSON key")
+	}
+	nk, err := NewKeyEncoded(string(buf[1 : len(buf)-1]))
+	if err != nil {
+		return err
+	}
+	*k = *nk
+	return nil
+}
+
+// Root returns the entity root for the given key.
+func (k *Key) Root() *Key {
+	if k != nil && len(k.toks) > 1 {
+		ret := *k
+		ret.toks = ret.toks[:1]
+		return &ret
+	}
+	return k
+}
+
+func (k *Key) Less(other *Key) bool {
+	if k == nil && other != nil {
+		return true
+	} else if k != nil && other == nil {
+		return false
+	}

[dnj, 2015/09/18 16:47:58]

They could still both be nil. Do:
if k == nil {
  return other != nil
}

[iannucci, 2015/09/18 22:25:48]

none of this exists anymore. You have nil you get boom.

+
+	if k.appID < other.appID {

[dnj, 2015/09/18 16:47:58]

Slightly better to avoid double string compare and do:
cmp := strings.Compare(k.appID, other.appID)
switch cmp {
  case -1:
    return true
  case 1:
    return false
}

"" for other string comparisons here.

[iannucci, 2015/09/18 22:25:48]

Which doesn't exist in 1.4 and whose 1.5 implementation does the exact same
thing as I do here: http://golang.org/src/strings/compare.go?s=491:520#L3

+		return true
+	} else if k.appID > other.appID {
+		return false
+	}
+
+	if k.namespace < other.namespace {
+		return true
+	} else if k.namespace > other.namespace {
+		return false
+	}
+
+	lim := len(k.toks)
+	if len(other.toks) < lim {
+		lim = len(other.toks)
+	}
+	for i := 0; i < lim; i++ {
+		a, b := k.toks[i], other.toks[i]
+		if a.Less(b) {
+			return true
+		} else if b.Less(a) {
+			return false
+		}
+	}
+	return len(k.toks) < len(other.toks)
+}
+
+// Equal returns true iff the two keys represent identical key values.
+func (k *Key) Equal(other *Key) (ret bool) {
+	ret = k == nil && other == nil
+	if !ret && k != nil && other != nil {
+		ret = (k.appID == other.appID &&
+			k.namespace == other.namespace &&
+			len(k.toks) == len(other.toks))
+		if ret {
+			for i, t := range k.toks {
+				if ret = t == other.toks[i]; !ret {
+					return
+				}
+			}
+		}
+	}
+	return
+}
+
+func (k *Key) Split() (appID, namespace string, toks []KeyTok) {
+	if k != nil {
+		appID = k.appID
+		namespace = k.namespace
+		toks = make([]KeyTok, len(k.toks))
+		copy(toks, k.toks)
+	}
+	return
+}