internal/datastore/spanner: force subject ID index

Fixes #1687.

internal/datastore/common/sql.go

@@ -494,6 +494,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 +514,16 @@ func (tqs QueryExecutor) ExecuteQuery(
  limit = int(*queryOpts.Limit)
  }
 
- toExecute := query.limit(uint64(limit))
- sql, args, err := toExecute.queryBuilder.ToSql()
+ queryBuilder := query.limit(uint64(limit)).queryBuilder
+
+ if fn != nil {
+ var err error
+ if queryBuilder, err = fn(queryBuilder); err != nil {
+ return nil, err
+ }
+ }
+
+ sql, args, err := queryBuilder.ToSql()
  if err != nil {
  return nil, err
  }
@@ -531,6 +540,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(sq.SelectBuilder) (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

@@ -3,9 +3,11 @@ package spanner
 import (
  "context"
  "fmt"
+ "strings"
  "time"
 
  "cloud.google.com/go/spanner"
+ sq "github.com/Masterminds/squirrel"
  "go.opentelemetry.io/otel/attribute"
  "go.opentelemetry.io/otel/trace"
  "google.golang.org/grpc/codes"
@@ -16,6 +18,27 @@ 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) (sq.SelectBuilder, error) {
+ sqlQuery, _, err := b.ToSql()
+ if err != nil {
+ return b, err
+ }
+
+ // The query is setting the subject ID and should be forced to use
+ // the `idx_relation_tuple_by_subject` index.
+ if strings.Contains(sqlQuery, "userset_object_id") {
+ // 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 +65,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 +89,7 @@ func (sr spannerReader) ReverseQueryRelationships(
 
  return sr.executor.ExecuteQuery(ctx,
  qBuilder,
+ forceIndex,
  options.WithLimit(queryOpts.LimitForReverse),
  options.WithAfter(queryOpts.AfterForReverse),
  options.WithSort(queryOpts.SortForReverse),