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 {
tracer trace.Tracer
attrs []attribute.KeyValue
tracer trace.Tracer
attrs []attribute.KeyValue
dbStmtEnabled bool
}
func NewTracingHook(opts ...Option) *TracingHook {
@ -28,6 +29,7 @@ func NewTracingHook(opts ...Option) *TracingHook {
attrs: []attribute.KeyValue{
semconv.DBSystemRedis,
},
dbStmtEnabled: true,
}
for _, opt := range opts {
opt.apply(cfg)
@ -37,7 +39,7 @@ func NewTracingHook(opts ...Option) *TracingHook {
defaultTracerName,
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) {
@ -48,9 +50,10 @@ func (th *TracingHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (cont
opts := []trace.SpanStartOption{
trace.WithSpanKind(trace.SpanKindClient),
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...)
@ -67,22 +70,26 @@ func (th *TracingHook) AfterProcess(ctx context.Context, cmd redis.Cmder) error
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() {
return ctx, nil
}
summary, cmdsString := rediscmd.CmdsString(cmds)
opts := []trace.SpanStartOption{
trace.WithSpanKind(trace.SpanKindClient),
trace.WithAttributes(th.attrs...),
trace.WithAttributes(
semconv.DBStatementKey.String(cmdsString),
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...)
return ctx, nil
@ -105,8 +112,9 @@ func recordError(ctx context.Context, span trace.Span, err error) {
}
type config struct {
tp trace.TracerProvider
attrs []attribute.KeyValue
tp trace.TracerProvider
attrs []attribute.KeyValue
dbStmtEnabled bool
}
// Option specifies instrumentation configuration options.
@ -136,3 +144,10 @@ func WithAttributes(attrs ...attribute.KeyValue) Option {
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])
}
}
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")
}
}
}