datastore: Update AllocateIDs to take keys.

impl/prod/raw_datastore.go

Index: impl/prod/raw_datastore.go
diff --git a/impl/prod/raw_datastore.go b/impl/prod/raw_datastore.go
index 0a638eecfd53c1bea34e99f08d3bf8bfac54d71b..f0dd01374d21ce050e44254eeb3ed1ec86916580 100644
--- a/impl/prod/raw_datastore.go
+++ b/impl/prod/raw_datastore.go
@@ -57,14 +57,25 @@ func idxCallbacker(err error, amt int, cb func(idx int, err error)) error {
 	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 {