ctx.go 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. // Copyright 2020 The etcd Authors
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package clientv3
  15. import (
  16. "context"
  17. "strings"
  18. "go.etcd.io/etcd/etcdserver/api/v3rpc/rpctypes"
  19. "go.etcd.io/etcd/version"
  20. "google.golang.org/grpc/metadata"
  21. )
  22. // WithRequireLeader requires client requests to only succeed
  23. // when the cluster has a leader.
  24. func WithRequireLeader(ctx context.Context) context.Context {
  25. md, ok := metadata.FromOutgoingContext(ctx)
  26. if !ok { // no outgoing metadata ctx key, create one
  27. md = metadata.Pairs(rpctypes.MetadataRequireLeaderKey, rpctypes.MetadataHasLeader)
  28. return metadata.NewOutgoingContext(ctx, md)
  29. }
  30. copied := md.Copy() // avoid racey updates
  31. // overwrite/add 'hasleader' key/value
  32. metadataSet(copied, rpctypes.MetadataRequireLeaderKey, rpctypes.MetadataHasLeader)
  33. return metadata.NewOutgoingContext(ctx, copied)
  34. }
  35. // embeds client version
  36. func withVersion(ctx context.Context) context.Context {
  37. md, ok := metadata.FromOutgoingContext(ctx)
  38. if !ok { // no outgoing metadata ctx key, create one
  39. md = metadata.Pairs(rpctypes.MetadataClientAPIVersionKey, version.APIVersion)
  40. return metadata.NewOutgoingContext(ctx, md)
  41. }
  42. copied := md.Copy() // avoid racey updates
  43. // overwrite/add version key/value
  44. metadataSet(copied, rpctypes.MetadataClientAPIVersionKey, version.APIVersion)
  45. return metadata.NewOutgoingContext(ctx, copied)
  46. }
  47. func metadataGet(md metadata.MD, k string) []string {
  48. k = strings.ToLower(k)
  49. return md[k]
  50. }
  51. func metadataSet(md metadata.MD, k string, vals ...string) {
  52. if len(vals) == 0 {
  53. return
  54. }
  55. k = strings.ToLower(k)
  56. md[k] = vals
  57. }