Merge pull request #2376 from josephschorr/crdb-index-logic

Improve CRDB index forcing logic to be more fine-grain

Author: josephschorr

internal/datastore/crdb/readwrite.go

@@ -130,7 +130,10 @@ func (rwt *crdbReadWriteTXN) insertQuery() sq.InsertBuilder {
 	return psql.Insert(rwt.schema.RelationshipTableName)
 }
 
-func (rwt *crdbReadWriteTXN) queryDeleteTuples(index common.IndexDefinition) sq.DeleteBuilder {
+func (rwt *crdbReadWriteTXN) queryDeleteTuples(index *common.IndexDefinition) sq.DeleteBuilder {
+	if index == nil {
+		return psql.Delete(rwt.schema.RelationshipTableName)
+	}
 	return psql.Delete(rwt.schema.RelationshipTableName + "@" + index.Name)
 }
 
@@ -281,7 +284,7 @@ func (rwt *crdbReadWriteTXN) WriteRelationships(ctx context.Context, mutations [
 	bulkTouch := rwt.queryTouchTuple()
 	var bulkTouchCount int64
 
-	bulkDelete := rwt.queryDeleteTuples(schema.IndexPrimaryKey)
+	bulkDelete := rwt.queryDeleteTuples(nil)
 	bulkDeleteOr := sq.Or{}
 	var bulkDeleteCount int64
 
@@ -533,7 +536,7 @@ func (rwt *crdbReadWriteTXN) DeleteNamespaces(ctx context.Context, nsNames ...st
 		return fmt.Errorf(errUnableToDeleteConfig, err)
 	}
 
-	deleteTupleSQL, deleteTupleArgs, err := rwt.queryDeleteTuples(schema.IndexPrimaryKey).Where(sq.Or(tplClauses)).ToSql()
+	deleteTupleSQL, deleteTupleArgs, err := rwt.queryDeleteTuples(nil).Where(sq.Or(tplClauses)).ToSql()
 	if err != nil {
 		return fmt.Errorf(errUnableToDeleteConfig, err)
 	}

internal/datastore/crdb/schema/indexes.go

@@ -86,20 +86,52 @@ func IndexingHintForQueryShape(schema common.SchemaInformation, qs queryshape.Sh
 	}
 }
 
-// IndexForFilter returns the index to use for a given relationships filter.
-func IndexForFilter(schema common.SchemaInformation, filter datastore.RelationshipsFilter) common.IndexDefinition {
-	if schema.IntegrityEnabled {
-		// Don't force anything since we don't have the other indexes.
-		return IndexPrimaryKey
+// IndexForFilter returns the index to use for a given relationships filter or nil if no index is forced.
+func IndexForFilter(schema common.SchemaInformation, filter datastore.RelationshipsFilter) *common.IndexDefinition {
+	resourceFieldDepth := 0
+	if filter.OptionalResourceType != "" {
+		resourceFieldDepth = 1
+		if len(filter.OptionalResourceIds) > 0 || filter.OptionalResourceIDPrefix != "" {
+			resourceFieldDepth = 2
+			if filter.OptionalResourceRelation != "" {
+				resourceFieldDepth = 3
+			}
+		}
+	}
+
+	subjectFieldDepth := 0
+	for _, subjectSelector := range filter.OptionalSubjectsSelectors {
+		sfd := 0
+		if len(subjectSelector.OptionalSubjectIds) > 0 {
+			sfd = 1
+			if subjectSelector.OptionalSubjectType != "" {
+				sfd = 2
+				if subjectSelector.RelationFilter.NonEllipsisRelation != "" {
+					sfd = 3
+				}
+			}
+		}
+		subjectFieldDepth = max(subjectFieldDepth, sfd)
+	}
+
+	if resourceFieldDepth == 0 && subjectFieldDepth == 0 {
+		return nil
 	}
 
-	// If there is no resource type specified, but there is a subject filter, use the
-	// inverse index.
-	if filter.OptionalResourceType == "" && len(filter.OptionalSubjectsSelectors) > 0 {
-		return IndexRelationshipBySubject
+	if resourceFieldDepth > subjectFieldDepth {
+		return &IndexPrimaryKey
+	}
+
+	if resourceFieldDepth < subjectFieldDepth {
+		if schema.IntegrityEnabled {
+			// Don't force this index since it doesn't exist for integrity-enabled schemas.
+			return nil
+		}
+
+		return &IndexRelationshipBySubject
 	}
 
-	return IndexPrimaryKey
+	return nil
 }
 
 // forcedIndex is an index hint that forces the use of a specific index.

internal/datastore/crdb/schema/indexes_test.go

@@ -0,0 +1,135 @@
+package schema
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/internal/datastore/common"
+	"github.com/authzed/spicedb/pkg/datastore"
+)
+
+func TestIndexForFilter(t *testing.T) {
+	tests := []struct {
+		name     string
+		filter   datastore.RelationshipsFilter
+		expected string
+	}{
+		{
+			name:     "no filter",
+			filter:   datastore.RelationshipsFilter{},
+			expected: "",
+		},
+		{
+			name:     "filter by resource type",
+			filter:   datastore.RelationshipsFilter{OptionalResourceType: "foo"},
+			expected: "pk_relation_tuple",
+		},
+		{
+			name: "filter by resource type and relation",
+			filter: datastore.RelationshipsFilter{
+				OptionalResourceType:     "foo",
+				OptionalResourceRelation: "bar",
+			},
+			expected: "pk_relation_tuple",
+		},
+		{
+			name: "filter by resource type, resource ID and relation",
+			filter: datastore.RelationshipsFilter{
+				OptionalResourceType:     "foo",
+				OptionalResourceIds:      []string{"baz"},
+				OptionalResourceRelation: "bar",
+			},
+			expected: "pk_relation_tuple",
+		},
+		{
+			name: "filter by subject type, subject ID and relation",
+			filter: datastore.RelationshipsFilter{
+				OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+					{
+						OptionalSubjectType: "foo",
+						OptionalSubjectIds:  []string{"bar"},
+						RelationFilter: datastore.SubjectRelationFilter{
+							NonEllipsisRelation: "baz",
+						},
+					},
+				},
+			},
+			expected: "ix_relation_tuple_by_subject",
+		},
+		{
+			name: "filter by subject type, subject ID",
+			filter: datastore.RelationshipsFilter{
+				OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+					{
+						OptionalSubjectType: "foo",
+						OptionalSubjectIds:  []string{"bar"},
+					},
+				},
+			},
+			expected: "ix_relation_tuple_by_subject",
+		},
+		{
+			name: "filter by subject relation, subject ID",
+			filter: datastore.RelationshipsFilter{
+				OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+					{
+						OptionalSubjectIds: []string{"bar"},
+						RelationFilter: datastore.SubjectRelationFilter{
+							NonEllipsisRelation: "baz",
+						},
+					},
+				},
+			},
+			expected: "ix_relation_tuple_by_subject",
+		},
+		{
+			name: "filter by subject type",
+			filter: datastore.RelationshipsFilter{
+				OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+					{
+						OptionalSubjectType: "foo",
+					},
+				},
+			},
+			expected: "",
+		},
+		{
+			name: "filter by resource type and subject type",
+			filter: datastore.RelationshipsFilter{
+				OptionalResourceType: "foo",
+				OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+					{
+						OptionalSubjectType: "foo",
+					},
+				},
+			},
+			expected: "pk_relation_tuple",
+		},
+		{
+			name: "filter by resource type and subject object ID",
+			filter: datastore.RelationshipsFilter{
+				OptionalResourceType: "foo",
+				OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+					{
+						OptionalSubjectIds: []string{"bar"},
+					},
+				},
+			},
+			expected: "",
+		},
+	}
+
+	schema := Schema(common.ColumnOptimizationOptionNone, false, false)
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			index := IndexForFilter(*schema, test.filter)
+			if test.expected == "" {
+				require.Nil(t, index)
+			} else {
+				require.NotNil(t, index)
+				require.Equal(t, test.expected, index.Name)
+			}
+		})
+	}
+}