datastore: Update AllocateIDs to take keys.

service/datastore/checkfilter.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 (
         "fmt"
 
         "github.com/luci/gae/service/info"
         "github.com/luci/luci-go/common/errors"
         "golang.org/x/net/context"
 )
 
 type checkFilter struct {
         RawInterface
 
         aid string
         ns  string
 }
 
-func (tcf *checkFilter) AllocateIDs(incomplete *Key, n int) (start int64, err error) {
-        if n <= 0 {
-                return 0, fmt.Errorf("datastore: invalid `n` parameter in AllocateIDs: %d", n)
+func (tcf *checkFilter) AllocateIDs(keys []*Key) error {
+        if len(keys) == 0 {
+                return nil
         }
-        if !incomplete.PartialValid(tcf.aid, tcf.ns) {
-                return 0, ErrInvalidKey
+
+        // Assert that all of the supplied keys are PartialValid, and that they all
+        // share the same entity.
+        //
+        // TODO: If GAE protobufs and/or API ever relaxes the same-kind requirement,
+        // we can drop that here and in the raw interface API too.
+        lme := errors.NewLazyMultiError(len(keys))
+        var exemplar *Key
+        for i, k := range keys {
+                if !k.PartialValid(tcf.aid, tcf.ns) {
+                        lme.Assign(i, ErrInvalidKey)
+                        continue
+                }
+
+                if exemplar == nil {
+                        exemplar = k
+                } else {
+                        if !k.SameKind(exemplar) {
+                                lme.Assign(i, ErrInvalidKey)
+                                continue
+                        }
+                }
         }
-        return tcf.RawInterface.AllocateIDs(incomplete, n)
+        if me := lme.Get(); me != nil {
+                return me
+        }
+
+        return tcf.RawInterface.AllocateIDs(keys)
 }
 
 func (tcf *checkFilter) RunInTransaction(f func(c context.Context) error, opts *TransactionOptions) error {
         if f == nil {
                 return fmt.Errorf("datastore: RunInTransaction function is nil")
         }
         return tcf.RawInterface.RunInTransaction(f, opts)
 }
 
 func (tcf *checkFilter) Run(fq *FinalizedQuery, cb RawRunCB) error {
@@ skipping 79 matching lines @@
                 return nil
         }
         return tcf.RawInterface.DeleteMulti(keys, cb)
 }
 
 func applyCheckFilter(c context.Context, i RawInterface) RawInterface {
         inf := info.Get(c)
         ns, _ := inf.GetNamespace()
         return &checkFilter{i, inf.FullyQualifiedAppID(), ns}
 }