datastore: Update AllocateIDs to take keys.

filter/count/rds.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 count
 
 import (
         "golang.org/x/net/context"
 
         ds "github.com/luci/gae/service/datastore"
@@ skipping 12 matching lines @@
 }
 
 type dsCounter struct {
         c *DSCounter
 
         ds ds.RawInterface
 }
 
 var _ ds.RawInterface = (*dsCounter)(nil)
 
-func (r *dsCounter) AllocateIDs(incomplete *ds.Key, n int) (int64, error) {
-        start, err := r.ds.AllocateIDs(incomplete, n)
-        return start, r.c.AllocateIDs.up(err)
+func (r *dsCounter) AllocateIDs(keys []*ds.Key, cb ds.PutMultiCB) error {
+        return r.c.AllocateIDs.up(r.ds.AllocateIDs(keys, cb))
 }
 
 func (r *dsCounter) DecodeCursor(s string) (ds.Cursor, error) {
         cursor, err := r.ds.DecodeCursor(s)
         return cursor, r.c.DecodeCursor.up(err)
 }
 
 func (r *dsCounter) Run(q *ds.FinalizedQuery, cb ds.RawRunCB) error {
         return r.c.Run.upFilterStop(r.ds.Run(q, cb))
 }
@@ skipping 35 matching lines @@
 // datastore.Stop will pass through this function, but, unlike other error
 // codes, will be counted as a success.
 func (e *Entry) upFilterStop(err error) error {
         upErr := err
         if upErr == ds.Stop {
                 upErr = nil
         }
         e.up(upErr)
         return err
 }