Update redis otel with option to not set raw command as an attribute

This commit is contained in:
Mitch Usher 2022-08-25 20:07:47 -06:00 committed by Vladimir Mihailenco
parent c561f3ca7e
commit 3129e09b29
2 changed files with 54 additions and 12 deletions

View File

@ -18,8 +18,9 @@ const (
) )
type TracingHook struct { type TracingHook struct {
tracer trace.Tracer tracer trace.Tracer
attrs []attribute.KeyValue attrs []attribute.KeyValue
dbStmtEnabled bool
} }
func NewTracingHook(opts ...Option) *TracingHook { func NewTracingHook(opts ...Option) *TracingHook {
@ -28,6 +29,7 @@ func NewTracingHook(opts ...Option) *TracingHook {
attrs: []attribute.KeyValue{ attrs: []attribute.KeyValue{
semconv.DBSystemRedis, semconv.DBSystemRedis,
}, },
dbStmtEnabled: true,
} }
for _, opt := range opts { for _, opt := range opts {
opt.apply(cfg) opt.apply(cfg)
@ -37,7 +39,7 @@ func NewTracingHook(opts ...Option) *TracingHook {
defaultTracerName, defaultTracerName,
trace.WithInstrumentationVersion("semver:"+redis.Version()), trace.WithInstrumentationVersion("semver:"+redis.Version()),
) )
return &TracingHook{tracer: tracer, attrs: cfg.attrs} return &TracingHook{tracer: tracer, attrs: cfg.attrs, dbStmtEnabled: cfg.dbStmtEnabled}
} }
func (th *TracingHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) { func (th *TracingHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) {
@ -48,9 +50,10 @@ func (th *TracingHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (cont
opts := []trace.SpanStartOption{ opts := []trace.SpanStartOption{
trace.WithSpanKind(trace.SpanKindClient), trace.WithSpanKind(trace.SpanKindClient),
trace.WithAttributes(th.attrs...), trace.WithAttributes(th.attrs...),
trace.WithAttributes( }
semconv.DBStatementKey.String(rediscmd.CmdString(cmd)),
), if th.dbStmtEnabled {
opts = append(opts, trace.WithAttributes(semconv.DBStatementKey.String(rediscmd.CmdString(cmd))))
} }
ctx, _ = th.tracer.Start(ctx, cmd.FullName(), opts...) ctx, _ = th.tracer.Start(ctx, cmd.FullName(), opts...)
@ -67,22 +70,26 @@ func (th *TracingHook) AfterProcess(ctx context.Context, cmd redis.Cmder) error
return nil return nil
} }
func (th *TracingHook) BeforeProcessPipeline(ctx context.Context, cmds []redis.Cmder) (context.Context, error) { func (th *TracingHook) BeforeProcessPipeline(
ctx context.Context, cmds []redis.Cmder,
) (context.Context, error) {
if !trace.SpanFromContext(ctx).IsRecording() { if !trace.SpanFromContext(ctx).IsRecording() {
return ctx, nil return ctx, nil
} }
summary, cmdsString := rediscmd.CmdsString(cmds)
opts := []trace.SpanStartOption{ opts := []trace.SpanStartOption{
trace.WithSpanKind(trace.SpanKindClient), trace.WithSpanKind(trace.SpanKindClient),
trace.WithAttributes(th.attrs...), trace.WithAttributes(th.attrs...),
trace.WithAttributes( trace.WithAttributes(
semconv.DBStatementKey.String(cmdsString),
attribute.Int("db.redis.num_cmd", len(cmds)), attribute.Int("db.redis.num_cmd", len(cmds)),
), ),
} }
summary, cmdsString := rediscmd.CmdsString(cmds)
if th.dbStmtEnabled {
opts = append(opts, trace.WithAttributes(semconv.DBStatementKey.String(cmdsString)))
}
ctx, _ = th.tracer.Start(ctx, "pipeline "+summary, opts...) ctx, _ = th.tracer.Start(ctx, "pipeline "+summary, opts...)
return ctx, nil return ctx, nil
@ -105,8 +112,9 @@ func recordError(ctx context.Context, span trace.Span, err error) {
} }
type config struct { type config struct {
tp trace.TracerProvider tp trace.TracerProvider
attrs []attribute.KeyValue attrs []attribute.KeyValue
dbStmtEnabled bool
} }
// Option specifies instrumentation configuration options. // Option specifies instrumentation configuration options.
@ -136,3 +144,10 @@ func WithAttributes(attrs ...attribute.KeyValue) Option {
cfg.attrs = append(cfg.attrs, attrs...) cfg.attrs = append(cfg.attrs, attrs...)
}) })
} }
// WithDBStatement tells the tracing hook not to log raw redis commands.
func WithDBStatement(on bool) Option {
return optionFunc(func(cfg *config) {
cfg.dbStmtEnabled = on
})
}

View File

@ -68,3 +68,30 @@ func TestNewWithAttributes(t *testing.T) {
t.Fatalf("expected attrs[2] to be semconv.DBStatementKey.String(\"ping\"), got: %v", attrs[2]) t.Fatalf("expected attrs[2] to be semconv.DBStatementKey.String(\"ping\"), got: %v", attrs[2])
} }
} }
func TestWithDBStatement(t *testing.T) {
provider := sdktrace.NewTracerProvider()
hook := redisotel.NewTracingHook(
redisotel.WithTracerProvider(provider),
redisotel.WithDBStatement(false),
)
ctx, span := provider.Tracer("redis-test").Start(context.TODO(), "redis-test")
cmd := redis.NewCmd(ctx, "ping")
defer span.End()
ctx, err := hook.BeforeProcess(ctx, cmd)
if err != nil {
t.Fatal(err)
}
err = hook.AfterProcess(ctx, cmd)
if err != nil {
t.Fatal(err)
}
attrs := trace.SpanFromContext(ctx).(sdktrace.ReadOnlySpan).Attributes()
for _, attr := range attrs {
if attr.Key == semconv.DBStatementKey {
t.Fatal("Attribute with db statement should not exist")
}
}
}