contrib/database/sql: sql comment tag injection experimental feature

README.md

@@ -80,5 +80,9 @@ might be running versions different from the vendored one, creating hard to debu
 To run integration tests locally, you should set the `INTEGRATION` environment variable. The dependencies of the integration tests are best run via Docker. To get an
 idea about the versions and the set-up take a look at our [CI config](./.circleci/config.yml).
 
-The best way to run the entire test suite is using the [CircleCI CLI](https://circleci.com/docs/2.0/local-jobs/). Simply run `circleci build`
-in the repository root. Note that you might have to increase the resources dedicated to Docker to around 4GB.
+The best way to run the entire test suite is using the [CircleCI CLI](https://circleci.com/docs/2.0/local-cli/). In order to run
+jobs locally, you'll first need to convert the Circle CI configuration to a format accepted by the `circleci` cli tool:
+  * `circleci config process .circleci/config.yml > process.yml` (from the repository root)
+
+Once you have a converted `process.yml`, simply run `circleci local execute -c process.yml --job <job-name>`. 
+Note that you might have to increase the resources dedicated to Docker to around 4GB.

contrib/database/sql/conn.go

@@ -15,6 +15,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/ddtrace/tracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 )
 
 var _ driver.Conn = (*tracedConn)(nil)
@@ -58,27 +59,34 @@ func (tc *tracedConn) BeginTx(ctx context.Context, opts driver.TxOptions) (tx dr
 
 func (tc *tracedConn) PrepareContext(ctx context.Context, query string) (stmt driver.Stmt, err error) {
 	start := time.Now()
+	mode := tc.cfg.commentInjectionMode
+	if mode == tracer.SQLInjectionModeFull {
+		// no context other than service in prepared statements
+		mode = tracer.SQLInjectionModeService
+	}
+	cquery, spanID := injectComments(ctx, query, mode)
 	if connPrepareCtx, ok := tc.Conn.(driver.ConnPrepareContext); ok {
-		stmt, err := connPrepareCtx.PrepareContext(ctx, query)
-		tc.tryTrace(ctx, queryTypePrepare, query, start, err)
+		stmt, err := connPrepareCtx.PrepareContext(ctx, cquery)
+		tc.tryTrace(ctx, queryTypePrepare, query, start, err, tracer.WithSpanID(spanID))
 		if err != nil {
 			return nil, err
 		}
-		return &tracedStmt{stmt, tc.traceParams, ctx, query}, nil
+		return &tracedStmt{Stmt: stmt, traceParams: tc.traceParams, ctx: ctx, query: query}, nil
 	}
-	stmt, err = tc.Prepare(query)
-	tc.tryTrace(ctx, queryTypePrepare, query, start, err)
+	stmt, err = tc.Prepare(cquery)
+	tc.tryTrace(ctx, queryTypePrepare, query, start, err, tracer.WithSpanID(spanID))
 	if err != nil {
 		return nil, err
 	}
-	return &tracedStmt{stmt, tc.traceParams, ctx, query}, nil
+	return &tracedStmt{Stmt: stmt, traceParams: tc.traceParams, ctx: ctx, query: query}, nil
 }
 
 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 {
-		r, err := execContext.ExecContext(ctx, query, args)
-		tc.tryTrace(ctx, queryTypeExec, query, start, err)
+		cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode)
+		r, err := execContext.ExecContext(ctx, cquery, args)
+		tc.tryTrace(ctx, queryTypeExec, query, start, err, tracer.WithSpanID(spanID))
 		return r, err
 	}
 	if execer, ok := tc.Conn.(driver.Execer); ok {
@@ -91,8 +99,9 @@ func (tc *tracedConn) ExecContext(ctx context.Context, query string, args []driv
 			return nil, ctx.Err()
 		default:
 		}
-		r, err = execer.Exec(query, dargs)
-		tc.tryTrace(ctx, queryTypeExec, query, start, err)
+		cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode)
+		r, err = execer.Exec(cquery, dargs)
+		tc.tryTrace(ctx, queryTypeExec, query, start, err, tracer.WithSpanID(spanID))
 		return r, err
 	}
 	return nil, driver.ErrSkip
@@ -111,8 +120,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 {
-		rows, err := queryerContext.QueryContext(ctx, query, args)
-		tc.tryTrace(ctx, queryTypeQuery, query, start, err)
+		cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode)
+		rows, err := queryerContext.QueryContext(ctx, cquery, args)
+		tc.tryTrace(ctx, queryTypeQuery, query, start, err, tracer.WithSpanID(spanID))
 		return rows, err
 	}
 	if queryer, ok := tc.Conn.(driver.Queryer); ok {
@@ -125,8 +135,9 @@ func (tc *tracedConn) QueryContext(ctx context.Context, query string, args []dri
 			return nil, ctx.Err()
 		default:
 		}
-		rows, err = queryer.Query(query, dargs)
-		tc.tryTrace(ctx, queryTypeQuery, query, start, err)
+		cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode)
+		rows, err = queryer.Query(cquery, dargs)
+		tc.tryTrace(ctx, queryTypeQuery, query, start, err, tracer.WithSpanID(spanID))
 		return rows, err
 	}
 	return nil, driver.ErrSkip
