namespaces: Add NamespaceFromEnv

Signed-off-by: Samuel Karp <skarp@amazon.com>

namespaces/context.go

@@ -1,16 +1,24 @@
 package namespaces
 
 import (
+	"os"
+
 	"github.com/pkg/errors"
 	"golang.org/x/net/context"
 )
 
+const (
+	namespaceEnvVar  = "CONTAINERD_NAMESPACE"
+	defaultNamespace = "default"
+)
+
 var (
 	errNamespaceRequired = errors.New("namespace is required")
 )
 
 type namespaceKey struct{}
 
+// WithNamespace sets a given namespace on the context
 func WithNamespace(ctx context.Context, namespace string) context.Context {
 	ctx = context.WithValue(ctx, namespaceKey{}, namespace) // set our key for namespace
 
@@ -19,6 +27,17 @@ func WithNamespace(ctx context.Context, namespace string) context.Context {
 	return withGRPCNamespaceHeader(ctx, namespace)
 }
 
+// NamespaceFromEnv uses the namespace defined in CONTAINERD_NAMESPACE or
+// default
+func NamespaceFromEnv(ctx context.Context) context.Context {
+	namespace := os.Getenv(namespaceEnvVar)
+	if namespace == "" {
+		namespace = defaultNamespace
+	}
+	return WithNamespace(ctx, namespace)
+}
+
+// Namespace returns the namespace from the context
 func Namespace(ctx context.Context) (string, bool) {
 	namespace, ok := ctx.Value(namespaceKey{}).(string)
 	if !ok {
@@ -28,10 +47,12 @@ func Namespace(ctx context.Context) (string, bool) {
 	return namespace, ok
 }
 
+// IsNamespaceRequired returns whether an error is caused by a missing namespace
 func IsNamespaceRequired(err error) bool {
 	return errors.Cause(err) == errNamespaceRequired
 }
 
+// NamespaceRequired returns the namespace or an error
 func NamespaceRequired(ctx context.Context) (string, error) {
 	namespace, ok := Namespace(ctx)
 	if !ok || namespace == "" {

namespaces/context_test.go

@@ -2,6 +2,7 @@ package namespaces
 
 import (
 	"context"
+	"os"
 	"testing"
 )
 
@@ -28,3 +29,31 @@ func TestContext(t *testing.T) {
 		t.Fatalf("unexpected namespace: %q != %q", namespace, expected)
 	}
 }
+
+func TestNamespaceFromEnv(t *testing.T) {
+	oldenv := os.Getenv(namespaceEnvVar)
+	defer os.Setenv(namespaceEnvVar, oldenv) // restore old env var
+
+	ctx := context.Background()
+	namespace, ok := Namespace(ctx)
+	if ok {
+		t.Fatal("namespace should not be present")
+	}
+
+	if namespace != "" {
+		t.Fatalf("namespace should not be defined: got %q", namespace)
+	}
+
+	expected := "test-namespace"
+	os.Setenv(namespaceEnvVar, expected)
+	nctx := NamespaceFromEnv(ctx)
+
+	namespace, ok = Namespace(nctx)
+	if !ok {
+		t.Fatal("expected to find a namespace")
+	}
+
+	if namespace != expected {
+		t.Fatalf("unexpected namespace: %q != %q", namespace, expected)
+	}
+}