Skip to content

Commit

Permalink
contrib/database/sql: SQL comment tag injection experimental feature (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
alexandre-normand committed Jun 15, 2022
1 parent f24f385 commit e2a9297
Show file tree
Hide file tree
Showing 12 changed files with 666 additions and 33 deletions.
8 changes: 6 additions & 2 deletions README.md
Expand Up @@ -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.
65 changes: 48 additions & 17 deletions contrib/database/sql/conn.go
Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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))
}
Expand Down
169 changes: 169 additions & 0 deletions 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'")
}
})
}
}

0 comments on commit e2a9297

Please sign in to comment.