@@ -3,39 +3,24 @@ package crdb
33import (
44 "context"
55 "fmt"
6- "math/rand"
7- "time"
6+ "slices"
87
98 "github.com/Masterminds/squirrel"
109 "github.com/jackc/pgx/v5"
11- "github.com/shopspring/decimal "
10+ "github.com/rs/zerolog/log "
1211
1312 pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
14- "github.com/authzed/spicedb/internal/datastore/revisions"
1513 "github.com/authzed/spicedb/pkg/datastore"
1614)
1715
1816const (
1917 tableMetadata = "metadata"
2018 colUniqueID = "unique_id"
21-
22- tableCounters = "relationship_estimate_counters"
23- colID = "id"
24- colCount = "count"
2519)
2620
2721var (
28- queryReadUniqueID = psql .Select (colUniqueID ).From (tableMetadata )
29- queryRelationshipEstimate = fmt .Sprintf ("SELECT COALESCE(SUM(%s), 0) FROM %s AS OF SYSTEM TIME follower_read_timestamp()" , colCount , tableCounters )
30-
31- upsertCounterQuery = psql .Insert (tableCounters ).Columns (
32- colID ,
33- colCount ,
34- ).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 ))
35-
36- rng = rand .NewSource (time .Now ().UnixNano ())
37-
38- uniqueID string
22+ queryReadUniqueID = psql .Select (colUniqueID ).From (tableMetadata )
23+ uniqueID string
3924)
4025
4126func (cds * crdbDatastore ) Statistics (ctx context.Context ) (datastore.Stats , error ) {
@@ -52,14 +37,6 @@ func (cds *crdbDatastore) Statistics(ctx context.Context) (datastore.Stats, erro
5237 }
5338
5439 var nsDefs []datastore.RevisionedNamespace
55- var relCount int64
56-
57- if err := cds .readPool .QueryRowFunc (ctx , func (ctx context.Context , row pgx.Row ) error {
58- return row .Scan (& relCount )
59- }, queryRelationshipEstimate ); err != nil {
60- return datastore.Stats {}, fmt .Errorf ("unable to read relationship count: %w" , err )
61- }
62-
6340 if err := cds .readPool .BeginTxFunc (ctx , pgx.TxOptions {AccessMode : pgx .ReadOnly }, func (tx pgx.Tx ) error {
6441 _ , err := tx .Exec (ctx , "SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()" )
6542 if err != nil {
@@ -76,37 +53,85 @@ func (cds *crdbDatastore) Statistics(ctx context.Context) (datastore.Stats, erro
7653 return datastore.Stats {}, err
7754 }
7855
79- // NOTE: this is a stop-gap solution to prevent panics in telemetry collection
80- if relCount < 0 {
81- relCount = 0
82- }
83-
84- return datastore.Stats {
85- UniqueID : uniqueID ,
86- EstimatedRelationshipCount : uint64 (relCount ),
87- ObjectTypeStatistics : datastore .ComputeObjectTypeStats (nsDefs ),
88- }, nil
89- }
56+ if cds .analyzeBeforeStatistics {
57+ if err := cds .readPool .BeginTxFunc (ctx , pgx.TxOptions {AccessMode : pgx .ReadOnly }, func (tx pgx.Tx ) error {
58+ if _ , err := tx .Exec (ctx , "ANALYZE " + tableTuple ); err != nil {
59+ return fmt .Errorf ("unable to analyze tuple table: %w" , err )
60+ }
9061
91- func updateCounter (ctx context.Context , tx pgx.Tx , change int64 ) (datastore.Revision , error ) {
92- counterID := make ([]byte , 2 )
93- // nolint:gosec
94- // G404 use of non cryptographically secure random number generator is not concern here,
95- // as this is only used to randomly distributed the counters across multiple rows and reduce write row contention
96- _ , err := rand .New (rng ).Read (counterID )
97- if err != nil {
98- return datastore .NoRevision , fmt .Errorf ("unable to select random counter: %w" , err )
62+ return nil
63+ }); err != nil {
64+ return datastore.Stats {}, err
65+ }
9966 }
10067
101- sql , args , err := upsertCounterQuery .Values (counterID , change ).ToSql ()
102- if err != nil {
103- return datastore .NoRevision , fmt .Errorf ("unable to prepare upsert counter sql: %w" , err )
104- }
68+ var estimatedRelCount uint64
69+ if err := cds .readPool .QueryFunc (ctx , func (ctx context.Context , rows pgx.Rows ) error {
70+ hasRows := false
71+
72+ for rows .Next () {
73+ hasRows = true
74+ values , err := rows .Values ()
75+ if err != nil {
76+ log .Warn ().Err (err ).Msg ("unable to read statistics" )
77+ return nil
78+ }
79+
80+ // Find the row whose column_names contains the expected columns for the
81+ // full relationship.
82+ isFullRelationshipRow := false
83+ for index , fd := range rows .FieldDescriptions () {
84+ if fd .Name != "column_names" {
85+ continue
86+ }
87+
88+ columnNames , ok := values [index ].([]any )
89+ if ! ok {
90+ log .Warn ().Msg ("unable to read column names" )
91+ return nil
92+ }
93+
94+ if slices .Contains (columnNames , "namespace" ) &&
95+ slices .Contains (columnNames , "object_id" ) &&
96+ slices .Contains (columnNames , "relation" ) &&
97+ slices .Contains (columnNames , "userset_namespace" ) &&
98+ slices .Contains (columnNames , "userset_object_id" ) &&
99+ slices .Contains (columnNames , "userset_relation" ) {
100+ isFullRelationshipRow = true
101+ break
102+ }
103+ }
104+
105+ if ! isFullRelationshipRow {
106+ continue
107+ }
108+
109+ // Read the estimated relationship count.
110+ for index , fd := range rows .FieldDescriptions () {
111+ if fd .Name != "row_count" {
112+ continue
113+ }
114+
115+ rowCount , ok := values [index ].(int64 )
116+ if ! ok {
117+ log .Warn ().Msg ("unable to read row count" )
118+ return nil
119+ }
120+
121+ estimatedRelCount = uint64 (rowCount )
122+ return nil
123+ }
124+ }
105125
106- var timestamp decimal.Decimal
107- if err := tx .QueryRow (ctx , sql , args ... ).Scan (& timestamp ); err != nil {
108- return datastore .NoRevision , fmt .Errorf ("unable to executed upsert counter query: %w" , err )
126+ log .Warn ().Bool ("has-rows" , hasRows ).Msg ("unable to find row count in statistics query result" )
127+ return nil
128+ }, "SHOW STATISTICS FOR TABLE relation_tuple;" ); err != nil {
129+ return datastore.Stats {}, fmt .Errorf ("unable to query unique estimated row count: %w" , err )
109130 }
110131
111- return revisions .NewForHLC (timestamp )
132+ return datastore.Stats {
133+ UniqueID : uniqueID ,
134+ EstimatedRelationshipCount : estimatedRelCount ,
135+ ObjectTypeStatistics : datastore .ComputeObjectTypeStats (nsDefs ),
136+ }, nil
112137}
0 commit comments