datastore: Update AllocateIDs to take keys.

impl/prod/raw_datastore.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 prod
 
 import (
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/luci-go/common/errors"
         "golang.org/x/net/context"
@@ skipping 39 matching lines @@
         me, ok := err.(errors.MultiError)
         if ok {
                 for i, err := range me {
                         cb(i, err)
                 }
                 return nil
         }
         return err
 }
 
-func (d rdsImpl) AllocateIDs(incomplete *ds.Key, n int) (start int64, err error) {
+func (d rdsImpl) AllocateIDs(keys []*ds.Key) error {
+        // Allocate a set of IDs in our key namespace. We know that "keys" has at
+        // least one entry and that all keys share the same namespace and entty
+        // type because checkfilter has already asserted this.
+        incomplete := keys[0]
         par, err := dsF2R(d.aeCtx, incomplete.Parent())
         if err != nil {
-                return
+                return err
         }
 
-        start, _, err = datastore.AllocateIDs(d.aeCtx, incomplete.Kind(), par, n)
-        return
+        start, _, err := datastore.AllocateIDs(d.aeCtx, incomplete.Kind(), par, len(keys))
+        if err != nil {
+                return err
+        }
+
+        for i, k := range keys {
+                keys[i] = ds.NewKey(k.AppID(), k.Namespace(), k.Kind(), "", start+int64(i), k.Parent())
+        }
+        return nil
 }
 
 func (d rdsImpl) DeleteMulti(ks []*ds.Key, cb ds.DeleteMultiCB) error {
         keys, err := dsMF2R(d.aeCtx, ks)
         if err == nil {
                 err = datastore.DeleteMulti(d.aeCtx, keys)
         }
         return idxCallbacker(err, len(ks), func(_ int, err error) {
                 cb(err)
         })
@@ skipping 152 matching lines @@
 func (d rdsImpl) RunInTransaction(f func(c context.Context) error, opts *ds.TransactionOptions) error {
         ropts := (*datastore.TransactionOptions)(opts)
         return datastore.RunInTransaction(d.aeCtx, func(c context.Context) error {
                 return f(context.WithValue(d.userCtx, prodContextKey, c))
         }, ropts)
 }
 
 func (d rdsImpl) Testable() ds.Testable {
         return nil
 }