vendor/github.com/go-kit/kit/transport/grpc/_grpc_test/context_metadata.go

   1 package test
   2
   3 import (
   4         "context"
   5         "fmt"
   6
   7         "google.golang.org/grpc/metadata"
   8 )
   9
  10 type metaContext string
  11
  12 const (
  13         correlationID     metaContext = "correlation-id"
  14         responseHDR       metaContext = "my-response-header"
  15         responseTRLR      metaContext = "my-response-trailer"
  16         correlationIDTRLR metaContext = "correlation-id-consumed"
  17 )
  18
  19 /* client before functions */
  20
  21 func injectCorrelationID(ctx context.Context, md *metadata.MD) context.Context {
  22         if hdr, ok := ctx.Value(correlationID).(string); ok {
  23                 fmt.Printf("\tClient found correlationID %q in context, set metadata header\n", hdr)
  24                 (*md)[string(correlationID)] = append((*md)[string(correlationID)], hdr)
  25         }
  26         return ctx
  27 }
  28
  29 func displayClientRequestHeaders(ctx context.Context, md *metadata.MD) context.Context {
  30         if len(*md) > 0 {
  31                 fmt.Println("\tClient >> Request Headers:")
  32                 for key, val := range *md {
  33                         fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
  34                 }
  35         }
  36         return ctx
  37 }
  38
  39 /* server before functions */
  40
  41 func extractCorrelationID(ctx context.Context, md metadata.MD) context.Context {
  42         if hdr, ok := md[string(correlationID)]; ok {
  43                 cID := hdr[len(hdr)-1]
  44                 ctx = context.WithValue(ctx, correlationID, cID)
  45                 fmt.Printf("\tServer received correlationID %q in metadata header, set context\n", cID)
  46         }
  47         return ctx
  48 }
  49
  50 func displayServerRequestHeaders(ctx context.Context, md metadata.MD) context.Context {
  51         if len(md) > 0 {
  52                 fmt.Println("\tServer << Request Headers:")
  53                 for key, val := range md {
  54                         fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
  55                 }
  56         }
  57         return ctx
  58 }
  59
  60 /* server after functions */
  61
  62 func injectResponseHeader(ctx context.Context, md *metadata.MD, _ *metadata.MD) context.Context {
  63         *md = metadata.Join(*md, metadata.Pairs(string(responseHDR), "has-a-value"))
  64         return ctx
  65 }
  66
  67 func displayServerResponseHeaders(ctx context.Context, md *metadata.MD, _ *metadata.MD) context.Context {
  68         if len(*md) > 0 {
  69                 fmt.Println("\tServer >> Response Headers:")
  70                 for key, val := range *md {
  71                         fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
  72                 }
  73         }
  74         return ctx
  75 }
  76
  77 func injectResponseTrailer(ctx context.Context, _ *metadata.MD, md *metadata.MD) context.Context {
  78         *md = metadata.Join(*md, metadata.Pairs(string(responseTRLR), "has-a-value-too"))
  79         return ctx
  80 }
  81
  82 func injectConsumedCorrelationID(ctx context.Context, _ *metadata.MD, md *metadata.MD) context.Context {
  83         if hdr, ok := ctx.Value(correlationID).(string); ok {
  84                 fmt.Printf("\tServer found correlationID %q in context, set consumed trailer\n", hdr)
  85                 *md = metadata.Join(*md, metadata.Pairs(string(correlationIDTRLR), hdr))
  86         }
  87         return ctx
  88 }
  89
  90 func displayServerResponseTrailers(ctx context.Context, _ *metadata.MD, md *metadata.MD) context.Context {
  91         if len(*md) > 0 {
  92                 fmt.Println("\tServer >> Response Trailers:")
  93                 for key, val := range *md {
  94                         fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
  95                 }
  96         }
  97         return ctx
  98 }
  99
 100 /* client after functions */
 101
 102 func displayClientResponseHeaders(ctx context.Context, md metadata.MD, _ metadata.MD) context.Context {
 103         if len(md) > 0 {
 104                 fmt.Println("\tClient << Response Headers:")
 105                 for key, val := range md {
 106                         fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
 107                 }
 108         }
 109         return ctx
 110 }
 111
 112 func displayClientResponseTrailers(ctx context.Context, _ metadata.MD, md metadata.MD) context.Context {
 113         if len(md) > 0 {
 114                 fmt.Println("\tClient << Response Trailers:")
 115                 for key, val := range md {
 116                         fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
 117                 }
 118         }
 119         return ctx
 120 }
 121
 122 func extractConsumedCorrelationID(ctx context.Context, _ metadata.MD, md metadata.MD) context.Context {
 123         if hdr, ok := md[string(correlationIDTRLR)]; ok {
 124                 fmt.Printf("\tClient received consumed correlationID %q in metadata trailer, set context\n", hdr[len(hdr)-1])
 125                 ctx = context.WithValue(ctx, correlationIDTRLR, hdr[len(hdr)-1])
 126         }
 127         return ctx
 128 }
 129
 130 /* CorrelationID context handlers */
 131
 132 func SetCorrelationID(ctx context.Context, v string) context.Context {
 133         return context.WithValue(ctx, correlationID, v)
 134 }
 135
 136 func GetConsumedCorrelationID(ctx context.Context) string {
 137         if trlr, ok := ctx.Value(correlationIDTRLR).(string); ok {
 138                 return trlr
 139         }
 140         return ""
 141 }