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.
11 "github.com/golang/protobuf/proto"
12 netcontext "golang.org/x/net/context"
15 var errNotAppEngineContext = errors.New("not an App Engine context")
17 type CallOverrideFunc func(ctx netcontext.Context, service, method string, in, out proto.Message) error
19 var callOverrideKey = "holds []CallOverrideFunc"
21 func 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...)
28 cofs = append(cofs, f)
29 return netcontext.WithValue(ctx, &callOverrideKey, cofs)
32 func callOverrideFromContext(ctx netcontext.Context) (CallOverrideFunc, netcontext.Context, bool) {
33 cofs, _ := ctx.Value(&callOverrideKey).([]CallOverrideFunc)
35 return nil, nil, false
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])
44 type logOverrideFunc func(level int64, format string, args ...interface{})
46 var logOverrideKey = "holds a logOverrideFunc"
48 func WithLogOverride(ctx netcontext.Context, f logOverrideFunc) netcontext.Context {
49 return netcontext.WithValue(ctx, &logOverrideKey, f)
52 var appIDOverrideKey = "holds a string, being the full app ID"
54 func WithAppIDOverride(ctx netcontext.Context, appID string) netcontext.Context {
55 return netcontext.WithValue(ctx, &appIDOverrideKey, appID)
58 var namespaceKey = "holds the namespace string"
60 func withNamespace(ctx netcontext.Context, ns string) netcontext.Context {
61 return netcontext.WithValue(ctx, &namespaceKey, ns)
64 func NamespaceFromContext(ctx netcontext.Context) string {
65 // If there's no namespace, return the empty string.
66 ns, _ := ctx.Value(&namespaceKey).(string)
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:").
73 func FullyQualifiedAppID(ctx netcontext.Context) string {
74 if id, ok := ctx.Value(&appIDOverrideKey).(string); ok {
77 return fullyQualifiedAppID(ctx)
80 func Logf(ctx netcontext.Context, level int64, format string, args ...interface{}) {
81 if f, ok := ctx.Value(&logOverrideKey).(logOverrideFunc); ok {
82 f(level, format, args...)
87 panic(errNotAppEngineContext)
89 logf(c, level, format, args...)
92 // NamespacedContext wraps a Context to support namespaces.
93 func NamespacedContext(ctx netcontext.Context, namespace string) netcontext.Context {
94 return withNamespace(ctx, namespace)
97 // SetTestEnv sets the env variables for testing background ticket in Flex.
98 func SetTestEnv() func() {
99 var environ = []struct {
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"},
109 for _, v := range environ {
110 old := os.Getenv(v.key)
111 os.Setenv(v.key, v.value)
114 return func() { // Restore old environment after the test completes.
115 for _, v := range environ {
120 os.Setenv(v.key, v.value)