add OTEL tracing to pgx Queries

go.mod

@@ -20,6 +20,7 @@ require (
 	github.com/ecordell/optgen v0.0.10-0.20230609182709-018141bf9698
 	github.com/emirpasic/gods v1.18.1
 	github.com/envoyproxy/protoc-gen-validate v1.0.2
+	github.com/exaring/otelpgx v0.5.0
 	github.com/fatih/color v1.15.0
 	github.com/go-co-op/gocron v1.30.1
 	github.com/go-logr/zerologr v1.2.3

go.sum

@@ -582,6 +582,8 @@ github.com/ettle/strcase v0.1.1 h1:htFueZyVeE1XNnMEfbqp5r67qAN/4r6ya1ysq8Q+Zcw=
 github.com/ettle/strcase v0.1.1/go.mod h1:hzDLsPC7/lwKyBOywSHEP89nt2pDgdy+No1NBA9o9VY=
 github.com/evanphx/json-patch v4.12.0+incompatible h1:4onqiflcdA9EOZ4RxV643DvftH5pOlLGNtQ5lPWQu84=
 github.com/evanphx/json-patch/v5 v5.6.0 h1:b91NhWfaz02IuVxO9faSllyAtNXHMPkC5J8sJCLunww=
+github.com/exaring/otelpgx v0.5.0 h1:XA8F1rWSyOleizQ8PlzZ3nYggIg4fwaVzFAJii0d5RE=
+github.com/exaring/otelpgx v0.5.0/go.mod h1:a3fAXoYxYGCID+2WHKNe9eoDIiWl01SCc26KtdTMVYk=
 github.com/fatih/color v1.15.0 h1:kOqh6YHBtK8aywxGerMG2Eq3H6Qgoqeo13Bk2Mv/nBs=
 github.com/fatih/color v1.15.0/go.mod h1:0h5ZqXfHYED7Bhv2ZJamyIOUej9KtShiJESRwBDUSsw=
 github.com/fatih/set v0.2.1 h1:nn2CaJyknWE/6txyUDGwysr3G5QC6xWB/PtVjPBbeaA=

internal/datastore/crdb/migrations/driver.go

@@ -35,6 +35,7 @@ func NewCRDBDriver(url string) (*CRDBDriver, error) {
 		return nil, fmt.Errorf(errUnableToInstantiate, err)
 	}
 	pgxcommon.ConfigurePGXLogger(connConfig)
+	pgxcommon.ConfigureOTELTracer(connConfig)
 
 	db, err := pgx.ConnectConfig(context.Background(), connConfig)
 	if err != nil {

internal/datastore/postgres/common/pgx.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/exaring/otelpgx"
 	zerologadapter "github.com/jackc/pgx-zerolog"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
@@ -105,7 +106,52 @@ func ConfigurePGXLogger(connConfig *pgx.ConnConfig) {
 	}
 
 	l := zerologadapter.NewLogger(log.Logger, zerologadapter.WithSubDictionary("pgx"))
-	connConfig.Tracer = &tracelog.TraceLog{Logger: levelMappingFn(l), LogLevel: tracelog.LogLevelInfo}
+	addTracer(connConfig, &tracelog.TraceLog{Logger: levelMappingFn(l), LogLevel: tracelog.LogLevelInfo})
+}
+
+// ConfigureOTELTracer adds OTEL tracing to a pgx.ConnConfig
+func ConfigureOTELTracer(connConfig *pgx.ConnConfig) {
+	addTracer(connConfig, otelpgx.NewTracer(otelpgx.WithTrimSQLInSpanName()))
+}
+
+func addTracer(connConfig *pgx.ConnConfig, tracer pgx.QueryTracer) {
+	composedTracer := addComposedTracer(connConfig)
+	composedTracer.Tracers = append(composedTracer.Tracers, tracer)
+}
+
+func addComposedTracer(connConfig *pgx.ConnConfig) *ComposedTracer {
+	var composedTracer *ComposedTracer
+	if connConfig.Tracer == nil {
+		composedTracer = &ComposedTracer{}
+		connConfig.Tracer = composedTracer
+	} else {
+		var ok bool
+		composedTracer, ok = connConfig.Tracer.(*ComposedTracer)
+		if !ok {
+			composedTracer.Tracers = append(composedTracer.Tracers, connConfig.Tracer)
+			connConfig.Tracer = composedTracer
+		}
+	}
+	return composedTracer
+}
+
+// ComposedTracer allows adding multiple tracers to a pgx.ConnConfig
+type ComposedTracer struct {
+	Tracers []pgx.QueryTracer
+}
+
+func (m *ComposedTracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+	for _, t := range m.Tracers {
+		ctx = t.TraceQueryStart(ctx, conn, data)
+	}
+
+	return ctx
+}
+
+func (m *ComposedTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
+	for _, t := range m.Tracers {
+		t.TraceQueryEnd(ctx, conn, data)
+	}
 }
 
 // DBFuncQuerier is satisfied by RetryPool and QuerierFuncs (which can wrap a pgxpool or transaction)
@@ -159,6 +205,7 @@ func (opts PoolOptions) ConfigurePgx(pgxConfig *pgxpool.Config) {
 	}
 
 	ConfigurePGXLogger(pgxConfig.ConnConfig)
+	ConfigureOTELTracer(pgxConfig.ConnConfig)
 }
 
 type QuerierFuncs struct {

internal/datastore/postgres/postgres_test.go

@@ -16,6 +16,9 @@ import (
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	"golang.org/x/sync/errgroup"
 
 	"github.com/authzed/spicedb/internal/datastore/common"
@@ -150,6 +153,15 @@ func TestPostgresDatastore(t *testing.T) {
 					MigrationPhase(config.migrationPhase),
 				))
 			}
+
+			t.Run("OTelTracing", createDatastoreTest(
+				b,
+				OTelTracingTest,
+				RevisionQuantization(0),
+				GCWindow(1*time.Millisecond),
+				WatchBufferLength(1),
+				MigrationPhase(config.migrationPhase),
+			))
 		})
 	}
 }
@@ -1054,6 +1066,43 @@ func RevisionInversionTest(t *testing.T, ds datastore.Datastore) {
 	require.False(commitFirstRev.Equal(commitLastRev))
 }
 
+func OTelTracingTest(t *testing.T, ds datastore.Datastore) {
+	require := require.New(t)
+
+	ctx := context.Background()
+	r, err := ds.ReadyState(ctx)
+	require.NoError(err)
+	require.True(r.IsReady)
+
+	// Setup OTel recording
+	defaultProvider := otel.GetTracerProvider()
+
+	provider := trace.NewTracerProvider(
+		trace.WithSampler(trace.AlwaysSample()),
+	)
+	spanrecorder := tracetest.NewSpanRecorder()
+	provider.RegisterSpanProcessor(spanrecorder)
+	otel.SetTracerProvider(provider)
+	defer func() {
+		otel.SetTracerProvider(defaultProvider)
+	}()
+
+	// Perform basic operation
+	_, err = ds.ReadWriteTx(ctx, func(rwt datastore.ReadWriteTransaction) error {
+		return rwt.WriteNamespaces(ctx, namespace.Namespace("resource"))
+	})
+	require.NoError(err)
+
+	ended := spanrecorder.Ended()
+	var present bool
+	for _, span := range ended {
+		if span.Name() == "query INSERT" {
+			present = true
+		}
+	}
+	require.True(present, "missing trace for Streaming gRPC call")
+}
+
 func WatchNotEnabledTest(t *testing.T, _ testdatastore.RunningEngineForTest) {
 	require := require.New(t)