Merge pull request #1656 from authzed/observable-watch-queries-crdb

CRDB: adds missing observability for Watch API

internal/datastore/crdb/pool/health.go

@@ -11,6 +11,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/time/rate"
 
+	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
 	log "github.com/authzed/spicedb/internal/logging"
 )
 
@@ -40,7 +41,7 @@ type NodeHealthTracker struct {
 
 // NewNodeHealthChecker builds a health checker that polls the cluster at the given url.
 func NewNodeHealthChecker(url string) (*NodeHealthTracker, error) {
-	connConfig, err := pgx.ParseConfig(url)
+	connConfig, err := pgxcommon.ParseConfigWithInstrumentation(url)
 	if err != nil {
 		return nil, err
 	}

internal/datastore/crdb/watch.go

@@ -9,7 +9,6 @@ import (
 	"sort"
 	"time"
 
-	"github.com/jackc/pgx/v5"
 	"github.com/prometheus/client_golang/prometheus"
 
 	"github.com/authzed/spicedb/internal/datastore/common"
@@ -67,7 +66,7 @@ func (cds *crdbDatastore) Watch(ctx context.Context, afterRevision datastore.Rev
 	// changefeed data, instead of using a connection pool as most client
 	// drivers do by default."
 	// see: https://www.cockroachlabs.com/docs/v22.2/changefeed-for#considerations
-	conn, err := pgx.Connect(ctx, cds.dburl)
+	conn, err := pgxcommon.ConnectWithInstrumentation(ctx, cds.dburl)
 	if err != nil {
 		errs <- err
 		return updates, errs
@@ -299,7 +298,7 @@ func (cds *crdbDatastore) watchSchemaWithoutRetry(ctx context.Context, afterRevi
 	// changefeed data, instead of using a connection pool as most client
 	// drivers do by default."
 	// see: https://www.cockroachlabs.com/docs/v22.2/changefeed-for#considerations
-	conn, err := pgx.Connect(ctx, cds.dburl)
+	conn, err := pgxcommon.ConnectWithInstrumentation(ctx, cds.dburl)
 	if err != nil {
 		processUpdate(nil, err)
 		return

internal/datastore/postgres/common/pgx.go

@@ -79,6 +79,29 @@ func queryTuples(ctx context.Context, sqlStatement string, args []any, span trac
 	return tuples, nil
 }
 
+// ParseConfigWithInstrumentation returns a pgx.ConnConfig that has been instrumented for observability
+func ParseConfigWithInstrumentation(url string) (*pgx.ConnConfig, error) {
+	connConfig, err := pgx.ParseConfig(url)
+	if err != nil {
+		return nil, err
+	}
+
+	ConfigurePGXLogger(connConfig)
+	ConfigureOTELTracer(connConfig)
+
+	return connConfig, nil
+}
+
+// ConnectWithInstrumentation returns a pgx.Conn that has been instrumented for observability
+func ConnectWithInstrumentation(ctx context.Context, url string) (*pgx.Conn, error) {
+	connConfig, err := ParseConfigWithInstrumentation(url)
+	if err != nil {
+		return nil, err
+	}
+
+	return pgx.ConnectConfig(ctx, connConfig)
+}
+
 // ConfigurePGXLogger sets zerolog global logger into the connection pool configuration, and maps
 // info level events to debug, as they are rather verbose for SpiceDB's info level
 func ConfigurePGXLogger(connConfig *pgx.ConnConfig) {