Add API to allow you to get the non-transactional datastore or taskqueue.

impl/memory/datastore.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memory
 
 import (
         "errors"
         "fmt"
 
         "golang.org/x/net/context"
 
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/info"
 )
 
 //////////////////////////////////// public ////////////////////////////////////
 
 // useRDS adds a gae.Datastore implementation to context, accessible
 // by gae.GetDS(c)
 func useRDS(c context.Context) context.Context {
-        return ds.SetRawFactory(c, func(ic context.Context) ds.RawInterface {
-                dsd := cur(ic).Get(memContextDSIdx)
+        return ds.SetRawFactory(c, func(ic context.Context, wantTxn bool) ds.RawInterface {
+                ns := curGID(ic).namespace
+                maybeTxnCtx := cur(ic)
 
-                ns := curGID(ic).namespace
+                needResetCtx := false
+                if !wantTxn {
+                        rootctx := curNoTxn(ic)
+                        if rootctx != maybeTxnCtx {
+                                needResetCtx = true
+                                maybeTxnCtx = rootctx
+                        }
+                }
+
+                dsd := maybeTxnCtx.Get(memContextDSIdx)
                 if x, ok := dsd.(*dataStoreData); ok {
+                        if needResetCtx {
+                                ic = context.WithValue(ic, memContextKey, maybeTxnCtx)

[iannucci, 2015/12/04 03:35:10]

this is necessary so that .RunInTransaction will fork off of the root context
instead of trying to fork off of the already-transactional context.

+                        }
                         return &dsImpl{x, ns, ic}
                 }
                 return &txnDsImpl{dsd.(*txnDataStoreData), ns}
         })
 }
 
 // NewDatastore creates a new standalone memory implementation of the datastore,
 // suitable for embedding for doing in-memory data organization.
 //
 // It's configured by default with the following settings:
@@ skipping 168 matching lines @@
         return countQuery(fq, d.data.parent.aid, d.ns, true, d.data.snap, d.data.snap)
 }
 
 func (*txnDsImpl) RunInTransaction(func(c context.Context) error, *ds.TransactionOptions) error {
         return errors.New("datastore: nested transactions are not supported")
 }
 
 func (*txnDsImpl) Testable() ds.Testable {
         return nil
 }