datastore: Update AllocateIDs to take keys.

service/datastore/key.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // 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
 }
 
-// Incomplete returns true iff this token doesn't define either a StringID or
+// IsIncomplete returns true iff this token doesn't define either a StringID or
 // an IntID.
-func (k KeyTok) Incomplete() bool {
+func (k KeyTok) IsIncomplete() bool {
         return k.StringID == "" && k.IntID == 0
 }
 
 // Special returns true iff this token begins and ends with "__"
 func (k KeyTok) Special() bool {
         return len(k.Kind) >= 2 && k.Kind[:2] == "__" && k.Kind[len(k.Kind)-2:] == "__"
 }
 
 // ID returns the 'active' id as a Property (either the StringID or the IntID).
 func (k KeyTok) ID() Property {
@@ skipping 170 matching lines @@
         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.LastTok().Incomplete()
+// IsIncomplete returns true iff k doesn't have an id yet.
+func (k *Key) IsIncomplete() bool {
+        return k.LastTok().IsIncomplete()
 }
 
 // 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 aid != k.appID || ns != k.namespace {
                 return false
         }
         for _, t := range k.toks {
-                if t.Incomplete() {
+                if t.IsIncomplete() {
                         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().
+// be IsIncomplete().
 func (k *Key) PartialValid(aid, ns string) bool {
-        if k.Incomplete() {
+        if k.IsIncomplete() {
                 k = NewKey(k.AppID(), k.Namespace(), k.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
@@ skipping 144 matching lines @@
                         fmt.Fprintf(ret, ", %s, %s", gqlQuoteString(t.Kind), gqlQuoteString(t.StringID))
                 }
         }
         if _, err := ret.WriteString(")"); err != nil {
                 panic(err)
         }
         return ret.String()
 }
 
 // Equal returns true iff the two keys represent identical key values.
-func (k *Key) Equal(other *Key) (ret bool) {
+func (k *Key) Equal(other *Key) bool {
+        return k.IncompleteEqual(other) && (k.LastTok() == other.LastTok())
+}
+
+// IncompleteEqual asserts that, were the two keys incomplete, they would be
+// equal.
+//
+// This asserts equality for the full lineage of the key, except for its last
+// token ID.
+func (k *Key) IncompleteEqual(other *Key) (ret bool) {
         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
+                        if i == len(k.toks)-1 {
+                                // Last token: check only Kind.
+                                if ret = (t.Kind == other.toks[i].Kind); !ret {
+                                        return
+                                }
+                        } else {
+                                if ret = t == other.toks[i]; !ret {
+                                        return
+                                }
                         }
                 }
         }
         return
 }
 
+// Incomplete returns an incomplete version of the key. The ID fields of the
+// last token will be set to zero/empty.
+func (k *Key) Incomplete() *Key {
+        if k.IsIncomplete() {
+                return k
+        }
+        return NewKey(k.appID, k.namespace, k.Kind(), "", 0, k.Parent())
+}
+
+// WithID returns the key generated by setting the ID of its last token to
+// the specified value.
+//
+// To generate this, k is reduced to its Incomplete form, then populated with a
+// new ID. The resulting key will have the same token linage as k (i.e., will
+// be IncompleteEqual).
+func (k *Key) WithID(stringID string, intID int64) *Key {
+        if k.StringID() == stringID && k.IntID() == intID {
+                return k
+        }
+        return NewKey(k.appID, k.namespace, k.Kind(), stringID, intID, k.Parent())
+}
+
 // Split componentizes the key into pieces (AppID, Namespace and tokens)
 //
 // Each token represents one piece of they key's 'path'.
 //
 // toks is guaranteed to be empty if and only if k is nil. If k is non-nil then
 // it contains at least one token.
 func (k *Key) Split() (appID, namespace string, toks []KeyTok) {
         appID = k.appID
         namespace = k.namespace
         toks = make([]KeyTok, len(k.toks))
@@ skipping 11 matching lines @@
         for _, t := range k.toks {
                 ret += int64(len(t.Kind))
                 if t.StringID != "" {
                         ret += int64(len(t.StringID))
                 } else {
                         ret += 8
                 }
         }
         return ret
 }