-
-
Notifications
You must be signed in to change notification settings - Fork 1.5k
/
grpcctx.go
117 lines (105 loc) · 2.92 KB
/
grpcctx.go
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
// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.
// Package grpcctx provides context feature for GRPC.
package grpcctx
import (
"context"
"google.golang.org/grpc/metadata"
"github.com/gogf/gf/v2/container/gmap"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/util/gconv"
)
type (
Ctx struct{}
)
func (c Ctx) NewIncoming(ctx context.Context, data ...g.Map) context.Context {
if len(data) > 0 {
incomingMd := make(metadata.MD)
for key, value := range data[0] {
incomingMd.Set(key, gconv.String(value))
}
return metadata.NewIncomingContext(ctx, incomingMd)
}
return metadata.NewIncomingContext(ctx, nil)
}
func (c Ctx) NewOutgoing(ctx context.Context, data ...g.Map) context.Context {
if len(data) > 0 {
outgoingMd := make(metadata.MD)
for key, value := range data[0] {
outgoingMd.Set(key, gconv.String(value))
}
return metadata.NewOutgoingContext(ctx, outgoingMd)
}
return metadata.NewOutgoingContext(ctx, nil)
}
func (c Ctx) IncomingToOutgoing(ctx context.Context, keys ...string) context.Context {
incomingMd, _ := metadata.FromIncomingContext(ctx)
if incomingMd == nil {
return ctx
}
outgoingMd, _ := metadata.FromOutgoingContext(ctx)
if outgoingMd == nil {
outgoingMd = make(metadata.MD)
}
if len(keys) > 0 {
for _, key := range keys {
outgoingMd[key] = append(outgoingMd[key], incomingMd.Get(key)...)
}
} else {
for key, values := range incomingMd {
outgoingMd[key] = append(outgoingMd[key], values...)
}
}
return metadata.NewOutgoingContext(ctx, outgoingMd)
}
func (c Ctx) IncomingMap(ctx context.Context) *gmap.Map {
var (
data = gmap.New()
incomingMd, _ = metadata.FromIncomingContext(ctx)
)
for key, values := range incomingMd {
if len(values) == 1 {
data.Set(key, values[0])
} else {
data.Set(key, values)
}
}
return data
}
func (c Ctx) OutgoingMap(ctx context.Context) *gmap.Map {
var (
data = gmap.New()
outgoingMd, _ = metadata.FromOutgoingContext(ctx)
)
for key, values := range outgoingMd {
if len(values) == 1 {
data.Set(key, values[0])
} else {
data.Set(key, values)
}
}
return data
}
func (c Ctx) SetIncoming(ctx context.Context, data g.Map) context.Context {
incomingMd, _ := metadata.FromIncomingContext(ctx)
if incomingMd == nil {
incomingMd = make(metadata.MD)
}
for key, value := range data {
incomingMd.Set(key, gconv.String(value))
}
return metadata.NewIncomingContext(ctx, incomingMd)
}
func (c Ctx) SetOutgoing(ctx context.Context, data g.Map) context.Context {
outgoingMd, _ := metadata.FromOutgoingContext(ctx)
if outgoingMd == nil {
outgoingMd = make(metadata.MD)
}
for key, value := range data {
outgoingMd.Set(key, gconv.String(value))
}
return metadata.NewOutgoingContext(ctx, outgoingMd)
}