Add ability to query relationships, sorted by subject

internal/datastore/benchmark/driver_bench_test.go

@@ -54,6 +54,11 @@ var skipped = []string{
 	spanner.Engine, // Not useful to benchmark a simulator
 }
 
+var sortOrders = map[string]options.SortOrder{
+	"ByResource": options.ByResource,
+	"BySubject":  options.BySubject,
+}
+
 func BenchmarkDatastoreDriver(b *testing.B) {
 	for _, driver := range drivers {
 		b.Run(driver.name+driver.suffix, func(b *testing.B) {
@@ -114,49 +119,64 @@ func BenchmarkDatastoreDriver(b *testing.B) {
 					}
 				})
 				b.Run("SortedSnapshotReadOnlyNamespace", func(b *testing.B) {
-					for n := 0; n < b.N; n++ {
-						iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
-							ResourceType: testfixtures.DocumentNS.Name,
-						}, options.WithSort(options.ByResource))
-						require.NoError(b, err)
-						var count int
-						for rel := iter.Next(); rel != nil; rel = iter.Next() {
-							count++
-						}
-						require.NoError(b, iter.Err())
-						iter.Close()
+					for orderName, order := range sortOrders {
+						order := order
+						b.Run(orderName, func(b *testing.B) {
+							for n := 0; n < b.N; n++ {
+								iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+									ResourceType: testfixtures.DocumentNS.Name,
+								}, options.WithSort(order))
+								require.NoError(b, err)
+								var count int
+								for rel := iter.Next(); rel != nil; rel = iter.Next() {
+									count++
+								}
+								require.NoError(b, iter.Err())
+								iter.Close()
+							}
+						})
 					}
 				})
 				b.Run("SortedSnapshotReadWithRelation", func(b *testing.B) {
-					for n := 0; n < b.N; n++ {
-						iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
-							ResourceType:             testfixtures.DocumentNS.Name,
-							OptionalResourceRelation: "viewer",
-						}, options.WithSort(options.ByResource))
-						require.NoError(b, err)
-						var count int
-						for rel := iter.Next(); rel != nil; rel = iter.Next() {
-							count++
-						}
-						require.NoError(b, iter.Err())
-						iter.Close()
+					for orderName, order := range sortOrders {
+						order := order
+						b.Run(orderName, func(b *testing.B) {
+							for n := 0; n < b.N; n++ {
+								iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+									ResourceType:             testfixtures.DocumentNS.Name,
+									OptionalResourceRelation: "viewer",
+								}, options.WithSort(order))
+								require.NoError(b, err)
+								var count int
+								for rel := iter.Next(); rel != nil; rel = iter.Next() {
+									count++
+								}
+								require.NoError(b, iter.Err())
+								iter.Close()
+							}
+						})
 					}
 				})
 				b.Run("SortedSnapshotReadAllResourceFields", func(b *testing.B) {
-					for n := 0; n < b.N; n++ {
-						randDocNum := rand.Intn(numDocuments)
-						iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
-							ResourceType:             testfixtures.DocumentNS.Name,
-							OptionalResourceIds:      []string{strconv.Itoa(randDocNum)},
-							OptionalResourceRelation: "viewer",
-						}, options.WithSort(options.ByResource))
-						require.NoError(b, err)
-						var count int
-						for rel := iter.Next(); rel != nil; rel = iter.Next() {
-							count++
-						}
-						require.NoError(b, iter.Err())
-						iter.Close()
+					for orderName, order := range sortOrders {
+						order := order
+						b.Run(orderName, func(b *testing.B) {
+							for n := 0; n < b.N; n++ {
+								randDocNum := rand.Intn(numDocuments)
+								iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+									ResourceType:             testfixtures.DocumentNS.Name,
+									OptionalResourceIds:      []string{strconv.Itoa(randDocNum)},
+									OptionalResourceRelation: "viewer",
+								}, options.WithSort(order))
+								require.NoError(b, err)
+								var count int
+								for rel := iter.Next(); rel != nil; rel = iter.Next() {
+									count++
+								}
+								require.NoError(b, iter.Err())
+								iter.Close()
+							}
+						})
 					}
 				})
 				b.Run("SnapshotReverseRead", func(b *testing.B) {

internal/datastore/common/sql.go

@@ -127,86 +127,114 @@ func (sqf SchemaQueryFilterer) TupleOrder(order options.SortOrder) SchemaQueryFi
 			sqf.schema.colUsersetObjectID,
 			sqf.schema.colUsersetRelation,
 		)
+
+	case options.BySubject:
+		sqf.queryBuilder = sqf.queryBuilder.OrderBy(
+			sqf.schema.colUsersetNamespace,
+			sqf.schema.colUsersetObjectID,
+			sqf.schema.colUsersetRelation,
+			sqf.schema.colNamespace,
+			sqf.schema.colObjectID,
+			sqf.schema.colRelation,
+		)
 	}
 
 	return sqf
 }
 
