mirror of https://github.com/coder/coder.git
46 lines
1.3 KiB
Go
46 lines
1.3 KiB
Go
package tracing
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"go.nhat.io/otelsql"
|
|
semconv "go.opentelemetry.io/otel/semconv/v1.14.0"
|
|
"go.opentelemetry.io/otel/trace"
|
|
"golang.org/x/xerrors"
|
|
)
|
|
|
|
// Postgres driver will register a new tracing sql driver and return the driver name.
|
|
func PostgresDriver(tp trace.TracerProvider, service string) (string, error) {
|
|
// Register the otelsql wrapper for the provided postgres driver.
|
|
driverName, err := otelsql.Register("postgres",
|
|
otelsql.WithDefaultAttributes(
|
|
semconv.ServiceNameKey.String(service),
|
|
),
|
|
otelsql.TraceQueryWithoutArgs(),
|
|
otelsql.WithSystem(semconv.DBSystemPostgreSQL),
|
|
otelsql.WithTracerProvider(tp),
|
|
otelsql.WithSpanNameFormatter(formatPostgresSpan),
|
|
)
|
|
if err != nil {
|
|
return "", xerrors.Errorf("registering postgres tracing driver: %w", err)
|
|
}
|
|
|
|
return driverName, nil
|
|
}
|
|
|
|
func formatPostgresSpan(ctx context.Context, op string) string {
|
|
const qPrefix = "-- name: "
|
|
q := otelsql.QueryFromContext(ctx)
|
|
if q == "" || !strings.HasPrefix(q, qPrefix) {
|
|
return strings.ToUpper(op)
|
|
}
|
|
|
|
// Remove the qPrefix and then grab the method name.
|
|
// We expect the first line of the query to be in
|
|
// the format "-- name: GetAPIKeyByID :one".
|
|
s := strings.SplitN(strings.TrimPrefix(q, qPrefix), " ", 2)[0]
|
|
return fmt.Sprintf("%s %s", strings.ToUpper(op), s)
|
|
}
|