From de4ecc5973f874484274f6fcd26e3b6a01d740f7 Mon Sep 17 00:00:00 2001 From: Jimmy Zelinskie Date: Wed, 20 Dec 2023 16:43:55 -0500 Subject: [PATCH] internal/datastore/spanner: force subject ID index Fixes #1687. --- internal/datastore/common/sql.go | 22 ++++++++++++++++++++-- internal/datastore/crdb/reader.go | 3 ++- internal/datastore/mysql/reader.go | 3 ++- internal/datastore/postgres/reader.go | 3 ++- internal/datastore/spanner/reader.go | 20 +++++++++++++++++++- 5 files changed, 45 insertions(+), 6 deletions(-) diff --git a/internal/datastore/common/sql.go b/internal/datastore/common/sql.go index 8ff329dce0..492f72d1b0 100644 --- a/internal/datastore/common/sql.go +++ b/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) diff --git a/internal/datastore/crdb/reader.go b/internal/datastore/crdb/reader.go index 5633a1d326..1b65c02655 100644 --- a/internal/datastore/crdb/reader.go +++ b/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)) diff --git a/internal/datastore/mysql/reader.go b/internal/datastore/mysql/reader.go index 285c764740..81ffa0e288 100644 --- a/internal/datastore/mysql/reader.go +++ b/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), diff --git a/internal/datastore/postgres/reader.go b/internal/datastore/postgres/reader.go index 290795191b..97f4d4f79a 100644 --- a/internal/datastore/postgres/reader.go +++ b/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), diff --git a/internal/datastore/spanner/reader.go b/internal/datastore/spanner/reader.go index 1642e182a4..e898556a54 100644 --- a/internal/datastore/spanner/reader.go +++ b/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),