From 54bf9150c922186bfc45a00bf9dfcb91a5063275 Mon Sep 17 00:00:00 2001
From: David Symonds <dsymonds@golang.org>
Date: Mon, 4 Jan 2016 14:31:56 +1100
Subject: [PATCH] Refactor namespace application so it gets correctly applied
 implicitly.

Previously, a call to appengine.Namespace would work because it used
both internal.WithNamespace and internal.WithCallOverride. However,
appengine.NewContext was only using internal.WithNamespace and so API
calls made under an implicit namespace (such as when handling a
taskqueue PUSH task) would not properly namespace its messages.

Change-Id: I0075425ae1877ae0520c997f1182dec6cc337190
---
 internal/api.go         |  2 +-
 internal/api_classic.go |  2 +-
 internal/api_common.go  | 22 +++++++++++++++++++++-
 namespace.go            | 19 +------------------
 4 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/internal/api.go b/internal/api.go
index 047793e..aa139d4 100644
--- a/internal/api.go
+++ b/internal/api.go
@@ -231,7 +231,7 @@ func fromContext(ctx netcontext.Context) *context {
 func withContext(parent netcontext.Context, c *context) netcontext.Context {
 	ctx := netcontext.WithValue(parent, &contextKey, c)
 	if ns := c.req.Header.Get(curNamespaceHeader); ns != "" {
-		ctx = WithNamespace(ctx, ns)
+		ctx = withNamespace(ctx, ns)
 	}
 	return ctx
 }
diff --git a/internal/api_classic.go b/internal/api_classic.go
index b40dd70..1c072e9 100644
--- a/internal/api_classic.go
+++ b/internal/api_classic.go
@@ -37,7 +37,7 @@ func withContext(parent netcontext.Context, c appengine.Context) netcontext.Cont
 	s := &basepb.StringProto{}
 	c.Call("__go__", "GetNamespace", &basepb.VoidProto{}, s, nil)
 	if ns := s.GetValue(); ns != "" {
-		ctx = WithNamespace(ctx, ns)
+		ctx = NamespacedContext(ctx, ns)
 	}
 
 	return ctx
diff --git a/internal/api_common.go b/internal/api_common.go
index 51c4dd4..ec5383e 100644
--- a/internal/api_common.go
+++ b/internal/api_common.go
@@ -52,7 +52,7 @@ func WithAppIDOverride(ctx netcontext.Context, appID string) netcontext.Context
 
 var namespaceKey = "holds the namespace string"
 
-func WithNamespace(ctx netcontext.Context, ns string) netcontext.Context {
+func withNamespace(ctx netcontext.Context, ns string) netcontext.Context {
 	return netcontext.WithValue(ctx, &namespaceKey, ns)
 }
 
@@ -79,3 +79,23 @@ func Logf(ctx netcontext.Context, level int64, format string, args ...interface{
 	}
 	logf(fromContext(ctx), level, format, args...)
 }
+
+// NamespacedContext wraps a Context to support namespaces.
+func NamespacedContext(ctx netcontext.Context, namespace string) netcontext.Context {
+	n := &namespacedContext{
+		namespace: namespace,
+	}
+	return withNamespace(WithCallOverride(ctx, n.call), namespace)
+}
+
+type namespacedContext struct {
+	namespace string
+}
+
+func (n *namespacedContext) call(ctx netcontext.Context, service, method string, in, out proto.Message) error {
+	// Apply any namespace mods.
+	if mod, ok := NamespaceMods[service]; ok {
+		mod(in, n.namespace)
+	}
+	return Call(ctx, service, method, in, out)
+}
diff --git a/namespace.go b/namespace.go
index 65b0a00..21860ca 100644
--- a/namespace.go
+++ b/namespace.go
@@ -8,7 +8,6 @@ import (
 	"fmt"
 	"regexp"
 
-	"github.com/golang/protobuf/proto"
 	"golang.org/x/net/context"
 
 	"google.golang.org/appengine/internal"
@@ -19,24 +18,8 @@ func Namespace(c context.Context, namespace string) (context.Context, error) {
 	if !validNamespace.MatchString(namespace) {
 		return nil, fmt.Errorf("appengine: namespace %q does not match /%s/", namespace, validNamespace)
 	}
-	n := &namespacedContext{
-		namespace: namespace,
-	}
-	return internal.WithNamespace(internal.WithCallOverride(c, n.call), namespace), nil
+	return internal.NamespacedContext(c, namespace), nil
 }
 
 // validNamespace matches valid namespace names.
 var validNamespace = regexp.MustCompile(`^[0-9A-Za-z._-]{0,100}$`)
-
-// namespacedContext wraps a Context to support namespaces.
-type namespacedContext struct {
-	namespace string
-}
-
-func (n *namespacedContext) call(ctx context.Context, service, method string, in, out proto.Message) error {
-	// Apply any namespace mods.
-	if mod, ok := internal.NamespaceMods[service]; ok {
-		mod(in, n.namespace)
-	}
-	return internal.Call(ctx, service, method, in, out)
-}