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} |