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

impl/prod/raw_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 prod
 
 import (
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/info"
         "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) ds.RawInterface {
-                return rdsImpl{ci, AEContext(ci), info.Get(ci).GetNamespace()}
+        return ds.SetRawFactory(c, func(ci context.Context, wantTxn bool) ds.RawInterface {
+                ns := info.Get(ci).GetNamespace()
+                maybeTxnCtx := AEContext(ci)
+
+                if wantTxn {
+                        return rdsImpl{ci, maybeTxnCtx, ns}
+                }
+                aeCtx := AEContextNoTX(ci)
+                if maybeTxnCtx != aeCtx {
+                        ci = context.WithValue(ci, prodContextKey, aeCtx)
+                        var err error
+                        ci, err = info.Get(ci).Namespace(ns)
+                        if err != nil {
+                                panic(err)
+                        }
+                }
+                return rdsImpl{ci, aeCtx, ns}
         })
 }
 
 ////////// Datastore
 
 type rdsImpl struct {
         // userCtx is the context that has the luci/gae services and user objects in
         // it.
         userCtx context.Context
 
@@ skipping 186 matching lines @@
         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 {
         ropts := (*datastore.TransactionOptions)(opts)
-        return datastore.RunInTransaction(d.aeCtx, func(aeCtx context.Context) error {
-                return f(context.WithValue(d.userCtx, prodContextKey, aeCtx))
+        return datastore.RunInTransaction(d.aeCtx, func(c context.Context) error {
+                return f(context.WithValue(d.userCtx, prodContextKey, c))
         }, ropts)
 }
 
 func (d rdsImpl) Testable() ds.Testable {
         return nil
 }