1// Copyright 2015 Google Inc. All rights reserved. 2// Use of this source code is governed by the Apache 2.0 3// license that can be found in the LICENSE file. 4 5package internal 6 7import ( 8 "errors" 9 "os" 10 11 "github.com/golang/protobuf/proto" 12 netcontext "golang.org/x/net/context" 13) 14 15var errNotAppEngineContext = errors.New("not an App Engine context") 16 17type CallOverrideFunc func(ctx netcontext.Context, service, method string, in, out proto.Message) error 18 19var callOverrideKey = "holds []CallOverrideFunc" 20 21func WithCallOverride(ctx netcontext.Context, f CallOverrideFunc) netcontext.Context { 22 // We avoid appending to any existing call override 23 // so we don't risk overwriting a popped stack below. 24 var cofs []CallOverrideFunc 25 if uf, ok := ctx.Value(&callOverrideKey).([]CallOverrideFunc); ok { 26 cofs = append(cofs, uf...) 27 } 28 cofs = append(cofs, f) 29 return netcontext.WithValue(ctx, &callOverrideKey, cofs) 30} 31 32func callOverrideFromContext(ctx netcontext.Context) (CallOverrideFunc, netcontext.Context, bool) { 33 cofs, _ := ctx.Value(&callOverrideKey).([]CallOverrideFunc) 34 if len(cofs) == 0 { 35 return nil, nil, false 36 } 37 // We found a list of overrides; grab the last, and reconstitute a 38 // context that will hide it. 39 f := cofs[len(cofs)-1] 40 ctx = netcontext.WithValue(ctx, &callOverrideKey, cofs[:len(cofs)-1]) 41 return f, ctx, true 42} 43 44type logOverrideFunc func(level int64, format string, args ...interface{}) 45 46var logOverrideKey = "holds a logOverrideFunc" 47 48func WithLogOverride(ctx netcontext.Context, f logOverrideFunc) netcontext.Context { 49 return netcontext.WithValue(ctx, &logOverrideKey, f) 50} 51 52var appIDOverrideKey = "holds a string, being the full app ID" 53 54func WithAppIDOverride(ctx netcontext.Context, appID string) netcontext.Context { 55 return netcontext.WithValue(ctx, &appIDOverrideKey, appID) 56} 57 58var namespaceKey = "holds the namespace string" 59 60func withNamespace(ctx netcontext.Context, ns string) netcontext.Context { 61 return netcontext.WithValue(ctx, &namespaceKey, ns) 62} 63 64func NamespaceFromContext(ctx netcontext.Context) string { 65 // If there's no namespace, return the empty string. 66 ns, _ := ctx.Value(&namespaceKey).(string) 67 return ns 68} 69 70// FullyQualifiedAppID returns the fully-qualified application ID. 71// This may contain a partition prefix (e.g. "s~" for High Replication apps), 72// or a domain prefix (e.g. "example.com:"). 73func FullyQualifiedAppID(ctx netcontext.Context) string { 74 if id, ok := ctx.Value(&appIDOverrideKey).(string); ok { 75 return id 76 } 77 return fullyQualifiedAppID(ctx) 78} 79 80func Logf(ctx netcontext.Context, level int64, format string, args ...interface{}) { 81 if f, ok := ctx.Value(&logOverrideKey).(logOverrideFunc); ok { 82 f(level, format, args...) 83 return 84 } 85 c := fromContext(ctx) 86 if c == nil { 87 panic(errNotAppEngineContext) 88 } 89 logf(c, level, format, args...) 90} 91 92// NamespacedContext wraps a Context to support namespaces. 93func NamespacedContext(ctx netcontext.Context, namespace string) netcontext.Context { 94 return withNamespace(ctx, namespace) 95} 96 97// SetTestEnv sets the env variables for testing background ticket in Flex. 98func SetTestEnv() func() { 99 var environ = []struct { 100 key, value string 101 }{ 102 {"GAE_LONG_APP_ID", "my-app-id"}, 103 {"GAE_MINOR_VERSION", "067924799508853122"}, 104 {"GAE_MODULE_INSTANCE", "0"}, 105 {"GAE_MODULE_NAME", "default"}, 106 {"GAE_MODULE_VERSION", "20150612t184001"}, 107 } 108 109 for _, v := range environ { 110 old := os.Getenv(v.key) 111 os.Setenv(v.key, v.value) 112 v.value = old 113 } 114 return func() { // Restore old environment after the test completes. 115 for _, v := range environ { 116 if v.value == "" { 117 os.Unsetenv(v.key) 118 continue 119 } 120 os.Setenv(v.key, v.value) 121 } 122 } 123} 124