Change CRDB driver to use new method for getting transaction timestamp

The existing call disables most optimizations on write transactions, but was necessary for legacy reasons. Following this change, any CRDB version 23 (or later) will use the newer, better call

Also changes the stats system to use the CRDB stats instead of a custom table

Author: josephschorr

internal/datastore/crdb/crdb.go

@@ -68,8 +68,10 @@ const (
 	errUnableToInstantiate = "unable to instantiate datastore"
 	errRevision            = "unable to find revision: %w"
 
-	querySelectNow      = "SELECT cluster_logical_timestamp()"
-	queryShowZoneConfig = "SHOW ZONE CONFIGURATION FOR RANGE default;"
+	querySelectNow            = "SELECT cluster_logical_timestamp()"
+	queryTransactionNowPreV23 = querySelectNow
+	queryTransactionNow       = "SHOW COMMIT TIMESTAMP"
+	queryShowZoneConfig       = "SHOW ZONE CONFIGURATION FOR RANGE default;"
 
 	livingTupleConstraint = "pk_relation_tuple"
 )
@@ -122,6 +124,12 @@ func newCRDBDatastore(ctx context.Context, url string, options ...Option) (datas
 		changefeedQuery = queryChangefeedPreV22
 	}
 
+	transactionNowQuery := queryTransactionNow
+	if version.Major < 23 {
+		log.Info().Object("version", version).Msg("using transaction now query for CRDB version < 23")
+		transactionNowQuery = queryTransactionNowPreV23
+	}
+
 	clusterTTLNanos, err := readClusterTTLNanos(initCtx, initPool)
 	if err != nil {
 		return nil, fmt.Errorf("unable to read cluster gc window: %w", err)
@@ -172,8 +180,9 @@ func newCRDBDatastore(ctx context.Context, url string, options ...Option) (datas
 		watchBufferWriteTimeout: config.watchBufferWriteTimeout,
 		writeOverlapKeyer:       keyer,
 		overlapKeyInit:          keySetInit,
-		disableStats:            config.disableStats,
 		beginChangefeedQuery:    changefeedQuery,
+		transactionNowQuery:     transactionNowQuery,
+		analyzeBeforeStatistics: config.analyzeBeforeStatistics,
 	}
 	ds.RemoteClockRevisions.SetNowFunc(ds.headRevisionInternal)
 
@@ -254,9 +263,10 @@ type crdbDatastore struct {
 	watchBufferWriteTimeout time.Duration
 	writeOverlapKeyer       overlapKeyer
 	overlapKeyInit          func(ctx context.Context) keySet
-	disableStats            bool
+	analyzeBeforeStatistics bool
 
 	beginChangefeedQuery string
+	transactionNowQuery  string
 
 	featureGroup singleflight.Group[string, *datastore.Features]
 
@@ -321,21 +331,11 @@ func (cds *crdbDatastore) ReadWriteTx(
 			}
 		}
 
-		if cds.disableStats {
-			var err error
-			commitTimestamp, err = readCRDBNow(ctx, querier)
-			if err != nil {
-				return fmt.Errorf("error getting commit timestamp: %w", err)
-			}
-			return nil
-		}
-
 		var err error
-		commitTimestamp, err = updateCounter(ctx, tx, rwt.relCountChange)
+		commitTimestamp, err = cds.readTransactionCommitRev(ctx, querier)
 		if err != nil {
-			return fmt.Errorf("error updating relationship counter: %w", err)
+			return fmt.Errorf("error getting commit timestamp: %w", err)
 		}
-
 		return nil
 	})
 	if err != nil {
@@ -371,7 +371,9 @@ func (cds *crdbDatastore) ReadyState(ctx context.Context) (datastore.ReadyState,
 		return datastore.ReadyState{}, err
 	}
 
-	if version != headMigration {
+	// TODO(jschorr): Remove the check for the older migration once we are confident
+	// that all users have migrated past it.
+	if version != headMigration && version != "add-caveats" {
 		return datastore.ReadyState{
 			Message: fmt.Sprintf(
 				"datastore is not migrated: currently at revision `%s`, but requires `%s`. Please run `spicedb migrate`.",
@@ -467,6 +469,20 @@ func (cds *crdbDatastore) features(ctx context.Context) (*datastore.Features, er
 	return &features, nil
 }
 
+func (cds *crdbDatastore) readTransactionCommitRev(ctx context.Context, reader pgxcommon.DBFuncQuerier) (datastore.Revision, error) {
+	ctx, span := tracer.Start(ctx, "readTransactionCommitRev")
+	defer span.End()
+
+	var hlcNow decimal.Decimal
+	if err := reader.QueryRowFunc(ctx, func(ctx context.Context, row pgx.Row) error {
+		return row.Scan(&hlcNow)
+	}, cds.transactionNowQuery); err != nil {
+		return datastore.NoRevision, fmt.Errorf("unable to read timestamp: %w", err)
+	}
+
+	return revisions.NewForHLC(hlcNow)
+}
+
 func readCRDBNow(ctx context.Context, reader pgxcommon.DBFuncQuerier) (datastore.Revision, error) {
 	ctx, span := tracer.Start(ctx, "readCRDBNow")
 	defer span.End()

internal/datastore/crdb/crdb_test.go

@@ -53,6 +53,7 @@ func TestCRDBDatastore(t *testing.T) {
 				RevisionQuantization(revisionQuantization),
 				WatchBufferLength(watchBufferLength),
 				OverlapStrategy(overlapStrategyPrefix),
+				DebugAnalyzeBeforeStatistics(),
 			)
 			require.NoError(t, err)
 			return ds
@@ -84,6 +85,7 @@ func TestCRDBDatastoreWithFollowerReads(t *testing.T) {
 					GCWindow(gcWindow),
 					RevisionQuantization(quantization),
 					FollowerReadDelay(followerReadDelay),
+					DebugAnalyzeBeforeStatistics(),
 				)
 				require.NoError(err)
 				return ds
@@ -134,15 +136,19 @@ func TestWatchFeatureDetection(t *testing.T) {
 				require.NoError(t, err)
 			},
 			expectEnabled: false,
-			expectMessage: "Range feeds must be enabled in CockroachDB and the user must have permission to create them in order to enable the Watch API: ERROR: user unprivileged does not have CHANGEFEED privilege on relation relation_tuple (SQLSTATE 42501)",
+			expectMessage: "(SQLSTATE 42501)",
 		},
 		{
 			name: "rangefeeds enabled, user has permission",
 			postInit: func(ctx context.Context, adminConn *pgx.Conn) {
 				_, err = adminConn.Exec(ctx, `SET CLUSTER SETTING kv.rangefeed.enabled = true;`)
 				require.NoError(t, err)
+
 				_, err = adminConn.Exec(ctx, fmt.Sprintf(`GRANT CHANGEFEED ON TABLE testspicedb.%s TO unprivileged;`, tableTuple))
 				require.NoError(t, err)
+
+				_, err = adminConn.Exec(ctx, fmt.Sprintf(`GRANT SELECT ON TABLE testspicedb.%s TO unprivileged;`, tableTuple))
+				require.NoError(t, err)
 			},
 			expectEnabled: true,
 		},

internal/datastore/crdb/migrations/zz_migration.0005_remove_stats_table.go

@@ -0,0 +1,25 @@
+package migrations
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+)
+
+const (
+	dropStatsTable = `DROP TABLE relationship_estimate_counters;`
+)
+
+func init() {
+	err := CRDBMigrations.Register("remove-stats-table", "add-caveats", removeStatsTable, noAtomicMigration)
+	if err != nil {
+		panic("failed to register migration: " + err.Error())
+	}
+}
+
+func removeStatsTable(ctx context.Context, conn *pgx.Conn) error {
+	if _, err := conn.Exec(ctx, dropStatsTable); err != nil {
+		return err
+	}
+	return nil
+}

internal/datastore/crdb/options.go

@@ -20,8 +20,8 @@ type crdbOptions struct {
 	maxRetries                  uint8
 	overlapStrategy             string
 	overlapKey                  string
-	disableStats                bool
 	enableConnectionBalancing   bool
+	analyzeBeforeStatistics     bool
 
 	enablePrometheusStats bool
 }
@@ -65,7 +65,6 @@ func generateConfig(options []Option) (crdbOptions, error) {
 		maxRetries:                  defaultMaxRetries,
 		overlapKey:                  defaultOverlapKey,
 		overlapStrategy:             defaultOverlapStrategy,
-		disableStats:                false,
 		enablePrometheusStats:       defaultEnablePrometheusStats,
 		enableConnectionBalancing:   defaultEnableConnectionBalancing,
 		connectRate:                 defaultConnectRate,
@@ -283,11 +282,6 @@ func OverlapKey(key string) Option {
 	return func(po *crdbOptions) { po.overlapKey = key }
 }
 
-// DisableStats disables recording counts to the stats table
-func DisableStats(disable bool) Option {
-	return func(po *crdbOptions) { po.disableStats = disable }
-}
-
 // WithEnablePrometheusStats marks whether Prometheus metrics provided by the Postgres
 // clients being used by the datastore are enabled.
 //
@@ -303,3 +297,12 @@ func WithEnablePrometheusStats(enablePrometheusStats bool) Option {
 func WithEnableConnectionBalancing(connectionBalancing bool) Option {
 	return func(po *crdbOptions) { po.enableConnectionBalancing = connectionBalancing }
 }
+
+// DebugAnalyzeBeforeStatistics signals to the Statistics method that it should
+// run Analyze on the database before returning statistics. This should only be
+// used for debug and testing.
+//
+// Disabled by default.
+func DebugAnalyzeBeforeStatistics() Option {
+	return func(po *crdbOptions) { po.analyzeBeforeStatistics = true }
+}

internal/datastore/crdb/stats.go

@@ -3,39 +3,24 @@ package crdb
 import (
 	"context"
 	"fmt"
-	"math/rand"
-	"time"
+	"slices"
 
 	"github.com/Masterminds/squirrel"
 	"github.com/jackc/pgx/v5"
-	"github.com/shopspring/decimal"
+	"github.com/rs/zerolog/log"
 
 	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
-	"github.com/authzed/spicedb/internal/datastore/revisions"
 	"github.com/authzed/spicedb/pkg/datastore"
 )
 
 const (
 	tableMetadata = "metadata"
 	colUniqueID   = "unique_id"
-
-	tableCounters = "relationship_estimate_counters"
-	colID         = "id"
-	colCount      = "count"
 )
 
 var (
-	queryReadUniqueID         = psql.Select(colUniqueID).From(tableMetadata)
-	queryRelationshipEstimate = fmt.Sprintf("SELECT COALESCE(SUM(%s), 0) FROM %s AS OF SYSTEM TIME follower_read_timestamp()", colCount, tableCounters)
-
-	upsertCounterQuery = psql.Insert(tableCounters).Columns(
-		colID,
-		colCount,
-	).Suffix(fmt.Sprintf("ON CONFLICT (%[1]s) DO UPDATE SET %[2]s = %[3]s.%[2]s + EXCLUDED.%[2]s RETURNING cluster_logical_timestamp()", colID, colCount, tableCounters))
-
-	rng = rand.NewSource(time.Now().UnixNano())
-
-	uniqueID string
+	queryReadUniqueID = psql.Select(colUniqueID).From(tableMetadata)
+	uniqueID          string
 )
 
 func (cds *crdbDatastore) Statistics(ctx context.Context) (datastore.Stats, error) {
@@ -52,14 +37,6 @@ func (cds *crdbDatastore) Statistics(ctx context.Context) (datastore.Stats, erro
 	}
 
 	var nsDefs []datastore.RevisionedNamespace
-	var relCount int64
-
-	if err := cds.readPool.QueryRowFunc(ctx, func(ctx context.Context, row pgx.Row) error {
-		return row.Scan(&relCount)
-	}, queryRelationshipEstimate); err != nil {
-		return datastore.Stats{}, fmt.Errorf("unable to read relationship count: %w", err)
-	}
-
 	if err := cds.readPool.BeginTxFunc(ctx, pgx.TxOptions{AccessMode: pgx.ReadOnly}, func(tx pgx.Tx) error {
 		_, err := tx.Exec(ctx, "SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()")
 		if err != nil {
@@ -76,37 +53,85 @@ func (cds *crdbDatastore) Statistics(ctx context.Context) (datastore.Stats, erro
 		return datastore.Stats{}, err
 	}
 
-	// NOTE: this is a stop-gap solution to prevent panics in telemetry collection
-	if relCount < 0 {
-		relCount = 0
-	}
-
-	return datastore.Stats{
-		UniqueID:                   uniqueID,
-		EstimatedRelationshipCount: uint64(relCount),
-		ObjectTypeStatistics:       datastore.ComputeObjectTypeStats(nsDefs),
-	}, nil
-}
+	if cds.analyzeBeforeStatistics {
+		if err := cds.readPool.BeginTxFunc(ctx, pgx.TxOptions{AccessMode: pgx.ReadOnly}, func(tx pgx.Tx) error {
+			if _, err := tx.Exec(ctx, "ANALYZE "+tableTuple); err != nil {
+				return fmt.Errorf("unable to analyze tuple table: %w", err)
+			}
 
-func updateCounter(ctx context.Context, tx pgx.Tx, change int64) (datastore.Revision, error) {
-	counterID := make([]byte, 2)
-	// nolint:gosec
-	// G404 use of non cryptographically secure random number generator is not concern here,
-	// as this is only used to randomly distributed the counters across multiple rows and reduce write row contention
-	_, err := rand.New(rng).Read(counterID)
-	if err != nil {
-		return datastore.NoRevision, fmt.Errorf("unable to select random counter: %w", err)
+			return nil
+		}); err != nil {
+			return datastore.Stats{}, err
+		}
 	}
 
-	sql, args, err := upsertCounterQuery.Values(counterID, change).ToSql()
-	if err != nil {
-		return datastore.NoRevision, fmt.Errorf("unable to prepare upsert counter sql: %w", err)
-	}
+	var estimatedRelCount uint64
+	if err := cds.readPool.QueryFunc(ctx, func(ctx context.Context, rows pgx.Rows) error {
+		hasRows := false
+
+		for rows.Next() {
+			hasRows = true
+			values, err := rows.Values()
+			if err != nil {
+				log.Warn().Err(err).Msg("unable to read statistics")
+				return nil
+			}
+
+			// Find the row whose column_names contains the expected columns for the
+			// full relationship.
+			isFullRelationshipRow := false
+			for index, fd := range rows.FieldDescriptions() {
+				if fd.Name != "column_names" {
+					continue
+				}
+
+				columnNames, ok := values[index].([]any)
+				if !ok {
+					log.Warn().Msg("unable to read column names")
+					return nil
+				}
+
+				if slices.Contains(columnNames, "namespace") &&
+					slices.Contains(columnNames, "object_id") &&
+					slices.Contains(columnNames, "relation") &&
+					slices.Contains(columnNames, "userset_namespace") &&
+					slices.Contains(columnNames, "userset_object_id") &&
+					slices.Contains(columnNames, "userset_relation") {
+					isFullRelationshipRow = true
+					break
+				}
+			}
+
+			if !isFullRelationshipRow {
+				continue
+			}
+
+			// Read the estimated relationship count.
+			for index, fd := range rows.FieldDescriptions() {
+				if fd.Name != "row_count" {
+					continue
+				}
+
+				rowCount, ok := values[index].(int64)
+				if !ok {
+					log.Warn().Msg("unable to read row count")
+					return nil
+				}
+
+				estimatedRelCount = uint64(rowCount)
+				return nil
+			}
+		}
 
-	var timestamp decimal.Decimal
-	if err := tx.QueryRow(ctx, sql, args...).Scan(&timestamp); err != nil {
-		return datastore.NoRevision, fmt.Errorf("unable to executed upsert counter query: %w", err)
+		log.Warn().Bool("has-rows", hasRows).Msg("unable to find row count in statistics query result")
+		return nil
+	}, "SHOW STATISTICS FOR TABLE relation_tuple;"); err != nil {
+		return datastore.Stats{}, fmt.Errorf("unable to query unique estimated row count: %w", err)
 	}
 
-	return revisions.NewForHLC(timestamp)
+	return datastore.Stats{
+		UniqueID:                   uniqueID,
+		EstimatedRelationshipCount: estimatedRelCount,
+		ObjectTypeStatistics:       datastore.ComputeObjectTypeStats(nsDefs),
+	}, nil
 }

internal/testserver/datastore/crdb.go

@@ -20,7 +20,7 @@ import (
 )
 
 const (
-	CRDBTestVersionTag = "v22.2.0"
+	CRDBTestVersionTag = "v23.1.16"
 
 	enableRangefeeds = `SET CLUSTER SETTING kv.rangefeed.enabled = true;`
 )

pkg/cmd/datastore/datastore.go

@@ -382,7 +382,6 @@ func newCRDBDatastore(ctx context.Context, opts Config) (datastore.Datastore, er
 		crdb.OverlapStrategy(opts.OverlapStrategy),
 		crdb.WatchBufferLength(opts.WatchBufferLength),
 		crdb.WatchBufferWriteTimeout(opts.WatchBufferWriteTimeout),
-		crdb.DisableStats(opts.DisableStats),
 		crdb.WithEnablePrometheusStats(opts.EnableDatastoreMetrics),
 		crdb.WithEnableConnectionBalancing(opts.EnableConnectionBalancing),
 		crdb.ConnectRate(opts.ConnectRate),