diff --git a/contrib/database/sql/conn.go b/contrib/database/sql/conn.go index 587edc720b..46d4ef65f6 100644 --- a/contrib/database/sql/conn.go +++ b/contrib/database/sql/conn.go @@ -34,6 +34,10 @@ const ( queryTypeRollback = "Rollback" ) +const ( + keyDBMTraceInjected = "_dd.dbm_trace_injected" +) + type tracedConn struct { driver.Conn *traceParams @@ -64,17 +68,17 @@ func (tc *tracedConn) PrepareContext(ctx context.Context, query string) (stmt dr // no context other than service in prepared statements mode = tracer.SQLInjectionModeService } - cquery, spanID := injectComments(ctx, query, mode) + cquery, spanID := tc.injectComments(ctx, query, mode) if connPrepareCtx, ok := tc.Conn.(driver.ConnPrepareContext); ok { stmt, err := connPrepareCtx.PrepareContext(ctx, cquery) - tc.tryTrace(ctx, queryTypePrepare, query, start, err, tracer.WithSpanID(spanID)) + tc.tryTrace(ctx, queryTypePrepare, query, start, err, append(withDBMTraceInjectedTag(mode), tracer.WithSpanID(spanID))...) if err != nil { return nil, err } return &tracedStmt{Stmt: stmt, traceParams: tc.traceParams, ctx: ctx, query: query}, nil } stmt, err = tc.Prepare(cquery) - tc.tryTrace(ctx, queryTypePrepare, query, start, err, tracer.WithSpanID(spanID)) + tc.tryTrace(ctx, queryTypePrepare, query, start, err, append(withDBMTraceInjectedTag(mode), tracer.WithSpanID(spanID))...) if err != nil { return nil, err } @@ -84,9 +88,9 @@ func (tc *tracedConn) PrepareContext(ctx context.Context, query string) (stmt dr func (tc *tracedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (r driver.Result, err error) { start := time.Now() if execContext, ok := tc.Conn.(driver.ExecerContext); ok { - cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode) + cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode) r, err := execContext.ExecContext(ctx, cquery, args) - tc.tryTrace(ctx, queryTypeExec, query, start, err, tracer.WithSpanID(spanID)) + tc.tryTrace(ctx, queryTypeExec, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...) return r, err } if execer, ok := tc.Conn.(driver.Execer); ok { @@ -99,9 +103,9 @@ func (tc *tracedConn) ExecContext(ctx context.Context, query string, args []driv return nil, ctx.Err() default: } - cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode) + cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode) r, err = execer.Exec(cquery, dargs) - tc.tryTrace(ctx, queryTypeExec, query, start, err, tracer.WithSpanID(spanID)) + tc.tryTrace(ctx, queryTypeExec, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...) return r, err } return nil, driver.ErrSkip @@ -120,9 +124,9 @@ func (tc *tracedConn) Ping(ctx context.Context) (err error) { func (tc *tracedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) { start := time.Now() if queryerContext, ok := tc.Conn.(driver.QueryerContext); ok { - cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode) + cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode) rows, err := queryerContext.QueryContext(ctx, cquery, args) - tc.tryTrace(ctx, queryTypeQuery, query, start, err, tracer.WithSpanID(spanID)) + tc.tryTrace(ctx, queryTypeQuery, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...) return rows, err } if queryer, ok := tc.Conn.(driver.Queryer); ok { @@ -135,9 +139,9 @@ func (tc *tracedConn) QueryContext(ctx context.Context, query string, args []dri return nil, ctx.Err() default: } - cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode) + cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode) rows, err = queryer.Query(cquery, dargs) - tc.tryTrace(ctx, queryTypeQuery, query, start, err, tracer.WithSpanID(spanID)) + tc.tryTrace(ctx, queryTypeQuery, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...) return rows, err } return nil, driver.ErrSkip @@ -181,7 +185,7 @@ func WithSpanTags(ctx context.Context, tags map[string]string) context.Context { // injectComments returns the query with SQL comments injected according to the comment injection mode along // with a span ID injected into SQL comments. The returned span ID should be used when the SQL span is created // following the traced database call. -func injectComments(ctx context.Context, query string, mode tracer.SQLCommentInjectionMode) (cquery string, spanID uint64) { +func (tc *tracedConn) injectComments(ctx context.Context, query string, mode tracer.SQLCommentInjectionMode) (cquery string, spanID uint64) { // The sql span only gets created after the call to the database because we need to be able to skip spans // when a driver returns driver.ErrSkip. In order to work with those constraints, a new span id is generated and // used during SQL comment injection and returned for the sql span to be used later when/if the span @@ -190,7 +194,7 @@ func injectComments(ctx context.Context, query string, mode tracer.SQLCommentInj if span, ok := tracer.SpanFromContext(ctx); ok { spanCtx = span.Context() } - carrier := tracer.SQLCommentCarrier{Query: query, Mode: mode} + carrier := tracer.SQLCommentCarrier{Query: query, Mode: mode, DBServiceName: tc.cfg.serviceName} if err := carrier.Inject(spanCtx); err != nil { // this should never happen log.Warn("contrib/database/sql: failed to inject query comments: %v", err) @@ -198,6 +202,14 @@ func injectComments(ctx context.Context, query string, mode tracer.SQLCommentInj return carrier.Query, carrier.SpanID } +func withDBMTraceInjectedTag(mode tracer.SQLCommentInjectionMode) []tracer.StartSpanOption { + if mode == tracer.SQLInjectionModeFull { + return []tracer.StartSpanOption{tracer.Tag(keyDBMTraceInjected, true)} + } + + return nil +} + // tryTrace will create a span using the given arguments, but will act as a no-op when err is driver.ErrSkip. func (tp *traceParams) tryTrace(ctx context.Context, qtype queryType, query string, startTime time.Time, err error, spanOpts ...ddtrace.StartSpanOption) { if err == driver.ErrSkip { diff --git a/contrib/database/sql/injection_test.go b/contrib/database/sql/injection_test.go index 2d980d2e9e..65d47c4c23 100644 --- a/contrib/database/sql/injection_test.go +++ b/contrib/database/sql/injection_test.go @@ -15,6 +15,7 @@ import ( "github.com/stretchr/testify/require" "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql/internal" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" ) @@ -51,7 +52,7 @@ func TestCommentInjection(t *testing.T) { _, err := db.PrepareContext(ctx, "SELECT 1 from DUAL") return err }, - prepared: []string{"/*dde='test-env',ddsn='test-service',ddsv='1.0.0'*/ SELECT 1 from DUAL"}, + prepared: []string{"/*dddbs='test.db',dde='test-env',ddsn='test-service',ddsv='1.0.0'*/ SELECT 1 from DUAL"}, }, { name: "prepare-full", @@ -60,7 +61,7 @@ func TestCommentInjection(t *testing.T) { _, err := db.PrepareContext(ctx, "SELECT 1 from DUAL") return err }, - prepared: []string{"/*dde='test-env',ddsn='test-service',ddsv='1.0.0'*/ SELECT 1 from DUAL"}, + prepared: []string{"/*dddbs='test.db',dde='test-env',ddsn='test-service',ddsv='1.0.0'*/ SELECT 1 from DUAL"}, }, { name: "query", @@ -87,7 +88,7 @@ func TestCommentInjection(t *testing.T) { _, err := db.QueryContext(ctx, "SELECT 1 from DUAL") return err }, - executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")}, + executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")}, }, { name: "query-full", @@ -96,7 +97,7 @@ func TestCommentInjection(t *testing.T) { _, err := db.QueryContext(ctx, "SELECT 1 from DUAL") return err }, - executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsid='\\d+',ddsn='test-service',ddsp='1',ddsv='1.0.0',ddtid='1'\\*/ SELECT 1 from DUAL")}, + executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddsid='\\d+',ddsn='test-service',ddsp='1',ddsv='1.0.0',ddtid='1'\\*/ SELECT 1 from DUAL")}, }, { name: "exec", @@ -123,7 +124,7 @@ func TestCommentInjection(t *testing.T) { _, err := db.ExecContext(ctx, "SELECT 1 from DUAL") return err }, - executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")}, + executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")}, }, { name: "exec-full", @@ -132,7 +133,7 @@ func TestCommentInjection(t *testing.T) { _, err := db.ExecContext(ctx, "SELECT 1 from DUAL") return err }, - executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsid='\\d+',ddsn='test-service',ddsp='1',ddsv='1.0.0',ddtid='1'\\*/ SELECT 1 from DUAL")}, + executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddsid='\\d+',ddsn='test-service',ddsp='1',ddsv='1.0.0',ddtid='1'\\*/ SELECT 1 from DUAL")}, }, } @@ -167,3 +168,126 @@ func TestCommentInjection(t *testing.T) { }) } } + +func TestDBMTraceContextTagging(t *testing.T) { + testCases := []struct { + name string + opts []RegisterOption + callDB func(ctx context.Context, db *sql.DB) error + spanType string + traceContextInjectedTag bool + }{ + { + name: "prepare", + opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)}, + callDB: func(ctx context.Context, db *sql.DB) error { + _, err := db.PrepareContext(ctx, "SELECT 1 from DUAL") + return err + }, + spanType: queryTypePrepare, + traceContextInjectedTag: false, + }, + { + name: "query-disabled", + opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)}, + callDB: func(ctx context.Context, db *sql.DB) error { + _, err := db.QueryContext(ctx, "SELECT 1 from DUAL") + return err + }, + spanType: queryTypeQuery, + traceContextInjectedTag: false, + }, + { + name: "query-service", + opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)}, + callDB: func(ctx context.Context, db *sql.DB) error { + _, err := db.QueryContext(ctx, "SELECT 1 from DUAL") + return err + }, + spanType: queryTypeQuery, + traceContextInjectedTag: false, + }, + { + name: "query-full", + opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)}, + callDB: func(ctx context.Context, db *sql.DB) error { + _, err := db.QueryContext(ctx, "SELECT 1 from DUAL") + return err + }, + spanType: queryTypeQuery, + traceContextInjectedTag: true, + }, + { + name: "exec-disabled", + opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)}, + callDB: func(ctx context.Context, db *sql.DB) error { + _, err := db.ExecContext(ctx, "SELECT 1 from DUAL") + return err + }, + spanType: queryTypeExec, + traceContextInjectedTag: false, + }, + { + name: "exec-service", + opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)}, + callDB: func(ctx context.Context, db *sql.DB) error { + _, err := db.ExecContext(ctx, "SELECT 1 from DUAL") + return err + }, + spanType: queryTypeExec, + traceContextInjectedTag: false, + }, + { + name: "exec-full", + opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)}, + callDB: func(ctx context.Context, db *sql.DB) error { + _, err := db.ExecContext(ctx, "SELECT 1 from DUAL") + return err + }, + spanType: queryTypeExec, + traceContextInjectedTag: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + tr := mocktracer.Start() + defer tr.Stop() + + d := &internal.MockDriver{} + Register("test", d, tc.opts...) + defer unregister("test") + + db, err := Open("test", "dn") + require.NoError(t, err) + + s, ctx := tracer.StartSpanFromContext(context.Background(), "test.call", tracer.WithSpanID(1)) + err = tc.callDB(ctx, db) + s.Finish() + + require.NoError(t, err) + spans := tr.FinishedSpans() + + sps := spansOfType(spans, tc.spanType) + for _, s := range sps { + tags := s.Tags() + if tc.traceContextInjectedTag { + assert.Equal(t, tags[keyDBMTraceInjected], true) + } else { + _, ok := tags[keyDBMTraceInjected] + assert.False(t, ok) + } + } + }) + } +} + +func spansOfType(spans []mocktracer.Span, spanType string) (filtered []mocktracer.Span) { + filtered = make([]mocktracer.Span, 0) + for _, s := range spans { + if s.Tag("sql.query_type") == spanType { + filtered = append(filtered, s) + } + } + return filtered +} diff --git a/ddtrace/tracer/sqlcomment.go b/ddtrace/tracer/sqlcomment.go index 05cd0263da..09da3d3ff1 100644 --- a/ddtrace/tracer/sqlcomment.go +++ b/ddtrace/tracer/sqlcomment.go @@ -12,6 +12,7 @@ import ( "gopkg.in/DataDog/dd-trace-go.v1/ddtrace" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext" "gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig" + "gopkg.in/DataDog/dd-trace-go.v1/internal/log" ) // SQLCommentInjectionMode represents the mode of SQL comment injection. @@ -34,6 +35,7 @@ const ( sqlCommentTraceID = "ddtid" sqlCommentSpanID = "ddsid" sqlCommentService = "ddsn" + sqlCommentDBService = "dddbs" sqlCommentVersion = "ddsv" sqlCommentEnv = "dde" ) @@ -42,9 +44,10 @@ const ( // of a sqlcommenter formatted comment prepended to the original query text. // See https://google.github.io/sqlcommenter/spec/ for more details. type SQLCommentCarrier struct { - Query string - Mode SQLCommentInjectionMode - SpanID uint64 + Query string + Mode SQLCommentInjectionMode + DBServiceName string + SpanID uint64 } // Inject injects a span context in the carrier's Query field as a comment. @@ -93,6 +96,7 @@ func (c *SQLCommentCarrier) Inject(spanCtx ddtrace.SpanContext) error { if version != "" { tags[sqlCommentVersion] = version } + tags[sqlCommentDBService] = c.DBServiceName } c.Query = commentQuery(c.Query, tags) return nil @@ -113,7 +117,7 @@ func commentQuery(query string, tags map[string]string) string { var b strings.Builder // the sqlcommenter specification dictates that tags should be sorted. Since we know all injected keys, // we skip a sorting operation by specifying the order of keys statically - orderedKeys := []string{sqlCommentEnv, sqlCommentSpanID, sqlCommentService, sqlCommentKeySamplingPriority, sqlCommentVersion, sqlCommentTraceID} + orderedKeys := []string{sqlCommentDBService, sqlCommentEnv, sqlCommentSpanID, sqlCommentService, sqlCommentKeySamplingPriority, sqlCommentVersion, sqlCommentTraceID} first := true for _, k := range orderedKeys { if v, ok := tags[k]; ok { @@ -141,6 +145,7 @@ func commentQuery(query string, tags map[string]string) string { if query == "" { return b.String() } + log.Debug("Injected sql comment: %s", b.String()) b.WriteRune(' ') b.WriteString(query) return b.String() diff --git a/ddtrace/tracer/sqlcomment_test.go b/ddtrace/tracer/sqlcomment_test.go index 449e11c0ee..86bfb2b14f 100644 --- a/ddtrace/tracer/sqlcomment_test.go +++ b/ddtrace/tracer/sqlcomment_test.go @@ -31,7 +31,7 @@ func TestSQLCommentCarrier(t *testing.T) { query: "SELECT * from FOO", mode: SQLInjectionModeFull, injectSpan: true, - expectedQuery: "/*dde='test-env',ddsid='',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsp='2',ddsv='1.0.0',ddtid='10'*/ SELECT * from FOO", + expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddsid='',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsp='2',ddsv='1.0.0',ddtid='10'*/ SELECT * from FOO", expectedSpanIDGen: true, }, { @@ -39,14 +39,14 @@ func TestSQLCommentCarrier(t *testing.T) { query: "SELECT * from FOO", mode: SQLInjectionModeService, injectSpan: true, - expectedQuery: "/*dde='test-env',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsv='1.0.0'*/ SELECT * from FOO", + expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsv='1.0.0'*/ SELECT * from FOO", expectedSpanIDGen: false, }, { name: "no-trace", query: "SELECT * from FOO", mode: SQLInjectionModeFull, - expectedQuery: "/*ddsid='',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsp='0',ddtid=''*/ SELECT * from FOO", + expectedQuery: "/*dddbs='whiskey-db',ddsid='',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsp='0',ddtid=''*/ SELECT * from FOO", expectedSpanIDGen: true, }, { @@ -54,7 +54,7 @@ func TestSQLCommentCarrier(t *testing.T) { query: "", mode: SQLInjectionModeFull, injectSpan: true, - expectedQuery: "/*dde='test-env',ddsid='',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsp='2',ddsv='1.0.0',ddtid='10'*/", + expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddsid='',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsp='2',ddsv='1.0.0',ddtid='10'*/", expectedSpanIDGen: true, }, { @@ -62,7 +62,7 @@ func TestSQLCommentCarrier(t *testing.T) { query: "SELECT * from FOO -- test query", mode: SQLInjectionModeFull, injectSpan: true, - expectedQuery: "/*dde='test-env',ddsid='',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsp='2',ddsv='1.0.0',ddtid='10'*/ SELECT * from FOO -- test query", + expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddsid='',ddsn='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddsp='2',ddsv='1.0.0',ddtid='10'*/ SELECT * from FOO -- test query", expectedSpanIDGen: true, }, } @@ -81,7 +81,7 @@ func TestSQLCommentCarrier(t *testing.T) { spanCtx = root.Context() } - carrier := SQLCommentCarrier{Query: tc.query, Mode: tc.mode} + carrier := SQLCommentCarrier{Query: tc.query, Mode: tc.mode, DBServiceName: "whiskey-db"} err := carrier.Inject(spanCtx) require.NoError(t, err)