Start on better indexing support

internal/datastore/common/index.go

@@ -0,0 +1,25 @@
+package common
+
+import "github.com/authzed/spicedb/pkg/datastore/queryshape"
+
+// IndexDefinition is a definition of an index for a datastore.
+type IndexDefinition struct {
+	// Name is the unique name for the index.
+	Name string
+
+	// ColumnsSQL is the SQL fragment of the columns over which this index will apply.
+	ColumnsSQL string
+
+	// Shapes are those query shapes for which this index should be used.
+	Shapes []queryshape.Shape
+}
+
+// matchesShape returns true if the index matches the given shape.
+func (id IndexDefinition) matchesShape(shape queryshape.Shape) bool {
+	for _, s := range id.Shapes {
+		if s == shape {
+			return true
+		}
+	}
+	return false
+}

internal/datastore/common/relationships.go

@@ -39,14 +39,70 @@ type closeRows interface {
 	Close()
 }
 
+func runExplainIfNecessary[R Rows](ctx context.Context, builder RelationshipsQueryBuilder, tx Querier[R], explainable datastore.Explainable) error {
+	if builder.sqlExplainCallback == nil {
+		return nil
+	}
+
+	// Determine the expected index names via the schema.
+	expectedIndexes := builder.Schema.expectedIndexesForShape(builder.queryShape)
+
+	// Run any pre-explain statements.
+	for _, statement := range explainable.PreExplainStatements() {
+		if err := tx.QueryFunc(ctx, func(ctx context.Context, rows R) error {
+			rows.Next()
+			return nil
+		}, statement); err != nil {
+			return fmt.Errorf(errUnableToQueryRels, err)
+		}
+	}
+
+	// Run the query with EXPLAIN ANALYZE.
+	sqlString, args, err := builder.SelectSQL()
+	if err != nil {
+		return fmt.Errorf(errUnableToQueryRels, err)
+	}
+
+	explainSQL, explainArgs, err := explainable.BuildExplainQuery(sqlString, args)
+	if err != nil {
+		return fmt.Errorf(errUnableToQueryRels, err)
+	}
+
+	err = tx.QueryFunc(ctx, func(ctx context.Context, rows R) error {
+		explainString := ""
+		for rows.Next() {
+			var explain string
+			if err := rows.Scan(&explain); err != nil {
+				return fmt.Errorf(errUnableToQueryRels, fmt.Errorf("scan err: %w", err))
+			}
+			explainString += explain + "\n"
+		}
+		if explainString == "" {
+			return fmt.Errorf("received empty explain")
+		}
+
+		builder.sqlExplainCallback(ctx, sqlString, args, builder.queryShape, explainString, expectedIndexes)
+		return nil
+	}, explainSQL, explainArgs...)
+	if err != nil {
+		return fmt.Errorf(errUnableToQueryRels, err)
+	}
+
+	return nil
+}
+
 // QueryRelationships queries relationships for the given query and transaction.
