/
template.go
144 lines (129 loc) · 3.51 KB
/
template.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
132
133
134
135
136
137
138
139
140
141
142
143
144
package stages
import (
"bytes"
"crypto/sha256"
"encoding/hex"
"errors"
"reflect"
"regexp"
"strings"
"text/template"
"time"
"github.com/go-kit/kit/log"
"github.com/go-kit/kit/log/level"
"github.com/mitchellh/mapstructure"
"github.com/prometheus/common/model"
"golang.org/x/crypto/sha3"
)
// Config Errors
const (
ErrEmptyTemplateStageConfig = "template stage config cannot be empty"
ErrTemplateSourceRequired = "template source value is required"
)
var (
functionMap = template.FuncMap{
"ToLower": strings.ToLower,
"ToUpper": strings.ToUpper,
"Replace": strings.Replace,
"Trim": strings.Trim,
"TrimLeft": strings.TrimLeft,
"TrimRight": strings.TrimRight,
"TrimPrefix": strings.TrimPrefix,
"TrimSuffix": strings.TrimSuffix,
"TrimSpace": strings.TrimSpace,
"Hash": func(salt string, input string) string {
hash := sha3.Sum256([]byte(salt + input))
return hex.EncodeToString(hash[:])
},
"Sha2Hash": func(salt string, input string) string {
hash := sha256.Sum256([]byte(salt + input))
return hex.EncodeToString(hash[:])
},
"regexReplaceAll": func(regex string, s string, repl string) string {
r := regexp.MustCompile(regex)
return r.ReplaceAllString(s, repl)
},
"regexReplaceAllLiteral": func(regex string, s string, repl string) string {
r := regexp.MustCompile(regex)
return r.ReplaceAllLiteralString(s, repl)
},
}
)
// TemplateConfig configures template value extraction
type TemplateConfig struct {
Source string `mapstructure:"source"`
Template string `mapstructure:"template"`
}
// validateTemplateConfig validates the templateStage config
func validateTemplateConfig(cfg *TemplateConfig) (*template.Template, error) {
if cfg == nil {
return nil, errors.New(ErrEmptyTemplateStageConfig)
}
if cfg.Source == "" {
return nil, errors.New(ErrTemplateSourceRequired)
}
return template.New("pipeline_template").Funcs(functionMap).Parse(cfg.Template)
}
// newTemplateStage creates a new templateStage
func newTemplateStage(logger log.Logger, config interface{}) (Stage, error) {
cfg := &TemplateConfig{}
err := mapstructure.Decode(config, cfg)
if err != nil {
return nil, err
}
t, err := validateTemplateConfig(cfg)
if err != nil {
return nil, err
}
return toStage(&templateStage{
cfgs: cfg,
logger: logger,
template: t,
}), nil
}
// templateStage will mutate the incoming entry and set it from extracted data
type templateStage struct {
cfgs *TemplateConfig
logger log.Logger
template *template.Template
}
// Process implements Stage
func (o *templateStage) Process(labels model.LabelSet, extracted map[string]interface{}, t *time.Time, entry *string) {
if o.cfgs == nil {
return
}
td := make(map[string]interface{})
for k, v := range extracted {
s, err := getString(v)
if err != nil {
if Debug {
level.Debug(o.logger).Log("msg", "extracted template could not be converted to a string", "err", err, "type", reflect.TypeOf(v))
}
continue
}
td[k] = s
if k == o.cfgs.Source {
td["Value"] = s
}
}
td["Entry"] = *entry
buf := &bytes.Buffer{}
err := o.template.Execute(buf, td)
if err != nil {
if Debug {
level.Debug(o.logger).Log("msg", "failed to execute template on extracted value", "err", err)
}
return
}
st := buf.String()
// If the template evaluates to an empty string, remove the key from the map
if st == "" {
delete(extracted, o.cfgs.Source)
} else {
extracted[o.cfgs.Source] = st
}
}
// Name implements Stage
func (o *templateStage) Name() string {
return StageTypeTemplate
}