Refactor keys and queries in datastore service and implementation.

service/datastore/key.go

+// 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
+}