github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/crdb/stats.go (about)

     1  package crdb
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"slices"
     7  
     8  	"github.com/Masterminds/squirrel"
     9  	"github.com/jackc/pgx/v5"
    10  	"github.com/rs/zerolog/log"
    11  
    12  	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
    13  	"github.com/authzed/spicedb/pkg/datastore"
    14  )
    15  
    16  const (
    17  	tableMetadata = "metadata"
    18  	colUniqueID   = "unique_id"
    19  )
    20  
    21  var (
    22  	queryReadUniqueID = psql.Select(colUniqueID).From(tableMetadata)
    23  	uniqueID          string
    24  )
    25  
    26  func (cds *crdbDatastore) Statistics(ctx context.Context) (datastore.Stats, error) {
    27  	if len(uniqueID) == 0 {
    28  		sql, args, err := queryReadUniqueID.ToSql()
    29  		if err != nil {
    30  			return datastore.Stats{}, fmt.Errorf("unable to prepare unique ID sql: %w", err)
    31  		}
    32  		if err := cds.readPool.QueryRowFunc(ctx, func(ctx context.Context, row pgx.Row) error {
    33  			return row.Scan(&uniqueID)
    34  		}, sql, args...); err != nil {
    35  			return datastore.Stats{}, fmt.Errorf("unable to query unique ID: %w", err)
    36  		}
    37  	}
    38  
    39  	var nsDefs []datastore.RevisionedNamespace
    40  	if err := cds.readPool.BeginTxFunc(ctx, pgx.TxOptions{AccessMode: pgx.ReadOnly}, func(tx pgx.Tx) error {
    41  		_, err := tx.Exec(ctx, "SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()")
    42  		if err != nil {
    43  			return fmt.Errorf("unable to read namespaces: %w", err)
    44  		}
    45  		nsDefs, err = loadAllNamespaces(ctx, pgxcommon.QuerierFuncsFor(tx), func(sb squirrel.SelectBuilder, fromStr string) squirrel.SelectBuilder {
    46  			return sb.From(fromStr)
    47  		})
    48  		if err != nil {
    49  			return fmt.Errorf("unable to read namespaces: %w", err)
    50  		}
    51  		return nil
    52  	}); err != nil {
    53  		return datastore.Stats{}, err
    54  	}
    55  
    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  			}
    61  
    62  			return nil
    63  		}); err != nil {
    64  			return datastore.Stats{}, err
    65  		}
    66  	}
    67  
    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  		}
   125  
   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)
   130  	}
   131  
   132  	return datastore.Stats{
   133  		UniqueID:                   uniqueID,
   134  		EstimatedRelationshipCount: estimatedRelCount,
   135  		ObjectTypeStatistics:       datastore.ComputeObjectTypeStats(nsDefs),
   136  	}, nil
   137  }