/
relic.go
51 lines (41 loc) · 1.2 KB
/
relic.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
package bunrelic
import (
"context"
"github.com/newrelic/go-agent/v3/newrelic"
"github.com/newrelic/go-agent/v3/newrelic/sqlparse"
"github.com/uptrace/bun"
)
type QueryHook struct {
baseSegment newrelic.DatastoreSegment
}
var _ bun.QueryHook = (*QueryHook)(nil)
type nrBunCtxKey string
const nrBunSegmentKey nrBunCtxKey = "nrbunsegment"
// NewQueryHook creates a new bun.QueryHook which reports database usage
// information to new relic.
func NewQueryHook(options ...Option) *QueryHook {
h := &QueryHook{}
for _, o := range options {
o(h)
}
return h
}
func (q *QueryHook) BeforeQuery(ctx context.Context, qe *bun.QueryEvent) context.Context {
segment := q.baseSegment
if qe.Model != nil {
if t, ok := qe.Model.(bun.TableModel); ok {
segment.Operation = qe.Operation()
segment.Collection = t.Table().Name
} else {
sqlparse.ParseQuery(&segment, qe.Query)
}
} else {
sqlparse.ParseQuery(&segment, qe.Query)
}
segment.StartTime = newrelic.FromContext(ctx).StartSegmentNow()
return context.WithValue(ctx, nrBunSegmentKey, &segment)
}
func (q *QueryHook) AfterQuery(ctx context.Context, qe *bun.QueryEvent) {
segment := ctx.Value(nrBunSegmentKey).(*newrelic.DatastoreSegment)
segment.End()
}