feat(bunotel): add options for set otel providers (#836)

uptrace · Sep 10, 2023 · 806e632 · 806e632
1 parent 16367aa
commit 806e632
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 15 deletions.
diff --git a/extra/bunotel/option.go b/extra/bunotel/option.go
@@ -2,7 +2,9 @@ package bunotel
 
 import (
 	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
 	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
+	"go.opentelemetry.io/otel/trace"
 )
 
 type Option func(h *QueryHook)
@@ -30,3 +32,23 @@ func WithFormattedQueries(format bool) Option {
 		h.formatQueries = format
 	}
 }
+
+// WithTracerProvider returns an Option to use the TracerProvider when
+// creating a Tracer.
+func WithTracerProvider(tp trace.TracerProvider) Option {
+	return func(h *QueryHook) {
+		if tp != nil {
+			h.tracer = tp.Tracer("github.com/uptrace/bun")
+		}
+	}
+}
+
+// WithMeterProvider returns an Option to use the MeterProvider when
+// creating a Meter.
+func WithMeterProvider(mp metric.MeterProvider) Option {
+	return func(h *QueryHook) {
+		if mp != nil {
+			h.meter = mp.Meter("github.com/uptrace/bun")
+		}
+	}
+}
diff --git a/extra/bunotel/otel.go b/extra/bunotel/otel.go
@@ -20,20 +20,12 @@ import (
 	"github.com/uptrace/opentelemetry-go-extra/otelsql"
 )
 
-var (
-	tracer = otel.Tracer("github.com/uptrace/bun")
-	meter  = otel.Meter("github.com/uptrace/bun")
-
-	queryHistogram, _ = meter.Int64Histogram(
-		"go.sql.query_timing",
-		metric.WithDescription("Timing of processed queries"),
-		metric.WithUnit("milliseconds"),
-	)
-)
-
 type QueryHook struct {
-	attrs         []attribute.KeyValue
-	formatQueries bool
+	attrs          []attribute.KeyValue
+	formatQueries  bool
+	tracer         trace.Tracer
+	meter          metric.Meter
+	queryHistogram metric.Int64Histogram
 }
 
 var _ bun.QueryHook = (*QueryHook)(nil)
@@ -43,6 +35,17 @@ func NewQueryHook(opts ...Option) *QueryHook {
 	for _, opt := range opts {
 		opt(h)
 	}
+	if h.tracer == nil {
+		h.tracer = otel.Tracer("github.com/uptrace/bun")
+	}
+	if h.meter == nil {
+		h.meter = otel.Meter("github.com/uptrace/bun")
+	}
+	h.queryHistogram, _ = h.meter.Int64Histogram(
+		"go.sql.query_timing",
+		metric.WithDescription("Timing of processed queries"),
+		metric.WithUnit("milliseconds"),
+	)
 	return h
 }
 
@@ -57,7 +60,7 @@ func (h *QueryHook) Init(db *bun.DB) {
 }
 
 func (h *QueryHook) BeforeQuery(ctx context.Context, event *bun.QueryEvent) context.Context {
-	ctx, _ = tracer.Start(ctx, "", trace.WithSpanKind(trace.SpanKindClient))
+	ctx, _ = h.tracer.Start(ctx, "", trace.WithSpanKind(trace.SpanKindClient))
 	return ctx
 }
 
@@ -75,7 +78,7 @@ func (h *QueryHook) AfterQuery(ctx context.Context, event *bun.QueryEvent) {
 	}
 
 	dur := time.Since(event.StartTime)
-	queryHistogram.Record(ctx, dur.Milliseconds(), metric.WithAttributes(labels...))
+	h.queryHistogram.Record(ctx, dur.Milliseconds(), metric.WithAttributes(labels...))
 
 	span := trace.SpanFromContext(ctx)
 	if !span.IsRecording() {