/
hook_es.go
71 lines (60 loc) · 1.67 KB
/
hook_es.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
package trace
import (
"bytes"
"io"
"net/http"
"strconv"
"github.com/opentracing/opentracing-go"
"github.com/opentracing/opentracing-go/ext"
"github.com/opentracing/opentracing-go/log"
)
const MaxContentLength = 1 << 16
var elasticComponent = opentracing.Tag{string(ext.Component), "elastic"}
type ESTracedTransport struct {
*http.Transport
}
func (t *ESTracedTransport) RoundTrip(r *http.Request) (resp *http.Response, err error) {
span, ctx := StartSpanFromContext(
r.Context(),
"elastic",
CustomTag("peer.service", "elastic"),
TagSpanKind("client"),
HeaderExtractor(r.Header),
CustomTag("http.url", r.URL.Path),
CustomTag("http.method", r.Method),
)
r = r.WithContext(ctx)
defer func() {
if err != nil {
span.SetTag("elastic.error", err.Error())
span.SetTag(string(ext.Error), true)
}
span.Finish()
}()
span.SetTag(string(ext.DBType), "elastic")
span.SetTag(string(ext.DBInstance), r.URL.Host)
span.SetTag("elastic.method", r.Method)
span.SetTag("elastic.url", r.URL.Path)
span.SetTag("elastic.params", r.URL.Query().Encode())
contentLength, _ := strconv.Atoi(r.Header.Get("Content-Length"))
if r.Body != nil && contentLength < MaxContentLength {
buf, err := io.ReadAll(r.Body)
if err != nil {
return nil, err
}
span.SetTag(string(ext.DBStatement), string(buf))
span.LogFields(log.String("params", string(buf)))
r.Body = io.NopCloser(bytes.NewBuffer(buf))
}
resp, err = t.Transport.RoundTrip(r)
if err != nil {
return nil, err
}
span.SetTag("elastic.status_code", resp.StatusCode)
return resp, err
}
func ESTraceServerInterceptor() *http.Client {
return &http.Client{
Transport: &ESTracedTransport{&http.Transport{}},
}
}