Add support for custom attributes function

redis · Nov 25, 2024 · 64bfb34 · 64bfb34
1 parent e63669e
commit 64bfb34
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/extra/redisotel/config.go b/extra/redisotel/config.go
@@ -1,6 +1,8 @@
 package redisotel
 
 import (
+	"context"
+
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
@@ -11,8 +13,9 @@ import (
 type config struct {
 	// Common options.
 
-	dbSystem string
-	attrs    []attribute.KeyValue
+	dbSystem  string
+	attrs     []attribute.KeyValue
+	attrsFunc func(context.Context) []attribute.KeyValue
 
 	// Tracing options.
 
@@ -51,8 +54,9 @@ func (fn option) metrics() {}
 
 func newConfig(opts ...baseOption) *config {
 	conf := &config{
-		dbSystem: "redis",
-		attrs:    []attribute.KeyValue{},
+		dbSystem:  "redis",
+		attrs:     []attribute.KeyValue{},
+		attrsFunc: func(ctx context.Context) []attribute.KeyValue { return []attribute.KeyValue{} },
 
 		tp:            otel.GetTracerProvider(),
 		mp:            otel.GetMeterProvider(),
@@ -81,6 +85,14 @@ func WithAttributes(attrs ...attribute.KeyValue) Option {
 	})
 }
 
+// WithAttributesFunc takes a function that returns additional attributes to be added using the context.
+// This is executed only in ProcessPipelineHook and ProcessHook
+func WithAttributesFunc(f func(context.Context) []attribute.KeyValue) Option {
+	return option(func(conf *config) {
+		conf.attrsFunc = f
+	})
+}
+
 //------------------------------------------------------------------------------
 
 type TracingOption interface {

diff --git a/extra/redisotel/metrics.go b/extra/redisotel/metrics.go
@@ -175,6 +175,7 @@ func addMetricsHook(rdb *redis.Client, conf *config) error {
 		createTime: createTime,
 		useTime:    useTime,
 		attrs:      conf.attrs,
+		attrsFunc: conf.attrsFunc,
 	})
 	return nil
 }
@@ -183,6 +184,7 @@ type metricsHook struct {
 	createTime metric.Float64Histogram
 	useTime    metric.Float64Histogram
 	attrs      []attribute.KeyValue
+	attrsFunc func(context.Context) []attribute.KeyValue
 }
 
 var _ redis.Hook = (*metricsHook)(nil)
@@ -214,6 +216,7 @@ func (mh *metricsHook) ProcessHook(hook redis.ProcessHook) redis.ProcessHook {
 
 		attrs := make([]attribute.KeyValue, 0, len(mh.attrs)+2)
 		attrs = append(attrs, mh.attrs...)
+		attrs = append(attrs, mh.attrsFunc(ctx)...)
 		attrs = append(attrs, attribute.String("type", "command"))
 		attrs = append(attrs, statusAttr(err))
 
@@ -235,6 +238,7 @@ func (mh *metricsHook) ProcessPipelineHook(
 
 		attrs := make([]attribute.KeyValue, 0, len(mh.attrs)+2)
 		attrs = append(attrs, mh.attrs...)
+		attrs = append(attrs, mh.attrsFunc(ctx)...)
 		attrs = append(attrs, attribute.String("type", "pipeline"))
 		attrs = append(attrs, statusAttr(err))