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

impl/memory/datastore.go

Index: impl/memory/datastore.go
diff --git a/impl/memory/datastore.go b/impl/memory/datastore.go
index efc61d3ba5f1a13ee1ac6bc8c89a04c206a508c9..4ebfac0d4a647b1a7933613b4833006021a337a0 100644
--- a/impl/memory/datastore.go
+++ b/impl/memory/datastore.go
@@ -19,11 +19,24 @@ import (
 // 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)
+
+		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)
+			}
 			return &dsImpl{x, ns, ic}
 		}
 		return &txnDsImpl{dsd.(*txnDataStoreData), ns}