feat: allow functionally generated span names

internal/constants.go

@@ -7,5 +7,5 @@ const (
 
 	// InstrumentationVersion is the version of the otelpgx library. This will
 	// be used as an attribute on each span.
-	InstrumentationVersion = "v0.4.1"
+	InstrumentationVersion = "v0.5.0"
 )

options.go

@@ -42,6 +42,21 @@ func WithTrimSQLInSpanName() Option {
 	})
 }
 
+// SpanNameFunc is a function that can be used to generate a span name for a
+// SQL. The function will be called with the SQL statement as a parameter.
+type SpanNameFunc func(stmt string) string
+
+// WithSpanNameFunc will use the provided function to generate the span name for
+// a SQL statement. The function will be called with the SQL statement as a
+// parameter.
+//
+// By default, the whole SQL statement is used as a span name, where applicable.
+func WithSpanNameFunc(fn SpanNameFunc) Option {
+	return optionFunc(func(cfg *tracerConfig) {
+		cfg.spanNameFunc = fn
+	})
+}
+
 // WithDisableSQLStatementInAttributes will disable logging the SQL statement in the span's
 // attributes.
 func WithDisableSQLStatementInAttributes() Option {

tracer.go

@@ -23,6 +23,7 @@ type Tracer struct {
 	tracer            trace.Tracer
 	attrs             []attribute.KeyValue
 	trimQuerySpanName bool
+	spanNameFunc      SpanNameFunc
 	logSQLStatement   bool
 	includeParams     bool
 }
@@ -31,6 +32,7 @@ type tracerConfig struct {
 	tp                trace.TracerProvider
 	attrs             []attribute.KeyValue
 	trimQuerySpanName bool
+	spanNameFunc      SpanNameFunc
 	logSQLStatement   bool
 	includeParams     bool
 }
@@ -43,6 +45,7 @@ func NewTracer(opts ...Option) *Tracer {
 			semconv.DBSystemPostgreSQL,
 		},
 		trimQuerySpanName: false,
+		spanNameFunc:      nil,
 		logSQLStatement:   true,
 		includeParams:     false,
 	}
@@ -55,6 +58,7 @@ func NewTracer(opts ...Option) *Tracer {
 		tracer:            cfg.tp.Tracer(internal.TracerName, trace.WithInstrumentationVersion(internal.InstrumentationVersion)),
 		attrs:             cfg.attrs,
 		trimQuerySpanName: cfg.trimQuerySpanName,
+		spanNameFunc:      cfg.spanNameFunc,
 		logSQLStatement:   cfg.logSQLStatement,
 		includeParams:     cfg.includeParams,
 	}
@@ -71,8 +75,15 @@ const sqlOperationUnknkown = "UNKNOWN"
 
 // sqlOperationName attempts to get the first 'word' from a given SQL query, which usually
 // is the operation name (e.g. 'SELECT').
