datastore: Update AllocateIDs to take keys.

service/datastore/datastore.go

Index: service/datastore/datastore.go
diff --git a/service/datastore/datastore.go b/service/datastore/datastore.go
index 752c99c55779e02a2f97cced14b351645f18a6af..2287b1ebef05a153bbbbbf662a5c2b5cdf2dfd92 100644
--- a/service/datastore/datastore.go
+++ b/service/datastore/datastore.go
@@ -48,6 +48,16 @@ func (d *datastoreImpl) NewKey(kind, stringID string, intID int64, parent *Key)
 	return NewKey(d.aid, d.ns, kind, stringID, intID, parent)
 }
 
+func (d *datastoreImpl) NewIncompleteKeys(count int, kind string, parent *Key) (keys []*Key) {
+	if count > 0 {
+		keys = make([]*Key, count)
+		for i := range keys {
+			keys[i] = d.NewKey(kind, "", 0, parent)
+		}
+	}
+	return
+}
+
 func (d *datastoreImpl) NewKeyToks(toks []KeyTok) *Key {
 	return NewKeyToks(d.aid, d.ns, toks)
 }
@@ -121,6 +131,21 @@ func runParseCallback(cbIface interface{}) (isKey, hasErr, hasCursorCB bool, mat
 	return
 }
 
+func (d *datastoreImpl) AllocateIDs(keys ...*Key) ([]*Key, error) {
+	if len(keys) == 0 {
+		return nil, nil
+	}
+
+	// Duplicate the keys slice, then pass to our raw interface.
+	allocKeys := make([]*Key, len(keys))
+	copy(allocKeys, keys)
+
+	if err := d.RawInterface.AllocateIDs(allocKeys); err != nil {
+		return nil, err
+	}
+	return allocKeys, nil
+}
+
 func (d *datastoreImpl) Run(q *Query, cbIface interface{}) error {
 	isKey, hasErr, hasCursorCB, mat := runParseCallback(cbIface)