From 3129e09b29f37ac8b8a1256deef1883140dbb255 Mon Sep 17 00:00:00 2001
From: Mitch Usher <musher@bolt.com>
Date: Thu, 25 Aug 2022 20:07:47 -0600
Subject: [PATCH] Update redis otel with option to not set raw command as an
 attribute

---
 extra/redisotel/redisotel.go      | 39 +++++++++++++++++++++----------
 extra/redisotel/redisotel_test.go | 27 +++++++++++++++++++++
 2 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/extra/redisotel/redisotel.go b/extra/redisotel/redisotel.go
index 53f7ec3a..94759ebc 100644
--- a/extra/redisotel/redisotel.go
+++ b/extra/redisotel/redisotel.go
@@ -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
+	})
+}
diff --git a/extra/redisotel/redisotel_test.go b/extra/redisotel/redisotel_test.go
index 4de0e713..8c0e4871 100644
--- a/extra/redisotel/redisotel_test.go
+++ b/extra/redisotel/redisotel_test.go
@@ -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")
+		}
+	}
+}