Interface update, per-method Contexts.

impl/prod/info.go

Index: impl/prod/info.go
diff --git a/impl/prod/info.go b/impl/prod/info.go
index 691f08186e6ab75d70a263a03c09c1ea337c4a83..94c8ba3501bb047f307142a77eefaf3b638b1947 100644
--- a/impl/prod/info.go
+++ b/impl/prod/info.go
@@ -40,12 +40,8 @@ func (g giImpl) AppID() string {
 func (g giImpl) FullyQualifiedAppID() string {
 	return getProbeCache(g.usrCtx).fqaid
 }
-func (g giImpl) GetNamespace() (string, bool) {
-	pc := getProbeCache(g.usrCtx)
-	if ns := pc.namespace; ns != nil {
-		return *ns, true
-	}
-	return "", false
+func (g giImpl) GetNamespace() string {
+	return getProbeCache(g.usrCtx).namespace
 }
 func (g giImpl) Datacenter() string {
 	return appengine.Datacenter(g.aeCtx)
@@ -72,14 +68,35 @@ func (g giImpl) ModuleName() (name string) {
 	return appengine.ModuleName(g.aeCtx)
 }
 func (g giImpl) Namespace(namespace string) (context.Context, error) {
-	aeCtx, err := appengine.Namespace(g.aeCtx, namespace)
+	c := g.usrCtx
+
+	pc := *getProbeCache(c)
+	if pc.namespace == namespace {
+		// Already using this namespace.
+		return c, nil
+	}
+	pc.namespace = namespace
+
+	// Apply the namespace to our retained GAE Contexts.
+	var err error
+	ps := getProdState(c)
+
+	// Apply to current GAE Context.
+	if ps.ctx, err = appengine.Namespace(ps.ctx, namespace); err != nil {
+		return c, err
+	}
+
+	// Apply to non-transactional Context. Since the previous one applied
+	// successfully, this must succeed.
+	ps.noTxnCtx, err = appengine.Namespace(ps.noTxnCtx, namespace)
 	if err != nil {
-		return g.usrCtx, err
+		panic(err)
 	}
-	usrCtx := context.WithValue(g.usrCtx, prodContextKey, aeCtx)
-	pc := *getProbeCache(usrCtx)
-	pc.namespace = &namespace
-	return withProbeCache(usrCtx, &pc), nil
+
+	// Update our user Context with the new namespace-imbued objects.
+	c = withProbeCache(c, &pc)
+	c = withProdState(c, ps)
+	return c, nil
 }
 func (g giImpl) PublicCertificates() ([]info.Certificate, error) {
 	certs, err := appengine.PublicCertificates(g.aeCtx)
@@ -113,34 +130,29 @@ func (g giImpl) VersionID() string {
 	return appengine.VersionID(g.aeCtx)
 }
 
-func (g giImpl) Testable() info.Testable {
-	return nil
-}
+func (g giImpl) GetTestable() info.Testable { return nil }
 
 type infoProbeCache struct {
-	namespace *string
+	namespace string
 	fqaid     string
 }
 
 func probe(aeCtx context.Context) *infoProbeCache {
 	probeKey := datastore.NewKey(aeCtx, "Kind", "id", 0, nil)
-	namespace := probeKey.Namespace()
 	ipb := infoProbeCache{
-		fqaid: probeKey.AppID(),
-	}
-	if namespace != "" {
-		ipb.namespace = &namespace
+		fqaid:     probeKey.AppID(),
+		namespace: probeKey.Namespace(),
 	}
 	return &ipb
 }
 
 func getProbeCache(c context.Context) *infoProbeCache {
-	if pc, ok := c.Value(probeCacheKey).(*infoProbeCache); ok {
+	if pc, ok := c.Value(&probeCacheKey).(*infoProbeCache); ok {
 		return pc
 	}
 	return nil
 }
 
 func withProbeCache(c context.Context, pc *infoProbeCache) context.Context {
-	return context.WithValue(c, probeCacheKey, pc)
+	return context.WithValue(c, &probeCacheKey, pc)
 }