From 0ffcf870d949fa0c7e22efe06d0ac3f189ed5133 Mon Sep 17 00:00:00 2001
From: Nathanael Demacon <nathanael.dmc@outlook.fr>
Date: Wed, 12 Jun 2024 15:04:25 +0200
Subject: [PATCH] feat: add an option to remove the query span name prefix
 (#34)

Signed-off-by: Nathanael DEMACON <quantumsheep@users.noreply.github.com>
Co-authored-by: Nathanael DEMACON <quantumsheep@users.noreply.github.com>
---
 options.go |  8 ++++++++
 tracer.go  | 60 ++++++++++++++++++++++++++++++++----------------------
 2 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/options.go b/options.go
index 53125e6..8780f31 100644
--- a/options.go
+++ b/options.go
@@ -57,6 +57,14 @@ func WithSpanNameFunc(fn SpanNameFunc) Option {
 	})
 }
 
+// WithDisableQuerySpanNamePrefix will disable the default prefix for the span
+// name. By default, the span name is prefixed with "batch query" or "query".
+func WithDisableQuerySpanNamePrefix() Option {
+	return optionFunc(func(cfg *tracerConfig) {
+		cfg.prefixQuerySpanName = false
+	})
+}
+
 // WithDisableSQLStatementInAttributes will disable logging the SQL statement in the span's
 // attributes.
 func WithDisableSQLStatementInAttributes() Option {
diff --git a/tracer.go b/tracer.go
index 255c83e..9f18fc1 100644
--- a/tracer.go
+++ b/tracer.go
@@ -40,21 +40,23 @@ const (
 // Tracer is a wrapper around the pgx tracer interfaces which instrument
 // queries.
 type Tracer struct {
-	tracer            trace.Tracer
-	attrs             []attribute.KeyValue
-	trimQuerySpanName bool
-	spanNameFunc      SpanNameFunc
-	logSQLStatement   bool
-	includeParams     bool
+	tracer              trace.Tracer
+	attrs               []attribute.KeyValue
+	trimQuerySpanName   bool
+	spanNameFunc        SpanNameFunc
+	prefixQuerySpanName bool
+	logSQLStatement     bool
+	includeParams       bool
 }
 
 type tracerConfig struct {
-	tp                trace.TracerProvider
-	attrs             []attribute.KeyValue
-	trimQuerySpanName bool
-	spanNameFunc      SpanNameFunc
-	logSQLStatement   bool
-	includeParams     bool
+	tp                  trace.TracerProvider
+	attrs               []attribute.KeyValue
+	trimQuerySpanName   bool
+	spanNameFunc        SpanNameFunc
+	prefixQuerySpanName bool
+	logSQLStatement     bool
+	includeParams       bool
 }
 
 // NewTracer returns a new Tracer.
@@ -64,10 +66,11 @@ func NewTracer(opts ...Option) *Tracer {
 		attrs: []attribute.KeyValue{
 			semconv.DBSystemPostgreSQL,
 		},
-		trimQuerySpanName: false,
-		spanNameFunc:      nil,
-		logSQLStatement:   true,
-		includeParams:     false,
+		trimQuerySpanName:   false,
+		spanNameFunc:        nil,
+		prefixQuerySpanName: true,
+		logSQLStatement:     true,
+		includeParams:       false,
 	}
 
 	for _, opt := range opts {
@@ -75,12 +78,13 @@ func NewTracer(opts ...Option) *Tracer {
 	}
 
 	return &Tracer{
-		tracer:            cfg.tp.Tracer(tracerName, trace.WithInstrumentationVersion(findOwnImportedVersion())),
-		attrs:             cfg.attrs,
-		trimQuerySpanName: cfg.trimQuerySpanName,
-		spanNameFunc:      cfg.spanNameFunc,
-		logSQLStatement:   cfg.logSQLStatement,
-		includeParams:     cfg.includeParams,
+		tracer:              cfg.tp.Tracer(tracerName, trace.WithInstrumentationVersion(findOwnImportedVersion())),
+		attrs:               cfg.attrs,
+		trimQuerySpanName:   cfg.trimQuerySpanName,
+		spanNameFunc:        cfg.spanNameFunc,
+		prefixQuerySpanName: cfg.prefixQuerySpanName,
+		logSQLStatement:     cfg.logSQLStatement,
+		includeParams:       cfg.includeParams,
 	}
 }
 
@@ -258,9 +262,17 @@ func (t *Tracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.T
 
 	}
 
-	spanName := "batch query " + data.SQL
+	var spanName string
 	if t.trimQuerySpanName {
-		spanName = "query " + t.sqlOperationName(data.SQL)
+		spanName = t.sqlOperationName(data.SQL)
+		if t.prefixQuerySpanName {
+			spanName = "query " + spanName
+		}
+	} else {
+		spanName = data.SQL
+		if t.prefixQuerySpanName {
+			spanName = "batch query " + spanName
+		}
 	}
 
 	_, span := t.tracer.Start(ctx, spanName, opts...)