feat: add tracing for sql (#1610)

coderd/tracing/postgres.go

@@ -0,0 +1,45 @@
+package tracing
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/nhatthm/otelsql"
+	semconv "go.opentelemetry.io/otel/semconv/v1.7.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)
+}