You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

123 lines
3.6 KiB

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