Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add context to custom trace name function to allow passing context va… #37

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

add context to custom trace name function to allow passing context va… #37

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions options.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package otelpgx

import (
"context"

"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
)
Expand Down Expand Up @@ -43,8 +45,8 @@ 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
// SQL. The function will be called with the current context and the SQL statement as a parameter.
type SpanNameFunc func(ctx context.Context, 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
Expand Down
10 changes: 5 additions & 5 deletions tracer.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,12 +102,12 @@ func recordError(span trace.Span, err error) {

// sqlOperationName attempts to get the first 'word' from a given SQL query, which usually
// is the operation name (e.g. 'SELECT').
func (t *Tracer) sqlOperationName(stmt string) string {
func (t *Tracer) sqlOperationName(ctx context.Context, stmt 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 t.spanNameFunc != nil {
return t.spanNameFunc(stmt)
return t.spanNameFunc(ctx, stmt)
}

parts := strings.Fields(stmt)
Expand Down Expand Up @@ -160,7 +160,7 @@ func (t *Tracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.T

spanName := data.SQL
if t.trimQuerySpanName {
spanName = t.sqlOperationName(data.SQL)
spanName = t.sqlOperationName(ctx, data.SQL)
}
if t.prefixQuerySpanName {
spanName = "query " + spanName
Expand Down Expand Up @@ -271,7 +271,7 @@ func (t *Tracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.T

var spanName string
if t.trimQuerySpanName {
spanName = t.sqlOperationName(data.SQL)
spanName = t.sqlOperationName(ctx, data.SQL)
if t.prefixQuerySpanName {
spanName = "query " + spanName
}
Expand Down Expand Up @@ -353,7 +353,7 @@ func (t *Tracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx

spanName := data.SQL
if t.trimQuerySpanName {
spanName = t.sqlOperationName(data.SQL)
spanName = t.sqlOperationName(ctx, data.SQL)
}
if t.prefixQuerySpanName {
spanName = "prepare " + spanName
Expand Down
25 changes: 23 additions & 2 deletions tracer_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,23 @@
package otelpgx

import (
"context"
"strings"
"testing"
)

type ctxKey string
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: let's use zero sized structs instead, to save allocations

Suggested change
type ctxKey string
type ctxKeySQLName struct{}

(even if this is just a test, it's probably a good idea to keep to good practices)


const sqlName = ctxKey("sql_name")

func TestTracer_sqlOperationName(t *testing.T) {
someCtx := context.WithValue(context.Background(), sqlName, "query name")
tests := []struct {
name string
tracer *Tracer
query string
expName string
ctx *context.Context
}{
{
name: "Spaces only",
Expand Down Expand Up @@ -78,11 +85,22 @@ func TestTracer_sqlOperationName(t *testing.T) {
tracer: NewTracer(WithSpanNameFunc(defaultSpanNameFunc)),
expName: sqlOperationUnknown,
},
{
name: "Custom SQL name query (invalid formatting)",
query: "foo \nSELECT * FROM users",
tracer: NewTracer(WithSpanNameFunc(defaultSpanNameFunc)),
ctx: &someCtx,
expName: "query name",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tr := tt.tracer
if got := tr.sqlOperationName(tt.query); got != tt.expName {
ctx := context.Background()
if tt.ctx != nil {
ctx = *tt.ctx
}
if got := tr.sqlOperationName(ctx, tt.query); got != tt.expName {
t.Errorf("Tracer.sqlOperationName() = %v, want %v", got, tt.expName)
}
})
Expand All @@ -91,7 +109,10 @@ func TestTracer_sqlOperationName(t *testing.T) {

// defaultSpanNameFunc is an utility function for testing that attempts to get
// the first name of the query from a given SQL statement.
var defaultSpanNameFunc SpanNameFunc = func(query string) string {
var defaultSpanNameFunc SpanNameFunc = func(ctx context.Context, query string) string {
if sqlName, ok := ctx.Value(sqlName).(string); ok {
return sqlName
}
for _, line := range strings.Split(query, "\n") {
var prefix string
switch {
Expand Down