datastore: Update AllocateIDs to take keys.

service/datastore/checkfilter.go

Index: service/datastore/checkfilter.go
diff --git a/service/datastore/checkfilter.go b/service/datastore/checkfilter.go
index fe9bc57f99f3a3ea52e7f8d232b5de783e33090b..2e93755c17aaf1e8a4b0bea006ac3998d9d6f132 100644
--- a/service/datastore/checkfilter.go
+++ b/service/datastore/checkfilter.go
@@ -19,14 +19,38 @@ type checkFilter struct {
 	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
+	}
+
+	// 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
+			}
+		}
 	}
-	if !incomplete.PartialValid(tcf.aid, tcf.ns) {
-		return 0, ErrInvalidKey
+	if me := lme.Get(); me != nil {
+		return me
 	}
-	return tcf.RawInterface.AllocateIDs(incomplete, n)
+
+	return tcf.RawInterface.AllocateIDs(keys)
 }
 
 func (tcf *checkFilter) RunInTransaction(f func(c context.Context) error, opts *TransactionOptions) error {