-
Notifications
You must be signed in to change notification settings - Fork 568
/
field.go
131 lines (121 loc) · 3.73 KB
/
field.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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package log
import (
"context"
"encoding/json"
"sort"
"strings"
"github.com/pachyderm/pachyderm/v2/src/constants"
"github.com/pachyderm/pachyderm/v2/src/internal/errors"
"github.com/pachyderm/pachyderm/v2/src/protoextensions"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
"golang.org/x/exp/maps"
"google.golang.org/grpc/metadata"
"google.golang.org/protobuf/encoding/protojson"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/anypb"
"google.golang.org/protobuf/types/known/durationpb"
"google.golang.org/protobuf/types/known/emptypb"
"google.golang.org/protobuf/types/known/timestamppb"
"google.golang.org/protobuf/types/known/wrapperspb"
)
// Proto is a Field containing a protocol buffer message.
func Proto(name string, msg proto.Message) Field {
switch x := msg.(type) {
case zapcore.ObjectMarshaler:
return zap.Object(name, x)
case *emptypb.Empty:
return zap.Skip()
case *timestamppb.Timestamp:
return zap.Time(name, x.AsTime())
case *durationpb.Duration:
return zap.Duration(name, x.AsDuration())
case *wrapperspb.BytesValue:
return zap.Object(name, protoextensions.ConciseBytes(x.GetValue()))
case *wrapperspb.Int64Value:
return zap.Int64(name, x.GetValue())
case *anypb.Any:
msg, err := x.UnmarshalNew()
if err != nil {
return zap.Any(name, x)
}
return Proto(name, msg)
}
if js, err := protoToJSONMap(msg); err == nil {
return zap.Any(name, js)
}
return zap.Any(name, msg)
}
func protoToJSONMap(msg proto.Message) (map[string]any, error) {
// This is a detail of dynamicpb.Message; if you pass in a raw &dynamicpb.Message{}, then
// marshalling panics. This avoids that panic.
if x, ok := msg.(interface{ IsValid() bool }); ok {
if !x.IsValid() {
return nil, errors.New("invalid dynamic message")
}
}
js, err := protojson.Marshal(msg)
if err != nil {
return nil, errors.Wrap(err, "marshal arbitrary message")
}
result := make(map[string]any)
if err := json.Unmarshal(js, &result); err != nil {
return nil, errors.Wrap(err, "unmarshal into map[string]any")
}
return result, nil
}
type attempt struct{ i, max int }
// MarshalLogObject implements zapcore.ObjectMarshaler.
func (a attempt) MarshalLogObject(enc zapcore.ObjectEncoder) error {
enc.AddInt("attempt", a.i)
enc.AddInt("totalAttempts", a.max)
return nil
}
// RetryAttempt is a Field that encodes the current retry (0-indexed) and the total number of
// retries. It's intended for a for loop where "i" is the loop iterator and "max" is the upper
// bound "i < max".
func RetryAttempt(i int, max int) Field {
return zap.Inline(&attempt{i: i, max: max})
}
// Metadata is a Field that logs the provided metadata (canonicalizing keys, collapsing
// single-element values to strings, and removing the Pachyderm auth token).
func Metadata(name string, md metadata.MD) Field {
return zap.Object(name, zapcore.ObjectMarshalerFunc(func(enc zapcore.ObjectEncoder) error {
keys := maps.Keys(md)
sort.Strings(keys)
for _, k := range keys {
v := md[k]
if k == constants.ContextTokenKey {
for i := range v {
v[i] = "[MASKED]"
}
}
switch len(v) {
case 0:
continue
case 1:
enc.AddString(strings.ToLower(k), v[0])
default:
if err := enc.AddArray(strings.ToLower(k), zapcore.ArrayMarshalerFunc(
func(enc zapcore.ArrayEncoder) error {
for _, x := range v {
enc.AppendString(x)
}
return nil
},
)); err != nil {
return errors.Wrap(err, "add metadata value array")
}
}
}
return nil
}))
}
// OutgoingMetadata is a Field that logs the outgoing metadata associated with the provided context.
func OutgoingMetadata(ctx context.Context) Field {
md, ok := metadata.FromOutgoingContext(ctx)
if !ok {
return zap.Skip()
}
return Metadata("metadata", md)
}