Adding Context to SpanNameFunc

[pquerna]

We're currently using an internal fork that includes #37

We are using this to propagate down a unique Query ID from our SQL Builder.

Is there an alternative idea for accomplishing this outcome? Context on the span works well for us but feels bad to be using a closed PR.

Example code:

type ctxKeyQueryName struct{}

var queryNameKey = ctxKeyQueryName{}

func WithQueryName(ctx context.Context, queryName string) context.Context {
	return context.WithValue(ctx, queryNameKey, queryName)
}

func QueryName(ctx context.Context) string {
	if name, ok := ctx.Value(queryNameKey).(string); ok {
		return name
	}
	return ""
}

const sqlOperationUnknown = "UNKNOWN"

// firstField is like strings.Fields(), but only gets the first value, eg strings.Fields()[0], but more efficient.
func firstField(input string) string {
	start := -1
	for i, char := range input {
		if unicode.IsSpace(char) {
			if start != -1 {
				// We've found the end of the first word
				return input[start:i]
			}
		} else if start == -1 {
			// We've found the start of the first word
			start = i
		}
	}
	// If we reach here, return the whole word if it exists, or empty string if input was empty/only spaces
	if start != -1 {
		return input[start:]
	}
	return ""
}

func querySpanNameFunc(ctx context.Context, stmt string) string {
	qn := QueryName(ctx)
	if qn != "" {
		return qn
	}

	qn = firstField(stmt)
	if qn != "" {
		return qn
	}
	return sqlOperationUnknown
}


func() biggerFunc() {
		tr := otelpgx.NewTracer(
			otelpgx.WithSpanNameFunc(querySpanNameFunc),
		)
}

[rcrowe]

@pquerna As an alternative, in case it helps, we use SQL comments (via sqlc in our case) to extract the query name. Could your query builder be changed to allow prefixing a comment?

For example...

-- name: DeleteUser :exec
UPDATE staff
SET deleted_at = $2
WHERE id=$1;

otelpgx.NewTracer(
  otelpgx.WithSpanNameFunc(func(stmt string) string {
      return queryName(stmt)
  }),
)

// queryName looks to see if the query has a name and returns it.
// It does this by looking for a prefixed comment on the SQL query,
// which means it supports SQLC out of the box.
//
// Formats supported listed in order of performance parsing:
// -- name: {name} :{type}
// -- name:{name}
// -- name: {name}
//
// PRs adding more formats are welcome, but need to show limited performance impact.
func queryName(q string) string {
	// regex is slow so we're doing it manually for a bunch of cases to maintain performance.

	if len(q) >= 10 && q[0:9] == "-- name: " {
		// -- name: {name} :{type}
		if index := strings.Index(q[9:], ":"); index != -1 {
			// remove leading space before :
			// optimised assuming this comment has been generated by SQLC
			return q[9 : 9+(index-1)]
		}
		// -- name: {name}
		if index := strings.Index(q[9:], "\n"); index != -1 {
			if q[len(q[9:9+index])-1] != ' ' {
				return q[9 : 9+index]
			}
			return strings.TrimSpace(q[9 : 9+index])
		}
	}

	if len(q) >= 9 && q[0:8] == "-- name:" {
		// -- name:{name}
		if index := strings.Index(q[8:], "\n"); index != -1 {
			if q[len(q[8:8+index])-1] != ' ' {
				return q[8 : 8+index]
			}
			return strings.TrimSpace(q[8 : 8+index])
		}
	}

	if len(q) >= 8 && q[0:7] == "--name:" {
		// -- name:{name}
		if index := strings.Index(q[7:], "\n"); index != -1 {
			if q[len(q[7:7+index])-1] != ' ' {
				return q[7 : 7+index]
			}
			return strings.TrimSpace(q[7 : 7+index])
		}
	}

	return ""
}

edit:
Had a quick go using a query builder to add the comment...

func main() {
	sb := sqlbuilder.NewUpdateBuilder()
	sb.Update("staff").
		Set(sb.Assign("deleted_at", "<timestamp>")).
		Where(sb.Assign("id", "<id>"))

	query := WithQueryName("DeleteUser", sb)
}

func WithQueryName(queryName string, builder sqlbuilder.Builder) sqlbuilder.Builder {
	return sqlbuilder.Buildf("-- name: %s\n%v", queryName, builder)
}

[robsonpeixoto]

An approach of allow access the context.Context without break API retro-compatibility: #67