Interface update, per-method Contexts.

impl/prod/raw_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 prod
 
 import (
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/luci-go/common/errors"
         "golang.org/x/net/context"
         "google.golang.org/appengine/datastore"
 )
 
 // useRDS adds a gae.RawDatastore implementation to context, accessible
 // by gae.GetDS(c)
 func useRDS(c context.Context) context.Context {
-        return ds.SetRawFactory(c, func(ci context.Context, wantTxn bool) ds.RawInterface {
-                maybeTxnCtx := AEContext(ci)
-
-                if wantTxn {
-                        return rdsImpl{ci, maybeTxnCtx}
+        return ds.SetRawFactory(c, func(ci context.Context) ds.RawInterface {
+                rds := rdsImpl{
+                        userCtx: ci,
+                        ps:      getProdState(ci),
                 }
-                aeCtx := AEContextNoTxn(ci)
-                if maybeTxnCtx != aeCtx {
-                        ci = context.WithValue(ci, prodContextKey, aeCtx)
-                }
-                return rdsImpl{ci, aeCtx}
+                rds.aeCtx = rds.ps.context(ci)
+                return &rds
         })
 }
 
 ////////// Datastore
 
 type rdsImpl struct {
         // userCtx is the context that has the luci/gae services and user objects in
         // it.
         userCtx context.Context
 
-        // aeCtx is the context with the appengine connection information in it.
+        // aeCtx is the AppEngine Context that will be used in method calls. This is
+        // derived from ps.
         aeCtx context.Context
+
+        // ps is the current production state.
+        ps prodState
 }
 
 func idxCallbacker(err error, amt int, cb func(idx int, err error) error) error {
         if err == nil {
                 for i := 0; i < amt; i++ {
                         if err := cb(i, nil); err != nil {
                                 return err
                         }
                 }
                 return nil
         }
         err = errors.Fix(err)
         me, ok := err.(errors.MultiError)
         if ok {
                 for i, err := range me {
                         if err := cb(i, err); err != nil {
                                 return err
                         }
                 }
                 return nil
         }
         return err
 }
 
-func (d rdsImpl) AllocateIDs(keys []*ds.Key, cb ds.NewKeyCB) error {
+func (d *rdsImpl) AllocateIDs(keys []*ds.Key, cb ds.NewKeyCB) error {
         // Map keys by entity type.
         entityMap := make(map[string][]int)
         for i, key := range keys {
                 ks := key.String()
                 entityMap[ks] = append(entityMap[ks], i)
         }
 
         // Allocate a set of IDs for each unique entity type.
         errors := errors.NewLazyMultiError(len(keys))
         setErrs := func(idxs []int, err error) {
@@ skipping 24 matching lines @@
         for i, key := range keys {
                 if err := errors.GetOne(i); err != nil {
                         cb(nil, err)
                 } else {
                         cb(key, nil)
                 }
         }
         return nil
 }
 
-func (d rdsImpl) DeleteMulti(ks []*ds.Key, cb ds.DeleteMultiCB) error {
+func (d *rdsImpl) DeleteMulti(ks []*ds.Key, cb ds.DeleteMultiCB) error {
         keys, err := dsMF2R(d.aeCtx, ks)
         if err == nil {
                 err = datastore.DeleteMulti(d.aeCtx, keys)
         }
         return idxCallbacker(err, len(ks), func(_ int, err error) error {
                 return cb(err)
         })
 }
 
-func (d rdsImpl) GetMulti(keys []*ds.Key, _meta ds.MultiMetaGetter, cb ds.GetMultiCB) error {
+func (d *rdsImpl) GetMulti(keys []*ds.Key, _meta ds.MultiMetaGetter, cb ds.GetMultiCB) error {
         vals := make([]datastore.PropertyLoadSaver, len(keys))
         rkeys, err := dsMF2R(d.aeCtx, keys)
         if err == nil {
                 for i := range keys {
                         vals[i] = &typeFilter{d.aeCtx, ds.PropertyMap{}}
                 }
                 err = datastore.GetMulti(d.aeCtx, rkeys, vals)
         }
         return idxCallbacker(err, len(keys), func(idx int, err error) error {
                 if pls := vals[idx]; pls != nil {
                         return cb(pls.(*typeFilter).pm, err)
                 }
                 return cb(nil, err)
         })
 }
 
-func (d rdsImpl) PutMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) error {
+func (d *rdsImpl) PutMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) error {
         rkeys, err := dsMF2R(d.aeCtx, keys)
         if err == nil {
                 rvals := make([]datastore.PropertyLoadSaver, len(vals))
                 for i, val := range vals {
                         rvals[i] = &typeFilter{d.aeCtx, val}
                 }
                 rkeys, err = datastore.PutMulti(d.aeCtx, rkeys, rvals)
         }
         return idxCallbacker(err, len(keys), func(idx int, err error) error {
                 k := (*ds.Key)(nil)
                 if err == nil {
                         k = dsR2F(rkeys[idx])
                 }
                 return cb(k, err)
         })
 }
 
-func (d rdsImpl) fixQuery(fq *ds.FinalizedQuery) (*datastore.Query, error) {
+func (d *rdsImpl) fixQuery(fq *ds.FinalizedQuery) (*datastore.Query, error) {
         ret := datastore.NewQuery(fq.Kind())
 
         start, end := fq.Bounds()
         if start != nil {
                 ret = ret.Start(start.(datastore.Cursor))
         }
         if end != nil {
                 ret = ret.End(end.(datastore.Cursor))
         }
 
@@ skipping 54 matching lines @@
         }
 
         ret = ret.Project(fq.Project()...)
         if fq.Distinct() {
                 ret = ret.Distinct()
         }
 
         return ret, nil
 }
 
-func (d rdsImpl) DecodeCursor(s string) (ds.Cursor, error) {
+func (d *rdsImpl) DecodeCursor(s string) (ds.Cursor, error) {
         return datastore.DecodeCursor(s)
 }
 
-func (d rdsImpl) Run(fq *ds.FinalizedQuery, cb ds.RawRunCB) error {
+func (d *rdsImpl) Run(fq *ds.FinalizedQuery, cb ds.RawRunCB) error {
         q, err := d.fixQuery(fq)
         if err != nil {
                 return err
         }
 
         t := q.Run(d.aeCtx)
 
         cfunc := func() (ds.Cursor, error) {
                 return t.Cursor()
         }
         tf := typeFilter{}
         for {
                 k, err := t.Next(&tf)
                 if err == datastore.Done {
                         return nil
                 }
                 if err != nil {
                         return err
                 }
                 if err := cb(dsR2F(k), tf.pm, cfunc); err != nil {
                         return err
                 }
         }
 }
 
-func (d rdsImpl) Count(fq *ds.FinalizedQuery) (int64, error) {
+func (d *rdsImpl) Count(fq *ds.FinalizedQuery) (int64, error) {
         q, err := d.fixQuery(fq)
         if err != nil {
                 return 0, err
         }
         ret, err := q.Count(d.aeCtx)
         return int64(ret), err
 }
 
-func (d rdsImpl) RunInTransaction(f func(c context.Context) error, opts *ds.TransactionOptions) error {
+func (d *rdsImpl) RunInTransaction(f func(c context.Context) error, opts *ds.TransactionOptions) error {
         ropts := (*datastore.TransactionOptions)(opts)
         return datastore.RunInTransaction(d.aeCtx, func(c context.Context) error {
-                return f(context.WithValue(d.userCtx, prodContextKey, c))
+                // Derive a prodState with this transaction Context.
+                ps := d.ps
+                ps.ctx = c
+                ps.inTxn = true
+
+                c = withProdState(d.userCtx, ps)
+                return f(c)
         }, ropts)
 }
 
-func (d rdsImpl) Testable() ds.Testable {
+func (d *rdsImpl) WithoutTransaction() context.Context {
+        c := d.userCtx
+        if d.ps.inTxn {
+                // We're in a transaction. Reset to non-transactional state.
+                ps := d.ps
+                ps.ctx = ps.noTxnCtx
+                ps.inTxn = false
+                c = withProdState(c, ps)
+        }
+        return c
+}
+
+func (d *rdsImpl) CurrentTransaction() ds.Transaction {
+        if d.ps.inTxn {
+                // Since we don't distinguish between transactions (yet), we just need this
+                // to be non-nil.
+                return struct{}{}
+        }
         return nil
 }
+
+func (d *rdsImpl) GetTestable() ds.Testable {
+        return nil
+}