Interface update, per-method Contexts.

impl/cloud/datastore.go

 // Copyright 2016 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 cloud
 
 import (
         "fmt"
         "reflect"
         "strings"
         "time"
 
         "github.com/luci/luci-go/common/errors"
 
+        ds "github.com/luci/gae/service/datastore"
+        "github.com/luci/gae/service/info"
+
         "cloud.google.com/go/datastore"
-        ds "github.com/luci/gae/service/datastore"
-        infoS "github.com/luci/gae/service/info"
 
         "golang.org/x/net/context"
 )
 
 type cloudDatastore struct {
         client *datastore.Client
 }
 
 func (cds *cloudDatastore) use(c context.Context) context.Context {
-        return ds.SetRawFactory(c, func(ic context.Context, wantTxn bool) ds.RawInterface {
-                inf := infoS.Get(ic)
-                if ns, ok := inf.GetNamespace(); ok {
+        return ds.SetRawFactory(c, func(ic context.Context) ds.RawInterface {
+                if ns := info.GetNamespace(ic); ns != "" {
                         ic = datastore.WithNamespace(ic, ns)
                 }
 
-                bds := boundDatastore{
+                return &boundDatastore{
                         Context:        ic,
                         cloudDatastore: cds,
-                        appID:          inf.FullyQualifiedAppID(),
+                        transaction:    datastoreTransaction(ic),
+                        kc:             ds.GetKeyContext(ic),
                 }
-                if wantTxn {
-                        bds.transaction = datastoreTransaction(ic)
-                }
-                return &bds
         })
 }
 
 // boundDatastore is a bound instance of the cloudDatastore installed in the
 // Context.
 type boundDatastore struct {
+        context.Context
+
         // Context is the bound user Context. It includes the datastore namespace, if
         // one is set.
-        context.Context
         *cloudDatastore
 
-        appID       string
         transaction *datastore.Transaction
+        kc          ds.KeyContext
 }
 
 func (bds *boundDatastore) AllocateIDs(keys []*ds.Key, cb ds.NewKeyCB) error {
         nativeKeys, err := bds.client.AllocateIDs(bds, bds.gaeKeysToNative(keys...))
         if err != nil {
                 return normalizeError(err)
         }
 
         keys = bds.nativeKeysToGAE(nativeKeys...)
         for _, key := range keys {
@@ skipping 47 matching lines @@
                         npls = bds.mkNPLS(nil)
                 }
                 nativeKey, err := it.Next(npls)
                 if err != nil {
                         if err == datastore.Done {
                                 return nil
                         }
                         return normalizeError(err)
                 }
 
-                if err := cb(bds.nativeKeysToGAE(nativeKey)[0], npls.pmap, cursorFn); err != nil {
+                var pmap ds.PropertyMap
+                if npls != nil {
+                        pmap = npls.pmap
+                }
+                if err := cb(bds.nativeKeysToGAE(nativeKey)[0], pmap, cursorFn); err != nil {
                         if err == ds.Stop {
                                 return nil
                         }
                         return normalizeError(err)
                 }
         }
 }
 
 func (bds *boundDatastore) Count(q *ds.FinalizedQuery) (int64, error) {
         v, err := bds.client.Count(bds, bds.prepareNativeQuery(q))
@@ skipping 26 matching lines @@
 }
 
 func (bds *boundDatastore) GetMulti(keys []*ds.Key, _meta ds.MultiMetaGetter, cb ds.GetMultiCB) error {
         nativeKeys := bds.gaeKeysToNative(keys...)
         nativePLS := make([]*nativePropertyLoadSaver, len(nativeKeys))
         for i := range nativePLS {
                 nativePLS[i] = bds.mkNPLS(nil)
         }
 
         var err error
-        if tx := bds.transaction; tx != nil {
+        if bds.transaction != nil {
                 // Transactional GetMulti.
-                err = tx.GetMulti(nativeKeys, nativePLS)
+                err = bds.transaction.GetMulti(nativeKeys, nativePLS)
         } else {
                 // Non-transactional GetMulti.
                 err = bds.client.GetMulti(bds, nativeKeys, nativePLS)
         }
 
         return idxCallbacker(err, len(nativePLS), func(idx int, err error) error {
                 return cb(nativePLS[idx].pmap, err)
         })
 }
 
 func (bds *boundDatastore) PutMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) error {
         nativeKeys := bds.gaeKeysToNative(keys...)
         nativePLS := make([]*nativePropertyLoadSaver, len(vals))
         for i := range nativePLS {
                 nativePLS[i] = bds.mkNPLS(vals[i])
         }
 
         var err error
-        if tx := bds.transaction; tx != nil {
+        if bds.transaction != nil {
                 // Transactional PutMulti.
                 //
                 // In order to simulate the presence of mid-transaction key allocation, we
                 // will identify any incomplete keys and allocate IDs for them. This is
                 // potentially wasteful in the event of failed or retried transactions, but
                 // it is required to maintain API compatibility with the datastore
                 // interface.
                 var incompleteKeys []*datastore.Key
                 var incompleteKeyMap map[int]int
                 for i, k := range nativeKeys {
@@ skipping 11 matching lines @@
                 if len(incompleteKeys) > 0 {
                         idKeys, err := bds.client.AllocateIDs(bds, incompleteKeys)
                         if err != nil {
                                 return err
                         }
                         for i, idKey := range idKeys {
                                 nativeKeys[incompleteKeyMap[i]] = idKey
                         }
                 }
 
-                _, err = tx.PutMulti(nativeKeys, nativePLS)
+                _, err = bds.transaction.PutMulti(nativeKeys, nativePLS)
         } else {
                 // Non-transactional PutMulti.
                 nativeKeys, err = bds.client.PutMulti(bds, nativeKeys, nativePLS)
         }
 
         return idxCallbacker(err, len(nativeKeys), func(idx int, err error) error {
                 if err == nil {
                         return cb(bds.nativeKeysToGAE(nativeKeys[idx])[0], nil)
                 }
                 return cb(nil, err)
         })
 }
 
 func (bds *boundDatastore) DeleteMulti(keys []*ds.Key, cb ds.DeleteMultiCB) error {
         nativeKeys := bds.gaeKeysToNative(keys...)
 
         var err error
-        if tx := bds.transaction; tx != nil {
+        if bds.transaction != nil {
                 // Transactional DeleteMulti.
-                err = tx.DeleteMulti(nativeKeys)
+                err = bds.transaction.DeleteMulti(nativeKeys)
         } else {
                 // Non-transactional DeleteMulti.
                 err = bds.client.DeleteMulti(bds, nativeKeys)
         }
 
         return idxCallbacker(err, len(nativeKeys), func(_ int, err error) error {
                 return cb(err)
         })
 }
 
-func (bds *boundDatastore) Testable() ds.Testable {
-        return nil
+func (bds *boundDatastore) WithoutTransaction() context.Context {
+        return withDatastoreTransaction(bds, nil)
 }
 
+func (bds *boundDatastore) CurrentTransaction() ds.Transaction { return bds.transaction }
+
+func (bds *boundDatastore) GetTestable() ds.Testable { return nil }
+
 func (bds *boundDatastore) prepareNativeQuery(fq *ds.FinalizedQuery) *datastore.Query {
         nq := datastore.NewQuery(fq.Kind())
         if bds.transaction != nil {
                 nq = nq.Transaction(bds.transaction)
         }
 
         // nativeFilter translates a filter field. If the translation fails, we'll
         // pass the result through to the underlying datastore and allow it to
         // reject it.
         nativeFilter := func(prop ds.Property) interface{} {
@@ skipping 56 matching lines @@
                 if ic.Descending {
                         prop = "-" + prop
                 }
                 nq = nq.Order(prop)
         }
 
         return nq
 }
 
 func (bds *boundDatastore) mkNPLS(base ds.PropertyMap) *nativePropertyLoadSaver {
-        return &nativePropertyLoadSaver{bds: bds, pmap: clonePropertyMap(base)}
+        return &nativePropertyLoadSaver{
+                bds:  bds,
+                pmap: clonePropertyMap(base),
+        }
 }
 
-func (bds *boundDatastore) gaePropertyToNative(name string, pdata ds.PropertyData) (nativeProp datastore.Property, err error) {
+func (bds *boundDatastore) gaePropertyToNative(name string, pdata ds.PropertyData) (
+        nativeProp datastore.Property, err error) {
+
         nativeProp.Name = name
 
         convert := func(prop *ds.Property) (interface{}, error) {
                 switch pt := prop.Type(); pt {
                 case ds.PTNull, ds.PTInt, ds.PTTime, ds.PTBool, ds.PTBytes, ds.PTString, ds.PTFloat:
                         return prop.Value(), nil
 
                 case ds.PTKey:
                         return bds.gaeKeysToNative(prop.Value().(*ds.Key))[0], nil
 
@@ skipping 28 matching lines @@
                 }
                 nativeProp.Value = multiProp
 
         default:
                 err = fmt.Errorf("unsupported PropertyData type for %q: %T", name, pdata)
         }
 
         return
 }
 
-func (bds *boundDatastore) nativePropertyToGAE(nativeProp datastore.Property) (name string, pdata ds.PropertyData, err error) {
+func (bds *boundDatastore) nativePropertyToGAE(nativeProp datastore.Property) (
+        name string, pdata ds.PropertyData, err error) {
+
         name = nativeProp.Name
 
         convert := func(nv interface{}, prop *ds.Property) error {
                 switch nvt := nv.(type) {
                 case nil:
                         nv = nil
 
                 case int64, bool, string, float64, []byte:
                         break
 
@@ skipping 52 matching lines @@
                         nativeKey = datastore.NewKey(bds, tok.Kind, tok.StringID, tok.IntID, nativeKey)
                 }
                 nativeKeys[i] = nativeKey
         }
         return nativeKeys
 }
 
 func (bds *boundDatastore) nativeKeysToGAE(nativeKeys ...*datastore.Key) []*ds.Key {
         keys := make([]*ds.Key, len(nativeKeys))
         toks := make([]ds.KeyTok, 1)
+
+        kc := bds.kc
         for i, nativeKey := range nativeKeys {
                 toks = toks[:0]
                 cur := nativeKey
                 for {
                         toks = append(toks, ds.KeyTok{Kind: cur.Kind(), IntID: cur.ID(), StringID: cur.Name()})
                         cur = cur.Parent()
                         if cur == nil {
                                 break
                         }
                 }
 
                 // Reverse "toks" so we have ancestor-to-child lineage.
                 for i := 0; i < len(toks)/2; i++ {
                         ri := len(toks) - i - 1
                         toks[i], toks[ri] = toks[ri], toks[i]
                 }
-                keys[i] = ds.NewKeyToks(bds.appID, nativeKey.Namespace(), toks)
+                kc.Namespace = nativeKey.Namespace()
+                keys[i] = kc.NewKeyToks(toks)
         }
         return keys
 }
 
 // nativePropertyLoadSaver is a ds.PropertyMap which implements
 // datastore.PropertyLoadSaver.
 //
 // It naturally converts between native and GAE properties and values.
 type nativePropertyLoadSaver struct {
         bds  *boundDatastore
@@ skipping 72 matching lines @@
         case datastore.ErrNoSuchEntity:
                 return ds.ErrNoSuchEntity
         case datastore.ErrConcurrentTransaction:
                 return ds.ErrConcurrentTransaction
         case datastore.ErrInvalidKey:
                 return ds.ErrInvalidKey
         default:
                 return err
         }
 }