internal/datastore/spanner: force subject ID index

Fixes #1687.

internal/datastore/common/sql.go

@@ -10,6 +10,7 @@ import (
 	"github.com/jzelinskie/stringz"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
+	"golang.org/x/exp/maps"
 
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
@@ -494,6 +495,7 @@ type QueryExecutor struct {
 func (tqs QueryExecutor) ExecuteQuery(
 	ctx context.Context,
 	query SchemaQueryFilterer,
+	fn QueryOptimizationFunc,
 	opts ...options.QueryOptionsOption,
 ) (datastore.RelationshipIterator, error) {
 	queryOpts := options.NewQueryOptionsWithOptions(opts...)
@@ -513,8 +515,18 @@ func (tqs QueryExecutor) ExecuteQuery(
 		limit = int(*queryOpts.Limit)
 	}
 
-	toExecute := query.limit(uint64(limit))
-	sql, args, err := toExecute.queryBuilder.ToSql()
+	limitedSchemaFilterer := query.limit(uint64(limit))
+	queryBuilder := limitedSchemaFilterer.queryBuilder
+
+	if fn != nil {
+		var err error
+		queryBuilder, err = fn(queryBuilder, maps.Keys(limitedSchemaFilterer.filteringColumnCounts))
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	sql, args, err := queryBuilder.ToSql()
 	if err != nil {
 		return nil, err
 	}
@@ -531,6 +543,12 @@ func (tqs QueryExecutor) ExecuteQuery(
 	return NewSliceRelationshipIterator(queryTuples, queryOpts.Sort), nil
 }
 
+// QueryOptimizationFunc is a function apply a final transform to a
+// ready-to-execute query.
+//
+// This is useful for scenarios like forcing a particular index.
+type QueryOptimizationFunc func(b sq.SelectBuilder, setColumns []string) (sq.SelectBuilder, error)
+
 // ExecuteQueryFunc is a function that can be used to execute a single rendered SQL query.
 type ExecuteQueryFunc func(ctx context.Context, sql string, args []any) ([]*core.RelationTuple, error)
 

internal/datastore/crdb/reader.go

@@ -100,7 +100,7 @@ func (cr *crdbReader) QueryRelationships(
 		return nil, err
 	}
 
-	return cr.executor.ExecuteQuery(ctx, qBuilder, opts...)
+	return cr.executor.ExecuteQuery(ctx, qBuilder, nil, opts...)
 }
 
 func (cr *crdbReader) ReverseQueryRelationships(
@@ -126,6 +126,7 @@ func (cr *crdbReader) ReverseQueryRelationships(
 	return cr.executor.ExecuteQuery(
 		ctx,
 		qBuilder,
+		nil,
 		options.WithLimit(queryOpts.LimitForReverse),
 		options.WithAfter(queryOpts.AfterForReverse),
 		options.WithSort(queryOpts.SortForReverse))

internal/datastore/mysql/reader.go

@@ -59,7 +59,7 @@ func (mr *mysqlReader) QueryRelationships(
 		return nil, err
 	}
 
-	return mr.executor.ExecuteQuery(ctx, qBuilder, opts...)
+	return mr.executor.ExecuteQuery(ctx, qBuilder, nil, opts...)
 }
 
 func (mr *mysqlReader) ReverseQueryRelationships(
@@ -85,6 +85,7 @@ func (mr *mysqlReader) ReverseQueryRelationships(
 	return mr.executor.ExecuteQuery(
 		ctx,
 		qBuilder,
+		nil,
 		options.WithLimit(queryOpts.LimitForReverse),
 		options.WithAfter(queryOpts.AfterForReverse),
 		options.WithSort(queryOpts.SortForReverse),

internal/datastore/postgres/reader.go

@@ -66,7 +66,7 @@ func (r *pgReader) QueryRelationships(
 		return nil, err
 	}
 
-	return r.executor.ExecuteQuery(ctx, qBuilder, opts...)
+	return r.executor.ExecuteQuery(ctx, qBuilder, nil, opts...)
 }
 
 func (r *pgReader) ReverseQueryRelationships(
@@ -90,6 +90,7 @@ func (r *pgReader) ReverseQueryRelationships(
 
 	return r.executor.ExecuteQuery(ctx,
 		qBuilder,
+		nil,
 		options.WithLimit(queryOpts.LimitForReverse),
 		options.WithAfter(queryOpts.AfterForReverse),
 		options.WithSort(queryOpts.SortForReverse),

internal/datastore/spanner/reader.go

@@ -6,8 +6,10 @@ import (
 	"time"
 
 	"cloud.google.com/go/spanner"
+	sq "github.com/Masterminds/squirrel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
+	"golang.org/x/exp/slices"
 	"google.golang.org/grpc/codes"
 
 	"github.com/authzed/spicedb/internal/datastore/common"
@@ -16,6 +18,21 @@ import (
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 )
 
+// forceIndex is used to force critical indices.
+//
+// You can find the documentation for forcing indices here:
+// https://cloud.google.com/spanner/docs/secondary-indexes#index-directive
+func forceIndex(b sq.SelectBuilder, setColumns []string) (sq.SelectBuilder, error) {
+	// The query is setting subject ID and not the resource ID; this is a
+	// reverse lookup.
+	if slices.Contains(setColumns, colUsersetObjectID) && !slices.Contains(setColumns, colObjectID) {
+		// When using the Postgres dialect, Spanner uses magic comments to
+		// force indices.
+		b = b.From(tableRelationship + " /*@ FORCE_INDEX=idx_relation_tuple_by_subject */")
+	}
+	return b, nil
+}
+
 // The underlying Spanner shared read transaction interface is not exposed, so we re-create
 // the subsection of it which we need here.
 // https://github.com/googleapis/google-cloud-go/blob/a33861fe46be42ae150d6015ad39dae6e35e04e8/spanner/transaction.go#L55
@@ -42,7 +59,7 @@ func (sr spannerReader) QueryRelationships(
 		return nil, err
 	}
 
-	return sr.executor.ExecuteQuery(ctx, qBuilder, opts...)
+	return sr.executor.ExecuteQuery(ctx, qBuilder, forceIndex, opts...)
 }
 
 func (sr spannerReader) ReverseQueryRelationships(
@@ -66,6 +83,7 @@ func (sr spannerReader) ReverseQueryRelationships(
 
 	return sr.executor.ExecuteQuery(ctx,
 		qBuilder,
+		forceIndex,
 		options.WithLimit(queryOpts.LimitForReverse),
 		options.WithAfter(queryOpts.AfterForReverse),
 		options.WithSort(queryOpts.SortForReverse),