-func sqlOperationName(query string) string {
-	parts := strings.Fields(query)
+func sqlOperationName(stmt string, fn func(string) string) string {
+	// If a custom function is provided, use that. Otherwise, fall back to the
+	// default implementation. This allows users to override the default
+	// behavior without having to reimplement it.
+	if fn != nil {
+		return fn(stmt)
+	}
+
+	parts := strings.Fields(stmt)
 	if len(parts) == 0 {
 		// Fall back to a fixed value to prevent creating lots of tracing operations
 		// differing only by the amount of whitespace in them (in case we'd fall back
@@ -116,7 +127,7 @@ func (t *Tracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.T
 
 	spanName := "query " + data.SQL
 	if t.trimQuerySpanName {
-		spanName = "query " + sqlOperationName(data.SQL)
+		spanName = "query " + sqlOperationName(data.SQL, t.spanNameFunc)
 	}
 
 	ctx, _ = t.tracer.Start(ctx, spanName, opts...)
@@ -216,7 +227,7 @@ func (t *Tracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.T
 
 	spanName := "batch query " + data.SQL
 	if t.trimQuerySpanName {
-		spanName = "query " + sqlOperationName(data.SQL)
+		spanName = "query " + sqlOperationName(data.SQL, t.spanNameFunc)
 	}
 
 	_, span := t.tracer.Start(ctx, spanName, opts...)
@@ -286,7 +297,7 @@ func (t *Tracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx
 
 	spanName := "prepare " + data.SQL
 	if t.trimQuerySpanName {
-		spanName = "prepare " + sqlOperationName(data.SQL)
+		spanName = "prepare " + sqlOperationName(data.SQL, t.spanNameFunc)
 	}
 
 	ctx, _ = t.tracer.Start(ctx, spanName, opts...)

tracer_test.go

@@ -1,46 +1,138 @@
 package otelpgx
 
-import "testing"
+import (
+	"strings"
+	"testing"
+)
 
 func TestSqlOperationName(t *testing.T) {
 	tests := []struct {
-		name    string
-		query   string
-		expName string
+		name         string
+		query        string
+		spanNameFunc func(string) string
+		expName      string
 	}{
 		{
-			name:    "Spaces only",
-			query:   "SELECT * FROM users",
-			expName: "SELECT",
+			name:         "Spaces only",
+			query:        "SELECT * FROM users",
+			spanNameFunc: nil,
+			expName:      "SELECT",
 		},
 		{
-			name:    "Newline and tab",
-			query:   "UPDATE\n\tfoo",
-			expName: "UPDATE",
+			name:         "Newline and tab",
+			query:        "UPDATE\n\tfoo",
+			spanNameFunc: nil,
+			expName:      "UPDATE",
 		},
 		{
-			name:    "Additional whitespace",
-			query:   " \n SELECT\n\t   *   FROM users  ",
-			expName: "SELECT",
+			name:         "Additional whitespace",
+			query:        " \n SELECT\n\t   *   FROM users  ",
+			spanNameFunc: nil,
+			expName:      "SELECT",
 		},
 		{
-			name:    "Whitespace-only query",
-			query:   " \n\t",
-			expName: "UNKNOWN",
+			name:         "Whitespace-only query",
+			query:        " \n\t",
+			spanNameFunc: nil,
+			expName:      sqlOperationUnknkown,
 		},
 		{
-			name:    "Empty query",
-			query:   "",
-			expName: "UNKNOWN",
+			name:         "Empty query",
+			query:        "",
+			spanNameFunc: nil,
+			expName:      sqlOperationUnknkown,
+		},
+		{
+			name:         "Functional span name (-- comment style)",
+			query:        "-- name: GetUsers :many\nSELECT * FROM users",
+			spanNameFunc: defaultSpanNameFunc(),
+			expName:      "GetUsers :many",
+		},
+		{
+			name:         "Functional span name (/**/ comment style)",
+			query:        "/* name: GetBooks :many */\nSELECT * FROM books",
+			spanNameFunc: defaultSpanNameFunc(),
+			expName:      "GetBooks :many",
+		},
+		{
+			name:         "Functional span name (# comment style)",
+			query:        "# name: GetRecords :many\nSELECT * FROM records",
+			spanNameFunc: defaultSpanNameFunc(),
+			expName:      "GetRecords :many",
+		},
+		{
+			name:         "Functional span name (no annotation)",
+			query:        "--\nSELECT * FROM user",
+			spanNameFunc: defaultSpanNameFunc(),
+			expName:      sqlOperationUnknkown,
+		},
+		{
+			name:         "Custom SQL name query (normal comment)",
+			query:        "-- foo \nSELECT * FROM users",
+			spanNameFunc: defaultSpanNameFunc(),
+			expName:      sqlOperationUnknkown,
+		},
+		{
+			name:         "Custom SQL name query (invalid formatting)",
+			query:        "foo \nSELECT * FROM users",
+			spanNameFunc: defaultSpanNameFunc(),
+			expName:      sqlOperationUnknkown,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			name := sqlOperationName(tt.query)
+			name := sqlOperationName(tt.query, tt.spanNameFunc)
 			if name != tt.expName {
 				t.Errorf("Got name %q, expected %q", name, tt.expName)
 			}
 		})
 	}
 }
+
+// defaultSpanNameFunc is an utility fucntion for testing that attempts to get
+// the first name of the query from a given SQL statement.
+func defaultSpanNameFunc() SpanNameFunc {
+	return func(query string) string {
+		for _, line := range strings.Split(query, "\n") {
+			var prefix string
+			switch {
+			case strings.HasPrefix(line, "--"):
+				prefix = "--"
+
+			case strings.HasPrefix(line, "/*"):
+				prefix = "/*"
+
+			case strings.HasPrefix(line, "#"):
+				prefix = "#"
+			default:
+				continue
+			}
+
+			rest := line[len(prefix):]
+			if !strings.HasPrefix(strings.TrimSpace(rest), "name") {
+				continue
+			}
+			if !strings.Contains(rest, ":") {
+				continue
+			}
+			if !strings.HasPrefix(rest, " name: ") {
+				return sqlOperationUnknkown
+			}
+
+			part := strings.Split(strings.TrimSpace(line), " ")
+			if prefix == "/*" {
+				part = part[:len(part)-1] // removes the trailing "*/" element
+			}
+			if len(part) == 2 {
+				return sqlOperationUnknkown
+			}
+
+			queryName := part[2]
+			queryType := strings.TrimSpace(part[3])
+
+			return queryName + " " + queryType
+		}
+		return sqlOperationUnknkown
+	}
+}