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
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
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
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
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
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 }
142
View as plain text