|
|
- // Copyright 2015 Google Inc. All rights reserved.
- // Use of this source code is governed by the Apache 2.0
- // license that can be found in the LICENSE file.
-
- package internal
-
- import (
- "errors"
- "os"
-
- "github.com/golang/protobuf/proto"
- netcontext "golang.org/x/net/context"
- )
-
- var errNotAppEngineContext = errors.New("not an App Engine context")
-
- type CallOverrideFunc func(ctx netcontext.Context, service, method string, in, out proto.Message) error
-
- var callOverrideKey = "holds []CallOverrideFunc"
-
- func WithCallOverride(ctx netcontext.Context, f CallOverrideFunc) netcontext.Context {
- // We avoid appending to any existing call override
- // so we don't risk overwriting a popped stack below.
- var cofs []CallOverrideFunc
- if uf, ok := ctx.Value(&callOverrideKey).([]CallOverrideFunc); ok {
- cofs = append(cofs, uf...)
- }
- cofs = append(cofs, f)
- return netcontext.WithValue(ctx, &callOverrideKey, cofs)
- }
-
- func callOverrideFromContext(ctx netcontext.Context) (CallOverrideFunc, netcontext.Context, bool) {
- cofs, _ := ctx.Value(&callOverrideKey).([]CallOverrideFunc)
- if len(cofs) == 0 {
- return nil, nil, false
- }
- // We found a list of overrides; grab the last, and reconstitute a
- // context that will hide it.
- f := cofs[len(cofs)-1]
- ctx = netcontext.WithValue(ctx, &callOverrideKey, cofs[:len(cofs)-1])
- return f, ctx, true
- }
-
- type logOverrideFunc func(level int64, format string, args ...interface{})
-
- var logOverrideKey = "holds a logOverrideFunc"
-
- func WithLogOverride(ctx netcontext.Context, f logOverrideFunc) netcontext.Context {
- return netcontext.WithValue(ctx, &logOverrideKey, f)
- }
-
- var appIDOverrideKey = "holds a string, being the full app ID"
-
- func WithAppIDOverride(ctx netcontext.Context, appID string) netcontext.Context {
- return netcontext.WithValue(ctx, &appIDOverrideKey, appID)
- }
-
- var namespaceKey = "holds the namespace string"
-
- func withNamespace(ctx netcontext.Context, ns string) netcontext.Context {
- return netcontext.WithValue(ctx, &namespaceKey, ns)
- }
-
- func NamespaceFromContext(ctx netcontext.Context) string {
- // If there's no namespace, return the empty string.
- ns, _ := ctx.Value(&namespaceKey).(string)
- return ns
- }
-
- // FullyQualifiedAppID returns the fully-qualified application ID.
- // This may contain a partition prefix (e.g. "s~" for High Replication apps),
- // or a domain prefix (e.g. "example.com:").
- func FullyQualifiedAppID(ctx netcontext.Context) string {
- if id, ok := ctx.Value(&appIDOverrideKey).(string); ok {
- return id
- }
- return fullyQualifiedAppID(ctx)
- }
-
- func Logf(ctx netcontext.Context, level int64, format string, args ...interface{}) {
- if f, ok := ctx.Value(&logOverrideKey).(logOverrideFunc); ok {
- f(level, format, args...)
- return
- }
- c := fromContext(ctx)
- if c == nil {
- panic(errNotAppEngineContext)
- }
- logf(c, level, format, args...)
- }
-
- // NamespacedContext wraps a Context to support namespaces.
- func NamespacedContext(ctx netcontext.Context, namespace string) netcontext.Context {
- return withNamespace(ctx, namespace)
- }
-
- // SetTestEnv sets the env variables for testing background ticket in Flex.
- func SetTestEnv() func() {
- var environ = []struct {
- key, value string
- }{
- {"GAE_LONG_APP_ID", "my-app-id"},
- {"GAE_MINOR_VERSION", "067924799508853122"},
- {"GAE_MODULE_INSTANCE", "0"},
- {"GAE_MODULE_NAME", "default"},
- {"GAE_MODULE_VERSION", "20150612t184001"},
- }
-
- for _, v := range environ {
- old := os.Getenv(v.key)
- os.Setenv(v.key, v.value)
- v.value = old
- }
- return func() { // Restore old environment after the test completes.
- for _, v := range environ {
- if v.value == "" {
- os.Unsetenv(v.key)
- continue
- }
- os.Setenv(v.key, v.value)
- }
- }
- }
|