-
Notifications
You must be signed in to change notification settings - Fork 429
/
gorm.go
156 lines (137 loc) · 4.89 KB
/
gorm.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
145
146
147
148
149
150
151
152
153
154
155
156
// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016 Datadog, Inc.
// Package gorm provides helper functions for tracing the jinzhu/gorm package (https://github.com/jinzhu/gorm).
//
// Deprecated: The underlying github.com/jinzhu/gorm packages has known security vulnerabilities and is no longer under
// active development.
// It is highly recommended that you update to the latest version available here as a contrib package "gorm.io/gorm.v1".
package gorm
import (
"context"
"math"
"time"
sqltraced "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
"gopkg.in/DataDog/dd-trace-go.v1/internal/telemetry"
"github.com/jinzhu/gorm"
)
const componentName = "jinzhu/gorm"
func init() {
telemetry.LoadIntegration(componentName)
tracer.MarkIntegrationImported("github.com/jinzhu/gorm")
}
const (
gormContextKey = "dd-trace-go:context"
gormConfigKey = "dd-trace-go:config"
gormSpanStartTimeKey = "dd-trace-go:span"
)
// Open opens a new (traced) database connection. The used dialect must be formerly registered
// using (gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql).Register.
func Open(dialect, source string, opts ...Option) (*gorm.DB, error) {
sqldb, err := sqltraced.Open(dialect, source)
if err != nil {
return nil, err
}
db, err := gorm.Open(dialect, sqldb)
if err != nil {
return db, err
}
return WithCallbacks(db, opts...), err
}
// WithCallbacks registers callbacks to the gorm.DB for tracing.
// It should be called once, after opening the db.
// The callbacks are triggered by Create, Update, Delete,
// Query and RowQuery operations.
func WithCallbacks(db *gorm.DB, opts ...Option) *gorm.DB {
afterFunc := func(operationName string) func(*gorm.Scope) {
return func(scope *gorm.Scope) {
after(scope, operationName)
}
}
cb := db.Callback()
cb.Create().Before("gorm:before_create").Register("dd-trace-go:before_create", before)
cb.Create().After("gorm:after_create").Register("dd-trace-go:after_create", afterFunc("gorm.create"))
cb.Update().Before("gorm:before_update").Register("dd-trace-go:before_update", before)
cb.Update().After("gorm:after_update").Register("dd-trace-go:after_update", afterFunc("gorm.update"))
cb.Delete().Before("gorm:before_delete").Register("dd-trace-go:before_delete", before)
cb.Delete().After("gorm:after_delete").Register("dd-trace-go:after_delete", afterFunc("gorm.delete"))
cb.Query().Before("gorm:query").Register("dd-trace-go:before_query", before)
cb.Query().After("gorm:after_query").Register("dd-trace-go:after_query", afterFunc("gorm.query"))
cb.RowQuery().Before("gorm:row_query").Register("dd-trace-go:before_row_query", before)
cb.RowQuery().After("gorm:row_query").Register("dd-trace-go:after_row_query", afterFunc("gorm.row_query"))
cfg := new(config)
defaults(cfg)
for _, fn := range opts {
fn(cfg)
}
log.Debug("contrib/jinzhu/gorm: Adding Callbacks: %#v", cfg)
return db.Set(gormConfigKey, cfg)
}
// WithContext attaches the specified context to the given db. The context will
// be used as a basis for creating new spans. An example use case is providing
// a context which contains a span to be used as a parent.
func WithContext(ctx context.Context, db *gorm.DB) *gorm.DB {
if ctx == nil {
return db
}
db = db.Set(gormContextKey, ctx)
return db
}
// ContextFromDB returns any context previously attached to db using WithContext,
// otherwise returning context.Background.
func ContextFromDB(db *gorm.DB) context.Context {
if v, ok := db.Get(gormContextKey); ok {
if ctx, ok := v.(context.Context); ok {
return ctx
}
}
return context.Background()
}
func before(scope *gorm.Scope) {
scope.Set(gormSpanStartTimeKey, time.Now())
}
func after(scope *gorm.Scope, operationName string) {
v, ok := scope.Get(gormContextKey)
if !ok {
return
}
ctx := v.(context.Context)
v, ok = scope.Get(gormConfigKey)
if !ok {
return
}
cfg := v.(*config)
v, ok = scope.Get(gormSpanStartTimeKey)
if !ok {
return
}
t, ok := v.(time.Time)
if !ok {
return
}
opts := []ddtrace.StartSpanOption{
tracer.StartTime(t),
tracer.ServiceName(cfg.serviceName),
tracer.SpanType(ext.SpanTypeSQL),
tracer.ResourceName(scope.SQL),
tracer.Tag(ext.Component, componentName),
}
if !math.IsNaN(cfg.analyticsRate) {
opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
}
if cfg.tagFns != nil {
for key, tagFn := range cfg.tagFns {
opts = append(opts, tracer.Tag(key, tagFn(scope)))
}
}
span, _ := tracer.StartSpanFromContext(ctx, operationName, opts...)
defer span.Finish()
if cfg.errCheck(scope.DB().Error) {
span.SetTag(ext.Error, scope.DB().Error)
}
}