@@ -167,8 +178,28 @@ func WithSpanTags(ctx context.Context, tags map[string]string) context.Context {
 	return context.WithValue(ctx, spanTagsKey, tags)
 }
 
+// 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) {
+	// 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
+	// gets created.
+	var spanCtx ddtrace.SpanContext
+	if span, ok := tracer.SpanFromContext(ctx); ok {
+		spanCtx = span.Context()
+	}
+	carrier := tracer.SQLCommentCarrier{Query: query, Mode: mode}
+	if err := carrier.Inject(spanCtx); err != nil {
+		// this should never happen
+		log.Warn("contrib/database/sql: failed to inject query comments: %v", err)
+	}
+	return carrier.Query, carrier.SpanID
+}
+
 // 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) {
+func (tp *traceParams) tryTrace(ctx context.Context, qtype queryType, query string, startTime time.Time, err error, spanOpts ...ddtrace.StartSpanOption) {
 	if err == driver.ErrSkip {
 		// Not a user error: driver is telling sql package that an
 		// optional interface method is not implemented. There is
@@ -180,11 +211,11 @@ func (tp *traceParams) tryTrace(ctx context.Context, qtype queryType, query stri
 		return
 	}
 	name := fmt.Sprintf("%s.query", tp.driverName)
-	opts := []ddtrace.StartSpanOption{
+	opts := append(spanOpts,
 		tracer.ServiceName(tp.cfg.serviceName),
 		tracer.SpanType(ext.SpanTypeSQL),
 		tracer.StartTime(startTime),
-	}
+	)
 	if !math.IsNaN(tp.cfg.analyticsRate) {
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, tp.cfg.analyticsRate))
 	}

contrib/database/sql/injection_test.go

@@ -0,0 +1,169 @@
+// 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 sql
+
+import (
+	"context"
+	"database/sql"
+	"regexp"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql/internal"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+)
+
+func TestCommentInjection(t *testing.T) {
+	testCases := []struct {
+		name     string
+		opts     []RegisterOption
+		callDB   func(ctx context.Context, db *sql.DB) error
+		prepared []string
+		executed []*regexp.Regexp
+	}{
+		{
+			name: "prepare",
+			opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
+			callDB: func(ctx context.Context, db *sql.DB) error {
+				_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
+				return err
+			},
+			prepared: []string{"SELECT 1 from DUAL"},
+		},
+		{
+			name: "prepare-disabled",
+			opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
+			callDB: func(ctx context.Context, db *sql.DB) error {
+				_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
+				return err
+			},
+			prepared: []string{"SELECT 1 from DUAL"},
+		},
+		{
+			name: "prepare-service",
+			opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)},
+			callDB: func(ctx context.Context, db *sql.DB) error {
+				_, 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"},
+		},
+		{
+			name: "prepare-full",
+			opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
+			callDB: func(ctx context.Context, db *sql.DB) error {
+				_, 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"},
+		},
+		{
+			name: "query",
+			opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
+			callDB: func(ctx context.Context, db *sql.DB) error {
+				_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
+				return err
+			},
+			executed: []*regexp.Regexp{regexp.MustCompile("SELECT 1 from DUAL")},
+		},
+		{
+			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
+			},
+			executed: []*regexp.Regexp{regexp.MustCompile("SELECT 1 from DUAL")},
+		},
+		{
+			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
+			},
+			executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")},
+		},
+		{
+			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
+			},
+			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")},
+		},
+		{
+			name: "exec",
+			opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
+			callDB: func(ctx context.Context, db *sql.DB) error {
+				_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
+				return err
+			},
+			executed: []*regexp.Regexp{regexp.MustCompile("SELECT 1 from DUAL")},
+		},
+		{
+			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
+			},
+			executed: []*regexp.Regexp{regexp.MustCompile("SELECT 1 from DUAL")},
+		},
+		{
+			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
+			},
+			executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")},
+		},
+		{
+			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
+			},
+			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")},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			tracer.Start(tracer.WithService("test-service"), tracer.WithEnv("test-env"), tracer.WithServiceVersion("1.0.0"))
+			defer tracer.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)
+			require.Len(t, d.Prepared, len(tc.prepared))
+			for i, e := range tc.prepared {
+				assert.Equal(t, e, d.Prepared[i])
+			}
+
+			require.Len(t, d.Executed, len(tc.executed))
+			for i, e := range tc.executed {
+				assert.Regexp(t, e, d.Executed[i])
+				// the injected span ID should not be the parent's span ID
+				assert.NotContains(t, d.Executed[i], "ddsid='1'")
+			}
+		})
+	}
+}