Codebase list golang-github-go-kit-kit / 4079d84 transport / grpc / _grpc_test / context_metadata.go
4079d84

Tree @4079d84 (Download .tar.gz)

context_metadata.go @4079d84 — raw · history · blame 

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
package test

import (
	"context"
	"fmt"

	"google.golang.org/grpc/metadata"
)

type metaContext string

const (
	correlationID     metaContext = "correlation-id"
	responseHDR       metaContext = "my-response-header"
	responseTRLR      metaContext = "my-response-trailer"
	correlationIDTRLR metaContext = "correlation-id-consumed"
)

/* client before functions */

func injectCorrelationID(ctx context.Context, md *metadata.MD) context.Context {
	if hdr, ok := ctx.Value(correlationID).(string); ok {
		fmt.Printf("\tClient found correlationID %q in context, set metadata header\n", hdr)
		(*md)[string(correlationID)] = append((*md)[string(correlationID)], hdr)
	}
	return ctx
}

func displayClientRequestHeaders(ctx context.Context, md *metadata.MD) context.Context {
	if len(*md) > 0 {
		fmt.Println("\tClient >> Request Headers:")
		for key, val := range *md {
			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
		}
	}
	return ctx
}

/* server before functions */

func extractCorrelationID(ctx context.Context, md metadata.MD) context.Context {
	if hdr, ok := md[string(correlationID)]; ok {
		cID := hdr[len(hdr)-1]
		ctx = context.WithValue(ctx, correlationID, cID)
		fmt.Printf("\tServer received correlationID %q in metadata header, set context\n", cID)
	}
	return ctx
}

func displayServerRequestHeaders(ctx context.Context, md metadata.MD) context.Context {
	if len(md) > 0 {
		fmt.Println("\tServer << Request Headers:")
		for key, val := range md {
			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
		}
	}
	return ctx
}

/* server after functions */

func injectResponseHeader(ctx context.Context, md *metadata.MD, _ *metadata.MD) context.Context {
	*md = metadata.Join(*md, metadata.Pairs(string(responseHDR), "has-a-value"))
	return ctx
}

func displayServerResponseHeaders(ctx context.Context, md *metadata.MD, _ *metadata.MD) context.Context {
	if len(*md) > 0 {
		fmt.Println("\tServer >> Response Headers:")
		for key, val := range *md {
			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
		}
	}
	return ctx
}

func injectResponseTrailer(ctx context.Context, _ *metadata.MD, md *metadata.MD) context.Context {
	*md = metadata.Join(*md, metadata.Pairs(string(responseTRLR), "has-a-value-too"))
	return ctx
}

func injectConsumedCorrelationID(ctx context.Context, _ *metadata.MD, md *metadata.MD) context.Context {
	if hdr, ok := ctx.Value(correlationID).(string); ok {
		fmt.Printf("\tServer found correlationID %q in context, set consumed trailer\n", hdr)
		*md = metadata.Join(*md, metadata.Pairs(string(correlationIDTRLR), hdr))
	}
	return ctx
}

func displayServerResponseTrailers(ctx context.Context, _ *metadata.MD, md *metadata.MD) context.Context {
	if len(*md) > 0 {
		fmt.Println("\tServer >> Response Trailers:")
		for key, val := range *md {
			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
		}
	}
	return ctx
}

/* client after functions */

func displayClientResponseHeaders(ctx context.Context, md metadata.MD, _ metadata.MD) context.Context {
	if len(md) > 0 {
		fmt.Println("\tClient << Response Headers:")
		for key, val := range md {
			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
		}
	}
	return ctx
}

func displayClientResponseTrailers(ctx context.Context, _ metadata.MD, md metadata.MD) context.Context {
	if len(md) > 0 {
		fmt.Println("\tClient << Response Trailers:")
		for key, val := range md {
			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
		}
	}
	return ctx
}

func extractConsumedCorrelationID(ctx context.Context, _ metadata.MD, md metadata.MD) context.Context {
	if hdr, ok := md[string(correlationIDTRLR)]; ok {
		fmt.Printf("\tClient received consumed correlationID %q in metadata trailer, set context\n", hdr[len(hdr)-1])
		ctx = context.WithValue(ctx, correlationIDTRLR, hdr[len(hdr)-1])
	}
	return ctx
}

/* CorrelationID context handlers */

func SetCorrelationID(ctx context.Context, v string) context.Context {
	return context.WithValue(ctx, correlationID, v)
}

func GetConsumedCorrelationID(ctx context.Context) string {
	if trlr, ok := ctx.Value(correlationIDTRLR).(string); ok {
		return trlr
	}
	return ""
}