+type nameAndValue struct {
+	name  string
+	value string
+}
+
 func (sqf SchemaQueryFilterer) After(cursor *core.RelationTuple, order options.SortOrder) SchemaQueryFilterer {
-	columnsAndValues := []struct {
-		name  string
-		value string
-	}{
-		{
-			sqf.schema.colNamespace, cursor.ResourceAndRelation.Namespace,
-		},
-		{
-			sqf.schema.colObjectID, cursor.ResourceAndRelation.ObjectId,
+	columnsAndValues := map[options.SortOrder][]nameAndValue{
+		options.ByResource: {
+			{
+				sqf.schema.colNamespace, cursor.ResourceAndRelation.Namespace,
+			},
+			{
+				sqf.schema.colObjectID, cursor.ResourceAndRelation.ObjectId,
+			},
+			{
+				sqf.schema.colRelation, cursor.ResourceAndRelation.Relation,
+			},
+			{
+				sqf.schema.colUsersetNamespace, cursor.Subject.Namespace,
+			},
+			{
+				sqf.schema.colUsersetObjectID, cursor.Subject.ObjectId,
+			},
+			{
+				sqf.schema.colUsersetRelation, cursor.Subject.Relation,
+			},
 		},
-		{
-			sqf.schema.colRelation, cursor.ResourceAndRelation.Relation,
+		options.BySubject: {
+			{
+				sqf.schema.colUsersetNamespace, cursor.Subject.Namespace,
+			},
+			{
+				sqf.schema.colUsersetObjectID, cursor.Subject.ObjectId,
+			},
+			{
+				sqf.schema.colUsersetRelation, cursor.Subject.Relation,
+			},
+			{
+				sqf.schema.colNamespace, cursor.ResourceAndRelation.Namespace,
+			},
+			{
+				sqf.schema.colObjectID, cursor.ResourceAndRelation.ObjectId,
+			},
+			{
+				sqf.schema.colRelation, cursor.ResourceAndRelation.Relation,
+			},
 		},
-		{
-			sqf.schema.colUsersetNamespace, cursor.Subject.Namespace,
-		},
-		{
-			sqf.schema.colUsersetObjectID, cursor.Subject.ObjectId,
-		},
-		{
-			sqf.schema.colUsersetRelation, cursor.Subject.Relation,
-		},
-	}
+	}[order]
 
 	switch sqf.schema.paginationFilterType {
 	case TupleComparison:
-		switch order {
-		case options.ByResource:
-			// For performance reasons, remove any column names that have static values in the query.
-			columnNames := make([]string, 0, len(columnsAndValues))
-			valueSlots := make([]any, 0, len(columnsAndValues))
-			comparisonSlotCount := 0
-
-			for _, cav := range columnsAndValues {
-				if sqf.filteringColumnCounts[cav.name] != 1 {
-					columnNames = append(columnNames, cav.name)
-					valueSlots = append(valueSlots, cav.value)
-					comparisonSlotCount++
-				}
+		// For performance reasons, remove any column names that have static values in the query.
+		columnNames := make([]string, 0, len(columnsAndValues))
+		valueSlots := make([]any, 0, len(columnsAndValues))
+		comparisonSlotCount := 0
+
+		for _, cav := range columnsAndValues {
+			if sqf.filteringColumnCounts[cav.name] != 1 {
+				columnNames = append(columnNames, cav.name)
+				valueSlots = append(valueSlots, cav.value)
+				comparisonSlotCount++
 			}
+		}
 
-			if comparisonSlotCount > 0 {
-				comparisonTuple := "(" + strings.Join(columnNames, ",") + ") > (" + strings.Repeat(",?", comparisonSlotCount)[1:] + ")"
-				sqf.queryBuilder = sqf.queryBuilder.Where(
-					comparisonTuple,
-					valueSlots...,
-				)
-			}
+		if comparisonSlotCount > 0 {
+			comparisonTuple := "(" + strings.Join(columnNames, ",") + ") > (" + strings.Repeat(",?", comparisonSlotCount)[1:] + ")"
+			sqf.queryBuilder = sqf.queryBuilder.Where(
+				comparisonTuple,
+				valueSlots...,
+			)
 		}
 
 	case ExpandedLogicComparison:
-		switch order {
-		case options.ByResource:
-			// For performance reasons, remove any column names that have static values in the query.
-			orClause := sq.Or{}
-
-			for index, cav := range columnsAndValues {
-				if sqf.filteringColumnCounts[cav.name] != 1 {
-					andClause := sq.And{}
-					for _, previous := range columnsAndValues[0:index] {
-						if sqf.filteringColumnCounts[previous.name] != 1 {
-							andClause = append(andClause, sq.Eq{previous.name: previous.value})
-						}
+		// For performance reasons, remove any column names that have static values in the query.
+		orClause := sq.Or{}
+
+		for index, cav := range columnsAndValues {
+			if sqf.filteringColumnCounts[cav.name] != 1 {
+				andClause := sq.And{}
+				for _, previous := range columnsAndValues[0:index] {
+					if sqf.filteringColumnCounts[previous.name] != 1 {
+						andClause = append(andClause, sq.Eq{previous.name: previous.value})
 					}
-
-					andClause = append(andClause, sq.Gt{cav.name: cav.value})
-					orClause = append(orClause, andClause)
 				}
-			}
 
-			if len(orClause) > 0 {
-				sqf.queryBuilder = sqf.queryBuilder.Where(orClause)
+				andClause = append(andClause, sq.Gt{cav.name: cav.value})
+				orClause = append(orClause, andClause)
 			}
 		}
+
+		if len(orClause) > 0 {
+			sqf.queryBuilder = sqf.queryBuilder.Where(orClause)
+		}
 	}
 
 	return sqf

internal/datastore/memdb/readonly.go

@@ -4,10 +4,14 @@ import (
 	"context"
 	"fmt"
 	"runtime"
+	"sort"
+
+	"github.com/authzed/spicedb/pkg/spiceerrors"
 
 	"github.com/hashicorp/go-memdb"
 	"github.com/jzelinskie/stringz"
 
+	"github.com/authzed/spicedb/internal/datastore/common"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -61,8 +65,20 @@ func (r *memdbReader) QueryRelationships(
 	)
 	filteredIterator := memdb.NewFilterIterator(bestIterator, matchingRelationshipsFilterFunc)
 
-	iter := newMemdbTupleIterator(filteredIterator, queryOpts.Limit, queryOpts.Sort)
-	return iter, nil
+	switch queryOpts.Sort {
+	case options.Unsorted:
+		fallthrough
+
+	case options.ByResource:
+		iter := newMemdbTupleIterator(filteredIterator, queryOpts.Limit, queryOpts.Sort)
+		return iter, nil
+
+	case options.BySubject:
+		return newSubjectSortedIterator(filteredIterator, queryOpts.Limit)
+
+	default:
+		return nil, spiceerrors.MustBugf("unsupported sort order: %v", queryOpts.Sort)
+	}
 }
 
 func mustHaveBeenClosed(iter *memdbTupleIterator) {
@@ -350,11 +366,50 @@ func makeCursorFilterFn(after *core.RelationTuple, order options.SortOrder) func
 						(less(tpl.subjectNamespace, tpl.subjectObjectID, tpl.subjectRelation, after.Subject) ||
 							eq(tpl.subjectNamespace, tpl.subjectObjectID, tpl.subjectRelation, after.Subject)))
 			}
+		case options.BySubject:
+			return func(tpl *relationship) bool {
+				return less(tpl.subjectNamespace, tpl.subjectObjectID, tpl.subjectRelation, after.Subject) ||
+					(eq(tpl.subjectNamespace, tpl.subjectObjectID, tpl.subjectRelation, after.Subject) &&
+						(less(tpl.namespace, tpl.resourceID, tpl.relation, after.ResourceAndRelation) ||
+							eq(tpl.subjectNamespace, tpl.subjectObjectID, tpl.subjectRelation, after.Subject)))
+			}
 		}
 	}
 	return noopCursorFilter
 }
 
+func newSubjectSortedIterator(it memdb.ResultIterator, limit *uint64) (datastore.RelationshipIterator, error) {
+	results := make([]*core.RelationTuple, 0)
+
+	// Coalesce all of the results into memory
+	for foundRaw := it.Next(); foundRaw != nil; foundRaw = it.Next() {
+		rt, err := foundRaw.(*relationship).RelationTuple()
+		if err != nil {
+			return nil, err
+		}
+
+		results = append(results, rt)
+	}
+
+	// Sort them by subject
+	sort.Slice(results, func(i, j int) bool {
+		lhsRes := results[i].ResourceAndRelation
+		lhsSub := results[i].Subject
+		rhsRes := results[j].ResourceAndRelation
+		rhsSub := results[j].Subject
+		return less(lhsSub.Namespace, lhsSub.ObjectId, lhsSub.Relation, rhsSub) ||
+			(eq(lhsSub.Namespace, lhsSub.ObjectId, lhsSub.Relation, rhsSub) &&
+				(less(lhsRes.Namespace, lhsRes.ObjectId, lhsRes.Relation, rhsRes)))
+	})
+
+	// Limit them if requested
+	if limit != nil && uint64(len(results)) > *limit {
+		results = results[0:*limit]
+	}
+
+	return common.NewSliceRelationshipIterator(results, options.BySubject), nil
+}
+
 func noopCursorFilter(_ *relationship) bool {
 	return false
 }

internal/testfixtures/datastore.go

@@ -320,7 +320,7 @@ func (tc TupleChecker) VerifyOrderedIteratorResults(iter datastore.RelationshipI
 		expectedStr := tuple.MustString(tpl)
 
 		found := iter.Next()
-		tc.Require.NotNil(found, "expected %s", expectedStr)
+		tc.Require.NotNil(found, "expected %s, but found no additional results", expectedStr)
 
 		foundStr := tuple.MustString(found)
 		tc.Require.Equal(expectedStr, foundStr)

pkg/datastore/options/options.go

@@ -17,9 +17,10 @@ const (
 	// ByResource sorts the relationships by the resource component first
 	ByResource
 
-	// NOTE: if you are intending to implement BySubject, it was fully implemented
-	// but not performance tested, and the code can be recovered from source control
-	// BySubject
+	// BySubject sorts the relationships by the subject component first. Note that
+	// BySubject might be quite a bit slower than ByResource, as relationships are
+	// indexed by resource.
+	BySubject
 )
 
 type Cursor *core.RelationTuple