Skip to content

Commit

Permalink
contrib/database/sql: Rename SQLCommentInjection to DBMPropagation (#…
Browse files Browse the repository at this point in the history 
…1556)
  • Loading branch information
@alexandre-normand
alexandre-normand committed Nov 3, 2022
1 parent ba96658 commit 61c5926
Show file tree
Hide file tree
Showing 7 changed files with 120 additions and 65 deletions.
29 changes: 14 additions & 15 deletions contrib/database/sql/conn.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,10 +63,10 @@ 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 {
mode := tc.cfg.dbmPropagationMode
if mode == tracer.DBMPropagationModeFull {
// no context other than service in prepared statements
mode = tracer.SQLInjectionModeService
mode = tracer.DBMPropagationModeService
}
cquery, spanID := tc.injectComments(ctx, query, mode)
if connPrepareCtx, ok := tc.Conn.(driver.ConnPrepareContext); ok {
Expand All @@ -88,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 := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode)
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.dbmPropagationMode)
r, err := execContext.ExecContext(ctx, cquery, args)
tc.tryTrace(ctx, queryTypeExec, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...)
tc.tryTrace(ctx, queryTypeExec, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.dbmPropagationMode), tracer.WithSpanID(spanID))...)
return r, err
}
if execer, ok := tc.Conn.(driver.Execer); ok {
Expand All @@ -103,9 +103,9 @@ func (tc *tracedConn) ExecContext(ctx context.Context, query string, args []driv
return nil, ctx.Err()
default:
}
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode)
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.dbmPropagationMode)
r, err = execer.Exec(cquery, dargs)
tc.tryTrace(ctx, queryTypeExec, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...)
tc.tryTrace(ctx, queryTypeExec, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.dbmPropagationMode), tracer.WithSpanID(spanID))...)
return r, err
}
return nil, driver.ErrSkip
Expand All @@ -124,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 := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode)
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.dbmPropagationMode)
rows, err := queryerContext.QueryContext(ctx, cquery, args)
tc.tryTrace(ctx, queryTypeQuery, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...)
tc.tryTrace(ctx, queryTypeQuery, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.dbmPropagationMode), tracer.WithSpanID(spanID))...)
return rows, err
}
if queryer, ok := tc.Conn.(driver.Queryer); ok {
Expand All @@ -139,9 +139,9 @@ func (tc *tracedConn) QueryContext(ctx context.Context, query string, args []dri
return nil, ctx.Err()
default:
}
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode)
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.dbmPropagationMode)
rows, err = queryer.Query(cquery, dargs)
tc.tryTrace(ctx, queryTypeQuery, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...)
tc.tryTrace(ctx, queryTypeQuery, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.dbmPropagationMode), tracer.WithSpanID(spanID))...)
return rows, err
}
return nil, driver.ErrSkip
Expand Down Expand Up @@ -185,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 (tc *tracedConn) injectComments(ctx context.Context, query string, mode tracer.SQLCommentInjectionMode) (cquery string, spanID uint64) {
func (tc *tracedConn) injectComments(ctx context.Context, query string, mode tracer.DBMPropagationMode) (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
Expand All @@ -202,11 +202,10 @@ func (tc *tracedConn) injectComments(ctx context.Context, query string, mode tra
return carrier.Query, carrier.SpanID
}

func withDBMTraceInjectedTag(mode tracer.SQLCommentInjectionMode) []tracer.StartSpanOption {
if mode == tracer.SQLInjectionModeFull {
func withDBMTraceInjectedTag(mode tracer.DBMPropagationMode) []tracer.StartSpanOption {
if mode == tracer.DBMPropagationModeFull {
return []tracer.StartSpanOption{tracer.Tag(keyDBMTraceInjected, true)}
}

return nil
}

Expand Down
19 changes: 19 additions & 0 deletions contrib/database/sql/example_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,3 +88,22 @@ func Example_sqlite() {
rows.Close()
span.Finish(tracer.WithError(err))
}

func Example_dbmPropagation() {
// The first step is to set the dbm propagation mode when registering the driver. Note that this can also
// be done on sqltrace.Open for more granular control over the feature.
sqltrace.Register("postgres", &pq.Driver{}, sqltrace.WithDBMPropagation(tracer.DBMPropagationModeFull))

// Followed by a call to Open.
db, err := sqltrace.Open("postgres", "postgres://pqgotest:password@localhost/pqgotest?sslmode=disable")
if err != nil {
log.Fatal(err)
}

// Then, we continue using the database/sql package as we normally would, with tracing.
rows, err := db.Query("SELECT name FROM users WHERE age=?", 27)
if err != nil {
log.Fatal(err)
}
defer rows.Close()
}
37 changes: 27 additions & 10 deletions contrib/database/sql/option.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,13 @@ import (
)

type config struct {
serviceName string
analyticsRate float64
dsn string
childSpansOnly bool
errCheck func(err error) bool
tags map[string]interface{}
commentInjectionMode tracer.SQLCommentInjectionMode
serviceName string
analyticsRate float64
dsn string
childSpansOnly bool
errCheck func(err error) bool
tags map[string]interface{}
dbmPropagationMode tracer.DBMPropagationMode
}

// Option represents an option that can be passed to Register, Open or OpenDB.
Expand All @@ -39,7 +39,11 @@ func defaults(cfg *config) {
} else {
cfg.analyticsRate = math.NaN()
}
cfg.commentInjectionMode = tracer.SQLCommentInjectionMode(os.Getenv("DD_TRACE_SQL_COMMENT_INJECTION_MODE"))
mode := os.Getenv("DD_DBM_PROPAGATION_MODE")
if mode == "" {
mode = os.Getenv("DD_TRACE_SQL_COMMENT_INJECTION_MODE")
}
cfg.dbmPropagationMode = tracer.DBMPropagationMode(mode)
}

// WithServiceName sets the given service name when registering a driver,
Expand Down Expand Up @@ -111,9 +115,22 @@ func WithCustomTag(key string, value interface{}) Option {

// WithSQLCommentInjection enables injection of tags as sql comments on traced queries.
// This includes dynamic values like span id, trace id and sampling priority which can make queries
// unique for some cache implementations. Use WithStaticTagsCommentInjection if this is a concern.
// unique for some cache implementations.
//
// Deprecated: Use WithDBMPropagation instead.
func WithSQLCommentInjection(mode tracer.SQLCommentInjectionMode) Option {
return WithDBMPropagation(tracer.DBMPropagationMode(mode))
}

// WithDBMPropagation enables injection of tags as sql comments on traced queries.
// This includes dynamic values like span id, trace id and the sampled flag which can make queries
// unique for some cache implementations. Use DBMPropagationModeService if this is a concern.
//
// Note that enabling sql comment propagation results in potentially confidential data (service names)
// being stored in the databases which can then be accessed by other 3rd parties that have been granted
// access to the database.
func WithDBMPropagation(mode tracer.DBMPropagationMode) Option {
return func(cfg *config) {
cfg.commentInjectionMode = mode
cfg.dbmPropagationMode = mode
}
}
42 changes: 21 additions & 21 deletions contrib/database/sql/injection_test.go → contrib/database/sql/propagation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import (
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

func TestCommentInjection(t *testing.T) {
func TestDBMPropagation(t *testing.T) {
testCases := []struct {
name string
opts []RegisterOption
Expand All @@ -29,7 +29,7 @@ func TestCommentInjection(t *testing.T) {
}{
{
name: "prepare",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -38,7 +38,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "prepare-disabled",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -47,7 +47,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "prepare-service",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeService)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -56,7 +56,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "prepare-full",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -65,7 +65,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "query",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -74,7 +74,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "query-disabled",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -83,7 +83,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "query-service",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeService)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -92,7 +92,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "query-full",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -101,7 +101,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "exec",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -110,7 +110,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "exec-disabled",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -119,7 +119,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "exec-service",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeService)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -128,7 +128,7 @@ func TestCommentInjection(t *testing.T) {
},
{
name: "exec-full",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
Expand Down Expand Up @@ -179,7 +179,7 @@ func TestDBMTraceContextTagging(t *testing.T) {
}{
{
name: "prepare",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -189,7 +189,7 @@ func TestDBMTraceContextTagging(t *testing.T) {
},
{
name: "query-disabled",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -199,7 +199,7 @@ func TestDBMTraceContextTagging(t *testing.T) {
},
{
name: "query-service",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeService)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -209,7 +209,7 @@ func TestDBMTraceContextTagging(t *testing.T) {
},
{
name: "query-full",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -219,7 +219,7 @@ func TestDBMTraceContextTagging(t *testing.T) {
},
{
name: "exec-disabled",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -229,7 +229,7 @@ func TestDBMTraceContextTagging(t *testing.T) {
},
{
name: "exec-service",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeService)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
Expand All @@ -239,7 +239,7 @@ func TestDBMTraceContextTagging(t *testing.T) {
},
{
name: "exec-full",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
Expand Down Expand Up @@ -272,7 +272,7 @@ func TestDBMTraceContextTagging(t *testing.T) {
for _, s := range sps {
tags := s.Tags()
if tc.traceContextInjectedTag {
assert.Equal(t, tags[keyDBMTraceInjected], true)
assert.Equal(t, true, tags[keyDBMTraceInjected])
} else {
_, ok := tags[keyDBMTraceInjected]
assert.False(t, ok)
Expand Down
4 changes: 2 additions & 2 deletions contrib/database/sql/sql.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -204,8 +204,8 @@ func OpenDB(c driver.Connector, opts ...Option) *sql.DB {
if cfg.errCheck == nil {
cfg.errCheck = rc.errCheck
}
if cfg.commentInjectionMode == tracer.SQLInjectionUndefined {
cfg.commentInjectionMode = rc.commentInjectionMode
if cfg.dbmPropagationMode == tracer.DBMPropagationModeUndefined {
cfg.dbmPropagationMode = rc.dbmPropagationMode
}
cfg.childSpansOnly = rc.childSpansOnly
tc := &tracedConnector{
Expand Down
Loading

0 comments on commit 61c5926

Please sign in to comment.