Start on better indexing support

internal/datastore/common/index.go

@@ -0,0 +1,28 @@
+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
+
+	// IsDeprecated is true if this index is deprecated and should not be used.
+	IsDeprecated bool
+}
+
+// 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/index_test.go

@@ -0,0 +1,18 @@
+package common
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
+)
+
+func TestMatchesShape(t *testing.T) {
+	id := IndexDefinition{
+		Shapes: []queryshape.Shape{queryshape.CheckPermissionSelectDirectSubjects},
+	}
+
+	require.True(t, id.matchesShape(queryshape.CheckPermissionSelectDirectSubjects))
+	require.False(t, id.matchesShape(queryshape.CheckPermissionSelectIndirectSubjects))
+}

internal/datastore/common/relationships.go

@@ -39,14 +39,69 @@ type closeRows interface {
 	Close()
 }
 
+func runExplainIfNecessary[R Rows](ctx context.Context, builder RelationshipsQueryBuilder, tx Querier[R], explainable datastore.Explainable) error {
+	if builder.SQLExplainCallbackForTest == 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")
+		}
+
+		return builder.SQLExplainCallbackForTest(ctx, sqlString, args, builder.queryShape, explainString, expectedIndexes)
+	}, 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/relationships_test.go

@@ -0,0 +1,81 @@
+package common
+
+import (
+	"context"
+	"testing"
+
+	sq "github.com/Masterminds/squirrel"
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/datastore/options"
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
+)
+
+type fakeQuerier struct {
+	queriesRun []string
+}
+
+func (fq *fakeQuerier) QueryFunc(ctx context.Context, f func(context.Context, Rows) error, sql string, args ...interface{}) error {
+	fq.queriesRun = append(fq.queriesRun, sql)
+	return nil
+}
+
+type fakeExplainable struct{}
+
+func (fakeExplainable) BuildExplainQuery(sql string, args []any) (string, []any, error) {
+	return "SOME EXPLAIN QUERY", nil, nil
+}
+
+func (fakeExplainable) ParseExplain(explain string) (datastore.ParsedExplain, error) {
+	return datastore.ParsedExplain{}, nil
+}
+
+func (fakeExplainable) PreExplainStatements() []string {
+	return []string{"SELECT SOMETHING"}
+}
+
+func TestRunExplainIfNecessaryWithoutEnabled(t *testing.T) {
+	fq := &fakeQuerier{}
+
+	err := runExplainIfNecessary(context.Background(), RelationshipsQueryBuilder{}, fq, fakeExplainable{})
+	require.Nil(t, err)
+	require.Nil(t, fq.queriesRun)
+}
+
+func TestRunExplainIfNecessaryWithEnabled(t *testing.T) {
+	fq := &fakeQuerier{}
+
+	schema := NewSchemaInformationWithOptions(
+		WithRelationshipTableName("relationtuples"),
+		WithColNamespace("ns"),
+		WithColObjectID("object_id"),
+		WithColRelation("relation"),
+		WithColUsersetNamespace("subject_ns"),
+		WithColUsersetObjectID("subject_object_id"),
+		WithColUsersetRelation("subject_relation"),
+		WithColCaveatName("caveat"),
+		WithColCaveatContext("caveat_context"),
+		WithColExpiration("expiration"),
+		WithPlaceholderFormat(sq.Question),
+		WithPaginationFilterType(TupleComparison),
+		WithColumnOptimization(ColumnOptimizationOptionStaticValues),
+		WithNowFunction("NOW"),
+	)
+
+	filterer := NewSchemaQueryFiltererForRelationshipsSelect(*schema, 100)
+	filterer = filterer.FilterToResourceID("test")
+
+	builder := RelationshipsQueryBuilder{
+		Schema: *schema,
+		SQLExplainCallbackForTest: func(ctx context.Context, sql string, args []any, shape queryshape.Shape, explain string, expectedIndexes options.SQLIndexInformation) error {
+			return nil
+		},
+		filteringValues:  filterer.filteringColumnTracker,
+		baseQueryBuilder: filterer,
+	}
+
+	err := runExplainIfNecessary(context.Background(), builder, fq, fakeExplainable{})
+	require.Nil(t, err)
+	require.Equal(t, fq.queriesRun, []string{"SELECT SOMETHING", "SOME EXPLAIN QUERY"})
+}

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/schema_test.go

@@ -0,0 +1,38 @@
+package common
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
+)
+
+func TestExpectedIndexesForShape(t *testing.T) {
+	schema := SchemaInformation{
+		Indexes: []IndexDefinition{
+			{
+				Name: "idx1",
+				Shapes: []queryshape.Shape{
+					queryshape.CheckPermissionSelectDirectSubjects,
+				},
+			},
+			{
+				Name: "idx2",
+				Shapes: []queryshape.Shape{
+					queryshape.CheckPermissionSelectDirectSubjects,
+					queryshape.CheckPermissionSelectIndirectSubjects,
+				},
+			},
+		},
+	}
+
+	expectedIndexes := schema.expectedIndexesForShape(queryshape.Unspecified)
+	require.Empty(t, expectedIndexes)
+
+	expectedIndexes = schema.expectedIndexesForShape(queryshape.CheckPermissionSelectDirectSubjects)
+	require.Equal(t, []string{"idx1", "idx2"}, expectedIndexes.ExpectedIndexNames)
+
+	expectedIndexes = schema.expectedIndexesForShape(queryshape.CheckPermissionSelectIndirectSubjects)
+	require.Equal(t, []string{"idx2"}, expectedIndexes.ExpectedIndexNames)
+}

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,
+		SQLCheckAssertionForTest:  queryOpts.SQLCheckAssertionForTest,
+		SQLExplainCallbackForTest: queryOpts.SQLExplainCallbackForTest,
+		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
+	SQLCheckAssertionForTest  options.SQLCheckAssertionForTest
+	SQLExplainCallbackForTest options.SQLExplainCallbackForTest
+	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.SQLCheckAssertionForTest != nil {
+		b.SQLCheckAssertionForTest(sql)
 	}
 
 	return sql, args, nil