@@ -2,159 +2,18 @@ package instrumentation
22
33import (
44 "context"
5- "fmt"
6- "strings"
75
8- "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
9- ddtrace "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
6+ "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
107 "gorm.io/gorm"
118)
129
13- type spanContextKey string
14-
15- var (
16- // ParentSpanGormKey is the name of the parent span key
17- parentSpanGormKey = spanContextKey ("trancingParentSpan" )
18- // SpanGormKey is the name of the span key
19- spanGormKey = spanContextKey ("tracingSpan" )
20- )
21-
22- // TraceDatabase sets span to gorm settings, returns cloned DB
10+ // TraceDatabase fetches the span from the context and injects it to a new
11+ // database session's statement context.
2312func TraceDatabase (ctx context.Context , db * gorm.DB ) * gorm.DB {
2413 if ctx == nil {
2514 return db
2615 }
2716
28- parentSpan , _ := ddtrace .SpanFromContext (ctx )
29-
30- return db .Session (& gorm.Session {
31- Context : context .WithValue (db .Statement .Context , parentSpanGormKey , parentSpan ),
32- })
33- }
34-
35- // InstrumentDatabase adds callbacks for tracing, call TraceDatabase to make it work
36- func InstrumentDatabase (db * gorm.DB , appName string ) {
37- callbacks := newCallbacks (appName )
38-
39- registerCallbacks (db , "create" , callbacks )
40- registerCallbacks (db , "query" , callbacks )
41- registerCallbacks (db , "update" , callbacks )
42- registerCallbacks (db , "delete" , callbacks )
43- registerCallbacks (db , "row" , callbacks )
44- }
45-
46- type callbacks struct {
47- serviceName string
48- }
49-
50- func newCallbacks (appName string ) * callbacks {
51- return & callbacks {
52- serviceName : fmt .Sprintf ("%s-%s" , appName , "mysql" ),
53- }
54- }
55-
56- func (c * callbacks ) beforeCreate (db * gorm.DB ) { c .before (db , "INSERT" , c .serviceName ) }
57- func (c * callbacks ) afterCreate (db * gorm.DB ) { c .after (db ) }
58- func (c * callbacks ) beforeQuery (db * gorm.DB ) { c .before (db , "SELECT" , c .serviceName ) }
59- func (c * callbacks ) afterQuery (db * gorm.DB ) { c .after (db ) }
60- func (c * callbacks ) beforeUpdate (db * gorm.DB ) { c .before (db , "UPDATE" , c .serviceName ) }
61- func (c * callbacks ) afterUpdate (db * gorm.DB ) { c .after (db ) }
62- func (c * callbacks ) beforeDelete (db * gorm.DB ) { c .before (db , "DELETE" , c .serviceName ) }
63- func (c * callbacks ) afterDelete (db * gorm.DB ) { c .after (db ) }
64- func (c * callbacks ) beforeRow (db * gorm.DB ) { c .before (db , "" , c .serviceName ) }
65- func (c * callbacks ) afterRow (db * gorm.DB ) { c .after (db ) }
66- func (c * callbacks ) before (db * gorm.DB , operationName string , serviceName string ) {
67- if db .Statement == nil || db .Statement .Context == nil {
68- return
69- }
70-
71- parentSpan , ok := db .Statement .Context .Value (parentSpanGormKey ).(ddtrace.Span )
72- if ! ok {
73- return
74- }
75-
76- spanOpts := []ddtrace.StartSpanOption {
77- ddtrace .ChildOf (parentSpan .Context ()),
78- ddtrace .SpanType (ext .SpanTypeSQL ),
79- ddtrace .ServiceName (serviceName ),
80- }
81- if operationName == "" {
82- operationName = strings .Split (db .Statement .SQL .String (), " " )[0 ]
83- }
84- sp := ddtrace .StartSpan (operationName , spanOpts ... )
85- db .Statement .Context = context .WithValue (db .Statement .Context , spanGormKey , sp )
86- }
87-
88- func (c * callbacks ) after (db * gorm.DB ) {
89- if db .Statement == nil || db .Statement .Context == nil {
90- return
91- }
92-
93- sp , ok := db .Statement .Context .Value (spanGormKey ).(ddtrace.Span )
94- if ! ok {
95- return
96- }
97-
98- sp .SetTag (ext .ResourceName , strings .ToUpper (db .Statement .SQL .String ()))
99- sp .SetTag ("db.table" , db .Statement .Table )
100- sp .SetTag ("db.query" , strings .ToUpper (db .Statement .SQL .String ()))
101- sp .SetTag ("db.err" , db .Error )
102- sp .SetTag ("db.count" , db .RowsAffected )
103- sp .Finish ()
104- }
105-
106- func registerCallbacks (db * gorm.DB , name string , c * callbacks ) {
107- var err error
108-
109- beforeName := fmt .Sprintf ("tracing:%v_before" , name )
110- afterName := fmt .Sprintf ("tracing:%v_after" , name )
111- gormCallbackName := fmt .Sprintf ("gorm:%v" , name )
112- // gorm does some magic, if you pass CallbackProcessor here - nothing works
113- switch name {
114- case "create" :
115- err = db .Callback ().Create ().Before (gormCallbackName ).Register (beforeName , c .beforeCreate )
116- if err != nil {
117- return
118- }
119- err = db .Callback ().Create ().After (gormCallbackName ).Register (afterName , c .afterCreate )
120- if err != nil {
121- return
122- }
123- case "query" :
124- err = db .Callback ().Query ().Before (gormCallbackName ).Register (beforeName , c .beforeQuery )
125- if err != nil {
126- return
127- }
128- err = db .Callback ().Query ().After (gormCallbackName ).Register (afterName , c .afterQuery )
129- if err != nil {
130- return
131- }
132- case "update" :
133- err = db .Callback ().Update ().Before (gormCallbackName ).Register (beforeName , c .beforeUpdate )
134- if err != nil {
135- return
136- }
137- err = db .Callback ().Update ().After (gormCallbackName ).Register (afterName , c .afterUpdate )
138- if err != nil {
139- return
140- }
141- case "delete" :
142- err = db .Callback ().Delete ().Before (gormCallbackName ).Register (beforeName , c .beforeDelete )
143- if err != nil {
144- return
145- }
146- err = db .Callback ().Delete ().After (gormCallbackName ).Register (afterName , c .afterDelete )
147- if err != nil {
148- return
149- }
150- case "row" :
151- err = db .Callback ().Row ().Before (gormCallbackName ).Register (beforeName , c .beforeRow )
152- if err != nil {
153- return
154- }
155- err = db .Callback ().Row ().After (gormCallbackName ).Register (afterName , c .afterRow )
156- if err != nil {
157- return
158- }
159- }
17+ parentSpan , _ := tracer .SpanFromContext (ctx )
18+ return db .WithContext (tracer .ContextWithSpan (ctx , parentSpan ))
16019}
0 commit comments