datastore: Update AllocateIDs to take keys.

service/datastore/multiarg.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"
         "reflect"
 
         "github.com/luci/luci-go/common/errors"
 )
 
+type metaMultiArgConstraints int
+
+const (
+        // mmaReadWrite allows a metaMultiArg to operate on any type that can be
+        // both read and written to.
+        mmaReadWrite metaMultiArgConstraints = iota
+        // mmaKeysOnly implies mmaReadWrite, with the further statement that the only
+        // operation that will be performed against the arguments will be key
+        // extraction.
+        mmaKeysOnly = iota
+        // mmaWriteKeys indicates that the caller is only going to write key
+        // values. This enables the same inputs as mmaReadWrite, but also allows
+        // []*Key.
+        mmaWriteKeys = iota
+)
+
+func (c metaMultiArgConstraints) allowSingleKey() bool {
+        return c == mmaKeysOnly
+}
+
+func (c metaMultiArgConstraints) keyOperationsOnly() bool {
+        return c >= mmaKeysOnly
+}
+
 type multiArgType struct {
         getMGS  func(slot reflect.Value) MetaGetterSetter
         getPLS  func(slot reflect.Value) PropertyLoadSaver
         newElem func() reflect.Value
 }
 
 func (mat *multiArgType) getKey(aid, ns string, slot reflect.Value) (*Key, error) {
         return newKeyObjErr(aid, ns, mat.getMGS(slot))
 }
 
 func (mat *multiArgType) getPM(slot reflect.Value) (PropertyMap, error) {
         return mat.getPLS(slot).Save(true)
 }
 
 func (mat *multiArgType) getMetaPM(slot reflect.Value) PropertyMap {
         return mat.getMGS(slot).GetAllMeta()
 }
 
 func (mat *multiArgType) setPM(slot reflect.Value, pm PropertyMap) error {
         return mat.getPLS(slot).Load(pm)
 }
 
