api_common.go 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  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. "github.com/golang/protobuf/proto"
  7. netcontext "golang.org/x/net/context"
  8. )
  9. type CallOverrideFunc func(ctx netcontext.Context, service, method string, in, out proto.Message) error
  10. var callOverrideKey = "holds []CallOverrideFunc"
  11. func WithCallOverride(ctx netcontext.Context, f CallOverrideFunc) netcontext.Context {
  12. // We avoid appending to any existing call override
  13. // so we don't risk overwriting a popped stack below.
  14. var cofs []CallOverrideFunc
  15. if uf, ok := ctx.Value(&callOverrideKey).([]CallOverrideFunc); ok {
  16. cofs = append(cofs, uf...)
  17. }
  18. cofs = append(cofs, f)
  19. return netcontext.WithValue(ctx, &callOverrideKey, cofs)
  20. }
  21. func callOverrideFromContext(ctx netcontext.Context) (CallOverrideFunc, netcontext.Context, bool) {
  22. cofs, _ := ctx.Value(&callOverrideKey).([]CallOverrideFunc)
  23. if len(cofs) == 0 {
  24. return nil, nil, false
  25. }
  26. // We found a list of overrides; grab the last, and reconstitute a
  27. // context that will hide it.
  28. f := cofs[len(cofs)-1]
  29. ctx = netcontext.WithValue(ctx, &callOverrideKey, cofs[:len(cofs)-1])
  30. return f, ctx, true
  31. }
  32. type logOverrideFunc func(level int64, format string, args ...interface{})
  33. var logOverrideKey = "holds a logOverrideFunc"
  34. func WithLogOverride(ctx netcontext.Context, f logOverrideFunc) netcontext.Context {
  35. return netcontext.WithValue(ctx, &logOverrideKey, f)
  36. }
  37. var appIDOverrideKey = "holds a string, being the full app ID"
  38. func WithAppIDOverride(ctx netcontext.Context, appID string) netcontext.Context {
  39. return netcontext.WithValue(ctx, &appIDOverrideKey, appID)
  40. }
  41. var namespaceKey = "holds the namespace string"
  42. func withNamespace(ctx netcontext.Context, ns string) netcontext.Context {
  43. return netcontext.WithValue(ctx, &namespaceKey, ns)
  44. }
  45. func NamespaceFromContext(ctx netcontext.Context) string {
  46. // If there's no namespace, return the empty string.
  47. ns, _ := ctx.Value(&namespaceKey).(string)
  48. return ns
  49. }
  50. // FullyQualifiedAppID returns the fully-qualified application ID.
  51. // This may contain a partition prefix (e.g. "s~" for High Replication apps),
  52. // or a domain prefix (e.g. "example.com:").
  53. func FullyQualifiedAppID(ctx netcontext.Context) string {
  54. if id, ok := ctx.Value(&appIDOverrideKey).(string); ok {
  55. return id
  56. }
  57. return fullyQualifiedAppID(ctx)
  58. }
  59. func Logf(ctx netcontext.Context, level int64, format string, args ...interface{}) {
  60. if f, ok := ctx.Value(&logOverrideKey).(logOverrideFunc); ok {
  61. f(level, format, args...)
  62. return
  63. }
  64. logf(fromContext(ctx), level, format, args...)
  65. }
  66. // NamespacedContext wraps a Context to support namespaces.
  67. func NamespacedContext(ctx netcontext.Context, namespace string) netcontext.Context {
  68. return withNamespace(ctx, namespace)
  69. }