datastore: Update AllocateIDs to take keys.

service/datastore/datastore.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package datastore
 
 import (
         "fmt"
         "io"
         "io/ioutil"
@@ skipping 30 matching lines @@
 }
 
 func (d *datastoreImpl) MakeKey(elems ...interface{}) *Key {
         return MakeKey(d.aid, d.ns, elems...)
 }
 
 func (d *datastoreImpl) NewKey(kind, stringID string, intID int64, parent *Key) *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)
 }
 
 // PopulateKey loads key into obj.
 //
 // obj is any object that Interface.Get is able to accept.
 //
+// Upon successful application, this method will return true. If the key could
+// not be applied to the object, this method will return false. It will panic if
+// obj is an invalid datastore model.
+//
 // This method will panic if obj is an invalid datastore model. If the key could
 // not be applied to the object, nothing will happen.
-func PopulateKey(obj interface{}, key *Key) {
-        populateKeyMGS(getMGS(obj), key)
+func PopulateKey(obj interface{}, key *Key) bool {
+        return populateKeyMGS(getMGS(obj), key)
 }
 
-func populateKeyMGS(mgs MetaGetterSetter, key *Key) {
-        if !mgs.SetMeta("key", key) {
-                lst := key.LastTok()
-                if lst.StringID != "" {
-                        mgs.SetMeta("id", lst.StringID)
-                } else {
-                        mgs.SetMeta("id", lst.IntID)
+func populateKeyMGS(mgs MetaGetterSetter, key *Key) bool {
+        if mgs.SetMeta("key", key) {
+                return true
+        }
+
+        lst := key.LastTok()
+        if lst.StringID != "" {
+                if !mgs.SetMeta("id", lst.StringID) {
+                        return false
                 }
-                mgs.SetMeta("kind", lst.Kind)
-                mgs.SetMeta("parent", key.Parent())
+        } else {
+                if !mgs.SetMeta("id", lst.IntID) {
+                        return false
+                }
         }
+
+        mgs.SetMeta("kind", lst.Kind)
+        mgs.SetMeta("parent", key.Parent())
+        return true
 }
 
 func checkMultiSliceType(v interface{}) error {
         if reflect.TypeOf(v).Kind() == reflect.Slice {
                 return nil
         }
         return fmt.Errorf("argument must be a slice, not %T", v)
 
 }
 
@@ skipping 18 matching lines @@
 
         numIn := cbTyp.NumIn()
         if numIn != 1 && numIn != 2 {
                 badSig()
         }
 
         firstArg := cbTyp.In(0)
         if firstArg == typeOfKey {
                 isKey = true
         } else {
-                mat = mustParseArg(firstArg)
+                mat = mustParseArg(firstArg, false)
                 if mat.newElem == nil {
                         badSig()
                 }
         }
 
         hasCursorCB = numIn == 2
         if hasCursorCB && cbTyp.In(1) != typeOfCursorCB {
                 badSig()
         }
 
         if cbTyp.NumOut() > 1 {
                 badSig()
         } else if cbTyp.NumOut() == 1 && cbTyp.Out(0) != typeOfError {
                 badSig()
         }
         hasErr = cbTyp.NumOut() == 1
 
         return
 }
 
+func (d *datastoreImpl) AllocateIDs(ent ...interface{}) error {
+        if len(ent) == 0 {
+                return nil
+        }
+
+        mma, err := makeMetaMultiArg(ent, mmaWriteKeys)
+        if err != nil {
+                panic(err)
+        }
+
+        keys, _, err := mma.getKeysPMs(d.aid, d.ns, false)
+        if err != nil {
+                return err
+        }
+        if len(keys) == 0 {
+                return nil
+        }
+
+        // Convert each key to be partial valid, assigning an integer ID of 0. Confirm
+        // that each object can be populated with such a key.
+        for i, key := range keys {
+                keys[i] = key.Incomplete()
+        }
+
+        var et errorTracker
+        it := mma.iterator(et.init(mma))
+        err = filterStop(d.RawInterface.AllocateIDs(keys, func(key *Key, err error) error {
+                it.next(func(mat *multiArgType, v reflect.Value) error {
+                        if err != nil {
+                                return err
+                        }
+
+                        if !mat.setKey(v, key) {
+                                return ErrInvalidKey
+                        }
+                        return nil
+                })
+
+                return nil
+        }))
+        if err == nil {
+                err = et.error()
+
+                if err != nil && len(ent) == 1 {
+                        // Single-argument Exists will return a single error.
+                        err = errors.SingleError(err)
+                }
+        }
+        return err
+}
+
 func (d *datastoreImpl) Run(q *Query, cbIface interface{}) error {
         isKey, hasErr, hasCursorCB, mat := runParseCallback(cbIface)
 
         if isKey {
                 q = q.KeysOnly(true)
         }
         fq, err := q.Finalize()
         if err != nil {
                 return err
         }
@@ skipping 108 matching lines @@
                         for i, e := range errs {
                                 me[i] = e
                         }
                         err = me
                 }
         }
         return err
 }
 
 func (d *datastoreImpl) Exists(ent ...interface{}) (*ExistsResult, error) {
-        mma, err := makeMetaMultiArg(ent, true)
+        if len(ent) == 0 {
+                return nil, nil
+        }
+
+        mma, err := makeMetaMultiArg(ent, mmaKeysOnly)
         if err != nil {
                 panic(err)
         }
 
         keys, _, err := mma.getKeysPMs(d.aid, d.ns, false)
         if err != nil {
                 return nil, err
         }
+        if len(keys) == 0 {
+                return nil, nil
+        }
 
-        i := 0
         var bt boolTracker
         it := mma.iterator(bt.init(mma))
         err = filterStop(d.RawInterface.GetMulti(keys, nil, func(_ PropertyMap, err error) error {
                 it.next(func(*multiArgType, reflect.Value) error {
                         return err
                 })
-                i++
                 return nil
         }))
         if err == nil {
                 err = bt.error()
 
                 if err != nil && len(ent) == 1 {
                         // Single-argument Exists will return a single error.
                         err = errors.SingleError(err)
                 }
         }
         return bt.result(), err
 }
 
 func (d *datastoreImpl) ExistsMulti(keys []*Key) (BoolList, error) {
         v, err := d.Exists(keys)
         if err != nil {
                 return nil, err
         }
         return v.List(0), nil
 }
 
 func (d *datastoreImpl) Get(dst ...interface{}) (err error) {
-        mma, err := makeMetaMultiArg(dst, false)
+        if len(dst) == 0 {
+                return nil
+        }
+
+        mma, err := makeMetaMultiArg(dst, mmaReadWrite)
         if err != nil {
                 panic(err)
         }
 
         keys, pms, err := mma.getKeysPMs(d.aid, d.ns, true)
         if err != nil {
                 return err
         }
+        if len(keys) == 0 {
+                return nil
+        }
 
-        i := 0
         var et errorTracker
         it := mma.iterator(et.init(mma))
         meta := NewMultiMetaGetter(pms)
         err = filterStop(d.RawInterface.GetMulti(keys, meta, func(pm PropertyMap, err error) error {
                 it.next(func(mat *multiArgType, slot reflect.Value) error {
                         if err != nil {
                                 return err
                         }
                         return mat.setPM(slot, pm)
                 })
-
-                i++
                 return nil
         }))
 
         if err == nil {
                 err = et.error()
 
                 if err != nil && len(dst) == 1 {
                         // Single-argument Get will return a single error.
                         err = errors.SingleError(err)
                 }
         }
         return err
 }
 
 func (d *datastoreImpl) GetMulti(dst interface{}) error {
         if err := checkMultiSliceType(dst); err != nil {
                 panic(err)
         }
         return d.Get(dst)
 }
 
 func (d *datastoreImpl) Put(src ...interface{}) (err error) {
-        mma, err := makeMetaMultiArg(src, false)
+        if len(src) == 0 {
+                return nil
+        }
+
+        mma, err := makeMetaMultiArg(src, mmaReadWrite)
         if err != nil {
                 panic(err)
         }
 
         keys, vals, err := mma.getKeysPMs(d.aid, d.ns, false)
         if err != nil {
                 return err
         }
+        if len(keys) == 0 {
+                return nil
+        }
 
         i := 0
         var et errorTracker
         it := mma.iterator(et.init(mma))
         err = filterStop(d.RawInterface.PutMulti(keys, vals, func(key *Key, err error) error {
                 it.next(func(mat *multiArgType, slot reflect.Value) error {
                         if err != nil {
                                 return err
                         }
                         if key != keys[i] {
@@ skipping 18 matching lines @@
 }
 
 func (d *datastoreImpl) PutMulti(src interface{}) error {
         if err := checkMultiSliceType(src); err != nil {
                 panic(err)
         }
         return d.Put(src)
 }
 
 func (d *datastoreImpl) Delete(ent ...interface{}) error {
-        mma, err := makeMetaMultiArg(ent, true)
+        if len(ent) == 0 {
+                return nil
+        }
+
+        mma, err := makeMetaMultiArg(ent, mmaKeysOnly)
         if err != nil {
                 panic(err)
         }
 
         keys, _, err := mma.getKeysPMs(d.aid, d.ns, false)
         if err != nil {
                 return err
         }
+        if len(keys) == 0 {
+                return nil
+        }
 
-        i := 0
         var et errorTracker
         it := mma.iterator(et.init(mma))
         err = filterStop(d.RawInterface.DeleteMulti(keys, func(err error) error {
                 it.next(func(*multiArgType, reflect.Value) error {
                         return err
                 })
-                i++
 
                 return nil
         }))
         if err == nil {
                 err = et.error()
 
                 if err != nil && len(ent) == 1 {
                         // Single-argument Delete will return a single error.
                         err = errors.SingleError(err)
                 }
@@ skipping 91 matching lines @@
                 currentDir = filepath.Dir(currentDir)
         }
 }
 
 func filterStop(err error) error {
         if err == Stop {
                 err = nil
         }
         return err
 }