-func (mat *multiArgType) setKey(slot reflect.Value, k *Key) {
-        PopulateKey(mat.getMGS(slot), k)
+func (mat *multiArgType) setKey(slot reflect.Value, k *Key) bool {
+        return populateKeyMGS(mat.getMGS(slot), k)
 }
 
 // parseArg checks that et is of type S, *S, I, P or *P, for some
 // struct type S, for some interface type I, or some non-interface non-pointer
 // type P such that P or *P implements PropertyLoadSaver.
 //
 // If et is a chan type that implements PropertyLoadSaver, new elements will be
 // allocated with a buffer of 0.
 //
-// If keysOnly is true, a read-only key extraction multiArgType will be
-// returned if et is a *Key.
-func parseArg(et reflect.Type, keysOnly bool) *multiArgType {
+// If allowKey is true, et may additional be type *Key. Only MetaGetterSetter
+// fields will be populated in the result (see keyMGS).
+func parseArg(et reflect.Type, allowKeys bool) *multiArgType {
         var mat multiArgType
 
-        if keysOnly && et == typeOfKey {
-                mat.getMGS = func(slot reflect.Value) MetaGetterSetter { return &keyExtractionMGS{key: slot.Interface().(*Key)} }
+        if et == typeOfKey {
+                if !allowKeys {
+                        return nil
+                }
+
+                mat.getMGS = func(slot reflect.Value) MetaGetterSetter { return &keyMGS{slot: slot} }
                 return &mat
         }
 
         // If we do identify a structCodec for this type, retain it so we don't
         // resolve it multiple times.
         var codec *structCodec
         initCodec := func(t reflect.Type) *structCodec {
                 if codec == nil {
                         codec = getCodec(t)
                 }
@@ skipping 128 matching lines @@
         return &mat
 }
 
 // mustParseMultiArg checks that v has type []S, []*S, []I, []P or []*P, for
 // some struct type S, for some interface type I, or some non-interface
 // non-pointer type P such that P or *P implements PropertyLoadSaver.
 func mustParseMultiArg(et reflect.Type) *multiArgType {
         if et.Kind() != reflect.Slice {
                 panic(fmt.Errorf("invalid argument type: expected slice, got %s", et))
         }
-        return mustParseArg(et.Elem())
+        return mustParseArg(et.Elem(), true)
 }
 
-func mustParseArg(et reflect.Type) *multiArgType {
+func mustParseArg(et reflect.Type, sliceArg bool) *multiArgType {
         if mat := parseArg(et, false); mat != nil {
                 return mat
         }
         panic(fmt.Errorf("invalid argument type: %s is not a PLS or pointer-to-struct", et))
 }
 
 func newKeyObjErr(aid, ns string, mgs MetaGetterSetter) (*Key, error) {
         if key, _ := GetMetaDefault(mgs, "key", nil).(*Key); key != nil {
                 return key, nil
         }
 
         // get kind
         kind := GetMetaDefault(mgs, "kind", "").(string)
         if kind == "" {
                 return nil, errors.New("unable to extract $kind")
         }
 
         // get id - allow both to be default for default keys
         sid := GetMetaDefault(mgs, "id", "").(string)
         iid := GetMetaDefault(mgs, "id", 0).(int64)
 
         // get parent
         par, _ := GetMetaDefault(mgs, "parent", nil).(*Key)
 
         return NewKey(aid, ns, kind, sid, iid, par), nil
 }
 
-func isOKSingleType(t reflect.Type, keysOnly bool) error {
+func isOKSingleType(t reflect.Type, allowKey bool) error {
         switch {
         case t == nil:
                 return errors.New("no type information")
         case t.Implements(typeOfPropertyLoadSaver):
                 return nil
-        case !keysOnly && t == typeOfKey:
+        case !allowKey && t == typeOfKey:
                 return errors.New("not user datatype")
 
         case t.Kind() != reflect.Ptr:
                 return errors.New("not a pointer")
         case t.Elem().Kind() != reflect.Struct:
                 return errors.New("does not point to a struct")
 
         default:
                 return nil
         }
 }
 
-// keyExtractionMGS is a MetaGetterSetter that wraps a key and only implements
-// GetMeta, and even then only retrieves the wrapped key meta value, "key".
-type keyExtractionMGS struct {
-        MetaGetterSetter
-
-        key *Key
+// keyMGS is a MetaGetterSetter that wraps a single key value/slot. It only
+// implements operations on the "key" key.
+//
+// GetMeta will be implemented, returning the *Key for the "key" meta.
+//
+// If slot is addressable, SetMeta will allow it to be set to the supplied
+// Value.
+type keyMGS struct {
+        slot reflect.Value
 }
 
-func (mgs *keyExtractionMGS) GetMeta(key string) (interface{}, bool) {
-        if key == "key" {
-                return mgs.key, true
+func (mgs *keyMGS) GetAllMeta() PropertyMap {
+        return PropertyMap{"$key": []Property{MkPropertyNI(mgs.slot.Interface())}}
+}
+
+func (mgs *keyMGS) GetMeta(key string) (interface{}, bool) {
+        if key != "key" {
+                return nil, false
         }
-        return nil, false
+        return mgs.slot.Interface(), true
+}
+
+func (mgs *keyMGS) SetMeta(key string, value interface{}) bool {
+        if !(key == "key" && mgs.slot.CanAddr()) {
+                return false
+        }
+        mgs.slot.Set(reflect.ValueOf(value))
+        return true
 }
 
 type metaMultiArgElement struct {
         arg  reflect.Value
         mat  *multiArgType
         size int // size is -1 if this element is not a slice.
 }
 
 type metaMultiArg struct {
         elems    []metaMultiArgElement
         keysOnly bool
 
         count int // total number of elements, flattening slices
 }
 
-func makeMetaMultiArg(args []interface{}, keysOnly bool) (*metaMultiArg, error) {
+// makeMetaMultiArg returns a metaMultiArg for the supplied args.
+//
+// If an arg is a slice, a slice metaMultiArg will be returned, and errors for
+// that slice will be written into a positional MultiError if they occur.
+//
+// If keysOnly is true, the caller is instructing metaMultiArg to only extract
+// the datastore *Key from args. *Key entries will be permitted, but the caller
+// may not write to them (since keys are read-only structs).
+func makeMetaMultiArg(args []interface{}, c metaMultiArgConstraints) (*metaMultiArg, error) {
         mma := metaMultiArg{
                 elems:    make([]metaMultiArgElement, len(args)),
-                keysOnly: keysOnly,
+                keysOnly: c.keyOperationsOnly(),
         }
 
         lme := errors.NewLazyMultiError(len(args))
         for i, arg := range args {
                 if arg == nil {
                         lme.Assign(i, errors.New("cannot use nil as single argument"))
                         continue
                 }
 
                 v := reflect.ValueOf(arg)
                 vt := v.Type()
                 mma.elems[i].arg = v
 
                 // Try and treat the argument as a single-value first. This allows slices
                 // that implement PropertyLoadSaver to be properly treated as a single
                 // element.
-                var err error
-                isSlice := false
-                mat := parseArg(vt, keysOnly)
-                if mat == nil {
+                var (
+                        mat     *multiArgType
+                        err     error
+                        isSlice = false
+                )
+                if mat = parseArg(vt, c.allowSingleKey()); mat == nil {
                         // If this is a slice, treat it as a slice of arg candidates.
+                        //
+                        // If only keys are being read/written, we allow a []*Key to be accepted
+                        // here, since slices are addressable (write).
                         if v.Kind() == reflect.Slice {
                                 isSlice = true
-                                mat = parseArg(vt.Elem(), keysOnly)
+                                mat = parseArg(vt.Elem(), c.keyOperationsOnly())
                         }
                 } else {
                         // Single types need to be able to be assigned to.
-                        err = isOKSingleType(vt, keysOnly)
+                        //
+                        // We only allow *Key here when the keys cannot be written to, since *Key
+                        // should not be modified in-place, as they are immutable.
+                        err = isOKSingleType(vt, c.allowSingleKey())
                 }
-                if mat == nil {
-                        err = errors.New("not a PLS or pointer-to-struct")
+                if mat == nil && err == nil {
+                        err = errors.New("not a PLS, pointer-to-struct, or slice thereof")
                 }
                 if err != nil {
                         lme.Assign(i, fmt.Errorf("invalid input type (%T): %s", arg, err))
                         continue
                 }
 
                 mma.elems[i].mat = mat
                 if isSlice {
                         l := v.Len()
                         mma.count += l
@@ skipping 181 matching lines @@
         default:
                 // Pass through to track as MultiError.
                 bt.errorTracker.trackError(it, err)
         }
 }
 
 func (bt *boolTracker) result() *ExistsResult {
         bt.res.updateSlices()
         return &bt.res
 }