Merge pull request #2326 from josephschorr/spanner-index-force

Force indexes on Spanner and change subject sort order to match index

Author: vroldanbet

internal/datastore/spanner/indexes.go

@@ -0,0 +1,51 @@
+package spanner
+
+import (
+	"github.com/authzed/spicedb/internal/datastore/common"
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
+)
+
+var IndexRelationshipBySubject = common.IndexDefinition{
+	Name:       "ix_relation_tuple_by_subject",
+	ColumnsSQL: `relation_tuple (userset_object_id, userset_namespace, userset_relation, namespace, relation)`,
+}
+
+var NoIndexingHint common.IndexingHint = nil
+
+// IndexingHintForQueryShape returns an indexing hint for the given query shape, if any.
+func IndexingHintForQueryShape(schema common.SchemaInformation, qs queryshape.Shape) common.IndexingHint {
+	if schema.IntegrityEnabled {
+		// Don't force anything since we don't have the other indexes.
+		return NoIndexingHint
+	}
+
+	switch qs {
+	case queryshape.MatchingResourcesForSubject:
+		return forcedIndex{IndexRelationshipBySubject}
+
+	case queryshape.FindSubjectOfType:
+		return forcedIndex{IndexRelationshipBySubject}
+
+	default:
+		return NoIndexingHint
+	}
+}
+
+// forcedIndex is an index hint that forces the use of a specific index.
+type forcedIndex struct {
+	index common.IndexDefinition
+}
+
+func (f forcedIndex) FromSQLSuffix() (string, error) {
+	return "", nil
+}
+
+func (f forcedIndex) FromTable(existingTableName string) (string, error) {
+	return existingTableName + "@{FORCE_INDEX=" + f.index.Name + "}", nil
+}
+
+func (f forcedIndex) SQLPrefix() (string, error) {
+	return "", nil
+}
+
+var _ common.IndexingHint = forcedIndex{}

internal/datastore/spanner/reader.go

@@ -140,6 +140,10 @@ func (sr spannerReader) QueryRelationships(
 		return nil, err
 	}
 
+	builtOpts := options.NewQueryOptionsWithOptions(opts...)
+	indexingHint := IndexingHintForQueryShape(sr.schema, builtOpts.QueryShape)
+	qBuilder = qBuilder.WithIndexingHint(indexingHint)
+
 	return sr.executor.ExecuteQuery(ctx, qBuilder, opts...)
 }
 
@@ -162,6 +166,9 @@ func (sr spannerReader) ReverseQueryRelationships(
 			FilterToRelation(queryOpts.ResRelation.Relation)
 	}
 
+	indexingHint := IndexingHintForQueryShape(sr.schema, queryOpts.QueryShapeForReverse)
+	qBuilder = qBuilder.WithIndexingHint(indexingHint)
+
 	return sr.executor.ExecuteQuery(ctx,
 		qBuilder,
 		options.WithLimit(queryOpts.LimitForReverse),

internal/datastore/spanner/spanner.go

@@ -207,6 +207,18 @@ func NewSpannerDatastore(ctx context.Context, database string, opts ...Option) (
 		common.WithPlaceholderFormat(sq.AtP),
 		common.WithNowFunction("CURRENT_TIMESTAMP"),
 		common.WithColumnOptimization(config.columnOptimizationOption),
+
+		// NOTE: this order differs from the default because the index
+		// used for sorting by subject (ix_relation_tuple_by_subject) is
+		// defined with the userset object ID first.
+		common.SetSortBySubjectColumnOrder([]string{
+			colUsersetObjectID,
+			colUsersetNamespace,
+			colUsersetRelation,
+			colNamespace,
+			colRelation,
+			colObjectID,
+		}),
 	)
 
 	ds := &spannerDatastore{

pkg/datastore/queryshape/queryshape.go

@@ -73,7 +73,7 @@ const (
 	// The query shape selects a subject of a given type, which is specified by
 	// providing the subject type. The other fields are never specified.
 	//
-	// ✅ resource_type, *️⃣ resource_id, *️⃣ resource_relation, ✅ subject_type, *️⃣ subject_id, *️⃣ subject_relation, *️⃣ caveat, *️⃣ expiration
+	// *️⃣ resource_type, *️⃣ resource_id, *️⃣ resource_relation, ✅ subject_type, *️⃣ subject_id, *️⃣ subject_relation, *️⃣ caveat, *️⃣ expiration
 	FindSubjectOfType = "find-subject-of-type"
 
 	// FindResourceOfTypeAndRelation indicates that the query is selecting a single