-func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, builder RelationshipsQueryBuilder, tx Querier[R]) (datastore.RelationshipIterator, error) {
+func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, builder RelationshipsQueryBuilder, tx Querier[R], explainable datastore.Explainable) (datastore.RelationshipIterator, error) {
 	span := trace.SpanFromContext(ctx)
 	sqlString, args, err := builder.SelectSQL()
 	if err != nil {
 		return nil, fmt.Errorf(errUnableToQueryRels, err)
 	}
 
+	if err := runExplainIfNecessary(ctx, builder, tx, explainable); err != nil {
+		return nil, err
+	}
+
 	var resourceObjectType string
 	var resourceObjectID string
 	var resourceRelation string

internal/datastore/common/schema.go

@@ -3,6 +3,8 @@ package common
 import (
 	sq "github.com/Masterminds/squirrel"
 
+	"github.com/authzed/spicedb/pkg/datastore/options"
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
 	"github.com/authzed/spicedb/pkg/spiceerrors"
 )
 
@@ -35,6 +37,9 @@ type SchemaInformation struct {
 	ColIntegrityHash      string `debugmap:"visible"`
 	ColIntegrityTimestamp string `debugmap:"visible"`
 
+	// Indexes are the indexes to use for this schema.
+	Indexes []IndexDefinition `debugmap:"visible"`
+
 	// PaginationFilterType is the type of pagination filter to use for this schema.
 	PaginationFilterType PaginationFilterType `debugmap:"visible"`
 
@@ -54,6 +59,17 @@ type SchemaInformation struct {
 	ExpirationDisabled bool `debugmap:"visible"`
 }
 
+// expectedIndexesForShape returns the expected index names for a given query shape.
+func (si SchemaInformation) expectedIndexesForShape(shape queryshape.Shape) options.SQLIndexInformation {
+	expectedIndexes := options.SQLIndexInformation{}
+	for _, index := range si.Indexes {
+		if index.matchesShape(shape) {
+			expectedIndexes.ExpectedIndexNames = append(expectedIndexes.ExpectedIndexNames, index.Name)
+		}
+	}
+	return expectedIndexes
+}
+
 func (si SchemaInformation) debugValidate() {
 	spiceerrors.DebugAssert(func() bool {
 		si.mustValidate()

internal/datastore/common/sql.go

@@ -16,6 +16,7 @@ import (
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
 	"github.com/authzed/spicedb/pkg/spiceerrors"
 )
 
@@ -663,12 +664,14 @@ func (exc QueryRelationshipsExecutor) ExecuteQuery(
 	query.queryBuilder = query.queryBuilder.From(from)
 
 	builder := RelationshipsQueryBuilder{
-		Schema:           query.schema,
-		SkipCaveats:      queryOpts.SkipCaveats,
-		SkipExpiration:   queryOpts.SkipExpiration,
-		sqlAssertion:     queryOpts.SQLAssertion,
-		filteringValues:  query.filteringColumnTracker,
-		baseQueryBuilder: query,
+		Schema:             query.schema,
+		SkipCaveats:        queryOpts.SkipCaveats,
+		SkipExpiration:     queryOpts.SkipExpiration,
+		sqlCheckAssertion:  queryOpts.SQLCheckAssertion,
+		sqlExplainCallback: queryOpts.SQLExplainCallback,
+		filteringValues:    query.filteringColumnTracker,
+		queryShape:         queryOpts.QueryShape,
+		baseQueryBuilder:   query,
 	}
 
 	return exc.Executor(ctx, builder)
@@ -681,9 +684,11 @@ type RelationshipsQueryBuilder struct {
 	SkipCaveats    bool
 	SkipExpiration bool
 
-	filteringValues  columnTrackerMap
-	baseQueryBuilder SchemaQueryFilterer
-	sqlAssertion     options.Assertion
+	filteringValues    columnTrackerMap
+	baseQueryBuilder   SchemaQueryFilterer
+	sqlCheckAssertion  options.SQLCheckAssertion
+	sqlExplainCallback options.SQLExplainCallback
+	queryShape         queryshape.Shape
 }
 
 // withCaveats returns true if caveats should be included in the query.
@@ -752,8 +757,8 @@ func (b RelationshipsQueryBuilder) SelectSQL() (string, []any, error) {
 		return "", nil, err
 	}
 
-	if b.sqlAssertion != nil {
-		b.sqlAssertion(sql)
+	if b.sqlCheckAssertion != nil {
+		b.sqlCheckAssertion(sql)
 	}
 
 	return sql, args, nil

internal/datastore/crdb/caveat.go

@@ -9,6 +9,7 @@ import (
 	sq "github.com/Masterminds/squirrel"
 	"github.com/jackc/pgx/v5"
 
+	"github.com/authzed/spicedb/internal/datastore/crdb/schema"
 	"github.com/authzed/spicedb/internal/datastore/revisions"
 	"github.com/authzed/spicedb/pkg/datastore"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -17,14 +18,14 @@ import (
 var (
 	upsertCaveatSuffix = fmt.Sprintf(
 		"ON CONFLICT (%s) DO UPDATE SET %s = excluded.%s",
-		colCaveatName,
-		colCaveatDefinition,
-		colCaveatDefinition,
+		schema.ColCaveatName,
+		schema.ColCaveatDefinition,
+		schema.ColCaveatDefinition,
 	)
-	writeCaveat  = psql.Insert(tableCaveat).Columns(colCaveatName, colCaveatDefinition).Suffix(upsertCaveatSuffix)
-	readCaveat   = psql.Select(colCaveatDefinition, colTimestamp)
-	listCaveat   = psql.Select(colCaveatName, colCaveatDefinition, colTimestamp).OrderBy(colCaveatName)
-	deleteCaveat = psql.Delete(tableCaveat)
+	writeCaveat  = psql.Insert(schema.TableCaveat).Columns(schema.ColCaveatName, schema.ColCaveatDefinition).Suffix(upsertCaveatSuffix)
+	readCaveat   = psql.Select(schema.ColCaveatDefinition, schema.ColTimestamp)
+	listCaveat   = psql.Select(schema.ColCaveatName, schema.ColCaveatDefinition, schema.ColTimestamp).OrderBy(schema.ColCaveatName)
+	deleteCaveat = psql.Delete(schema.TableCaveat)
 )
 
 const (
@@ -35,7 +36,7 @@ const (
 )
 
 func (cr *crdbReader) ReadCaveatByName(ctx context.Context, name string) (*core.CaveatDefinition, datastore.Revision, error) {
-	query := cr.addFromToQuery(readCaveat.Where(sq.Eq{colCaveatName: name}), tableCaveat)
+	query := cr.addFromToQuery(readCaveat.Where(sq.Eq{schema.ColCaveatName: name}), schema.TableCaveat)
 	sql, args, err := query.ToSql()
 	if err != nil {
 		return nil, datastore.NoRevision, fmt.Errorf(errReadCaveat, name, err)
@@ -80,9 +81,9 @@ type bytesAndTimestamp struct {
 }
 
 func (cr *crdbReader) lookupCaveats(ctx context.Context, caveatNames []string) ([]datastore.RevisionedCaveat, error) {
-	caveatsWithNames := cr.addFromToQuery(listCaveat, tableCaveat)
+	caveatsWithNames := cr.addFromToQuery(listCaveat, schema.TableCaveat)
 	if len(caveatNames) > 0 {
-		caveatsWithNames = caveatsWithNames.Where(sq.Eq{colCaveatName: caveatNames})
+		caveatsWithNames = caveatsWithNames.Where(sq.Eq{schema.ColCaveatName: caveatNames})
 	}
 
 	sql, args, err := caveatsWithNames.ToSql()
@@ -158,7 +159,7 @@ func (rwt *crdbReadWriteTXN) WriteCaveats(ctx context.Context, caveats []*core.C
 }
 
 func (rwt *crdbReadWriteTXN) DeleteCaveats(ctx context.Context, names []string) error {
-	deleteCaveatClause := deleteCaveat.Where(sq.Eq{colCaveatName: names})
+	deleteCaveatClause := deleteCaveat.Where(sq.Eq{schema.ColCaveatName: names})
 	sql, args, err := deleteCaveatClause.ToSql()
 	if err != nil {
 		return fmt.Errorf(errDeleteCaveats, err)