Chromium Code Reviews| 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) |
|
iannucci
2015/12/04 03:35:10
this is necessary so that .RunInTransaction will f
|
| + } |
| return &dsImpl{x, ns, ic} |
| } |
| return &txnDsImpl{dsd.(*txnDataStoreData), ns} |