Add extended relationships filtering to watch API

e2e/go.mod

@@ -5,7 +5,7 @@ go 1.21
 toolchain go1.21.1
 
 require (
-	github.com/authzed/authzed-go v0.10.1
+	github.com/authzed/authzed-go v0.10.2-0.20240206183056-781a5f5d1b3c
 	github.com/authzed/grpcutil v0.0.0-20240123092924-129dc0a6a6e1
 	github.com/authzed/spicedb v1.29.1
 	github.com/brianvoe/gofakeit/v6 v6.23.2
@@ -50,7 +50,7 @@ require (
 	go.opentelemetry.io/otel/metric v1.24.0 // indirect
 	go.opentelemetry.io/otel/trace v1.24.0 // indirect
 	golang.org/x/crypto v0.18.0 // indirect
-	golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
+	golang.org/x/exp v0.0.0-20240112132812-db7319d0e0e3 // indirect
 	golang.org/x/mod v0.14.0 // indirect
 	golang.org/x/net v0.20.0 // indirect
 	golang.org/x/sync v0.6.0 // indirect

e2e/go.sum

@@ -20,8 +20,8 @@ github.com/Masterminds/squirrel v1.5.4 h1:uUcX/aBc8O7Fg9kaISIUsHXdKuqehiXAMQTYX8
 github.com/Masterminds/squirrel v1.5.4/go.mod h1:NNaOrjSoIDfDA40n7sr2tPNZRfjzjA400rg+riTZj10=
 github.com/antlr/antlr4/runtime/Go/antlr/v4 v4.0.0-20230512164433-5d1fd1a340c9 h1:goHVqTbFX3AIo0tzGr14pgfAW2ZfPChKO21Z9MGf/gk=
 github.com/antlr/antlr4/runtime/Go/antlr/v4 v4.0.0-20230512164433-5d1fd1a340c9/go.mod h1:pSwJ0fSY5KhvocuWSx4fz3BA8OrA1bQn+K1Eli3BRwM=
-github.com/authzed/authzed-go v0.10.1 h1:0aX2Ox9PPPknID92kLs/FnmhCmfl6Ni16v3ZTLsds5M=
-github.com/authzed/authzed-go v0.10.1/go.mod h1:ZsaFPCiMjwT0jLW0gCyYzh3elHqhKDDGGRySyykXwqc=
+github.com/authzed/authzed-go v0.10.2-0.20240206183056-781a5f5d1b3c h1:w71KppS/+mCDkNXR8vmwXGVt/1dLt+axcIq+qK7f7To=
+github.com/authzed/authzed-go v0.10.2-0.20240206183056-781a5f5d1b3c/go.mod h1:bS4eeTw/ZpCunZHePrt1MAcvOggAwL8Djh8cx+CR33g=
 github.com/authzed/cel-go v0.17.5 h1:lfpkNrR99B5QRHg5qdG9oLu/kguVlZC68VJuMk8tH9Y=
 github.com/authzed/cel-go v0.17.5/go.mod h1:XL/zEq5hKGVF8aOdMbG7w+BQPihLjY2W8N+UIygDA2I=
 github.com/authzed/grpcutil v0.0.0-20240123092924-129dc0a6a6e1 h1:zBfQzia6Hz45pJBeURTrv1b6HezmejB6UmiGuBilHZM=
@@ -256,8 +256,8 @@ golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPh
 golang.org/x/crypto v0.18.0 h1:PGVlW0xEltQnzFZ55hkuX5+KLyrMYhHld1YHO4AKcdc=
 golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1mg=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
-golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
-golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
+golang.org/x/exp v0.0.0-20240112132812-db7319d0e0e3 h1:hNQpMuAJe5CtcUqCXaWga3FHu+kQvCqcsoVaQgSV60o=
+golang.org/x/exp v0.0.0-20240112132812-db7319d0e0e3/go.mod h1:idGWGoKP1toJGkd5/ig9ZLuPcZBC3ewk7SzmH0uou08=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=

internal/services/v1/relationships.go

@@ -319,7 +319,7 @@ func (ps *permissionServer) WriteRelationships(ctx context.Context, req *v1.Writ
 
 		// Validate the preconditions.
 		for _, precond := range req.OptionalPreconditions {
-			if err := ps.validateRelationshipsFilter(ctx, precond.Filter, rwt); err != nil {
+			if err := ps.validatePrecondition(ctx, precond, rwt); err != nil {
 				return err
 			}
 		}
@@ -395,7 +395,7 @@ func (ps *permissionServer) DeleteRelationships(ctx context.Context, req *v1.Del
 		})
 
 		for _, precond := range req.OptionalPreconditions {
-			if err := ps.validatePrecondition(precond); err != nil {
+			if err := ps.validatePrecondition(ctx, precond, rwt); err != nil {
 				return err
 			}
 		}
@@ -464,14 +464,10 @@ func (ps *permissionServer) DeleteRelationships(ctx context.Context, req *v1.Del
 	}, nil
 }
 
-func (ps *permissionServer) validatePrecondition(precond *v1.Precondition) error {
-	if precond.Filter == nil {
+func (ps *permissionServer) validatePrecondition(ctx context.Context, precond *v1.Precondition, reader datastore.Reader) error {
+	if precond.EqualVT(&v1.Precondition{}) || precond.Filter == nil {
 		return NewEmptyPreconditionErr()
 	}
 
-	if precond.Filter.EqualVT(&v1.RelationshipFilter{}) {
-		return NewEmptyPreconditionErr()
-	}
-
-	return nil
+	return ps.validateRelationshipsFilter(ctx, precond.Filter, reader)
 }

internal/services/v1/relationships_test.go

@@ -643,7 +643,7 @@ func TestInvalidWriteRelationship(t *testing.T) {
 			[]*v1.RelationshipFilter{{}},
 			nil,
 			codes.InvalidArgument,
-			"one of the specified preconditions is empty",
+			"the relationship filter provided is not valid",
 		},
 		{
 			"good precondition, invalid update",
@@ -1057,11 +1057,7 @@ func TestDeleteRelationships(t *testing.T) {
 				},
 			},
 			expectedCode: codes.OK,
-			deleted: map[string]struct{}{
-				"document:specialplan#editor@user:multiroleguy":              {},
-				"document:specialplan#viewer_and_editor@user:missingrolegal": {},
-				"document:specialplan#viewer_and_editor@user:multiroleguy":   {},
-			},
+			deleted:      map[string]struct{}{},
 		},
 		{
 			name: "delete unknown resource type",

internal/services/v1/watch.go

@@ -13,6 +13,7 @@ import (
 	"github.com/authzed/spicedb/internal/middleware/usagemetrics"
 	"github.com/authzed/spicedb/internal/services/shared"
 	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	dispatchv1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
 	"github.com/authzed/spicedb/pkg/tuple"
@@ -38,14 +39,24 @@ func NewWatchServer(heartbeatDuration time.Duration) v1.WatchServiceServer {
 }
 
 func (ws *watchServer) Watch(req *v1.WatchRequest, stream v1.WatchService_WatchServer) error {
-	ctx := stream.Context()
-	ds := datastoremw.MustFromContext(ctx)
+	if len(req.GetOptionalObjectTypes()) > 0 && len(req.OptionalRelationshipFilters) > 0 {
+		return status.Errorf(codes.InvalidArgument, "cannot specify both object types and relationship filters")
+	}
+
+	objectTypes := mapz.NewSet[string](req.GetOptionalObjectTypes()...)
+	filters := make([]datastore.RelationshipsFilter, 0, len(req.OptionalRelationshipFilters))
+	for _, filter := range req.OptionalRelationshipFilters {
+		dsFilter, err := datastore.RelationshipsFilterFromPublicFilter(filter)
+		if err != nil {
+			return status.Errorf(codes.InvalidArgument, "failed to parse relationship filter: %s", err)
+		}
 
-	objectTypesMap := make(map[string]struct{})
-	for _, objectType := range req.GetOptionalObjectTypes() {
-		objectTypesMap[objectType] = struct{}{}
+		filters = append(filters, dsFilter)
 	}
 
+	ctx := stream.Context()
+	ds := datastoremw.MustFromContext(ctx)
+
 	var afterRevision datastore.Revision
 	if req.OptionalStartCursor != nil && req.OptionalStartCursor.Token != "" {
 		decodedRevision, err := zedtoken.DecodeRevision(req.OptionalStartCursor, ds)
@@ -74,7 +85,7 @@ func (ws *watchServer) Watch(req *v1.WatchRequest, stream v1.WatchService_WatchS
 		select {
 		case update, ok := <-updates:
 			if ok {
-				filtered := filterUpdates(objectTypesMap, update.RelationshipChanges)
+				filtered := filterUpdates(objectTypes, filters, update.RelationshipChanges)
 				if len(filtered) > 0 {
 					if err := stream.Send(&v1.WatchResponse{
 						Updates:        filtered,
@@ -97,20 +108,38 @@ func (ws *watchServer) Watch(req *v1.WatchRequest, stream v1.WatchService_WatchS
 	}
 }
 
-func filterUpdates(objectTypes map[string]struct{}, candidates []*core.RelationTupleUpdate) []*v1.RelationshipUpdate {
+func filterUpdates(objectTypes *mapz.Set[string], filters []datastore.RelationshipsFilter, candidates []*core.RelationTupleUpdate) []*v1.RelationshipUpdate {
 	updates := tuple.UpdatesToRelationshipUpdates(candidates)
 
-	if len(objectTypes) == 0 {
+	if objectTypes.IsEmpty() && len(filters) == 0 {
 		return updates
 	}
 
-	var filtered []*v1.RelationshipUpdate
+	filtered := make([]*v1.RelationshipUpdate, 0, len(updates))
 	for _, update := range updates {
 		objectType := update.GetRelationship().GetResource().GetObjectType()
 
-		if _, ok := objectTypes[objectType]; ok {
-			filtered = append(filtered, update)
+		if !objectTypes.IsEmpty() && !objectTypes.Has(objectType) {
+			continue
+		}
+
+		if len(filters) > 0 {
+			// If there are filters, we need to check if the update matches any of them.
+			matched := false
+			for _, filter := range filters {
+				// TODO(jschorr): Maybe we should add TestRelationship to avoid the conversion?
+				if filter.Test(tuple.MustFromRelationship(update.GetRelationship())) {
+					matched = true
+					break
+				}
+			}
+
+			if !matched {
+				continue
+			}
 		}
+
+		filtered = append(filtered, update)
 	}
 
 	return filtered

internal/services/v1/watch_test.go

@@ -51,12 +51,13 @@ func update(
 
 func TestWatch(t *testing.T) {
 	testCases := []struct {
-		name              string
-		objectTypesFilter []string
-		startCursor       *v1.ZedToken
-		mutations         []*v1.RelationshipUpdate
-		expectedCode      codes.Code
-		expectedUpdates   []*v1.RelationshipUpdate
+		name                string
+		objectTypesFilter   []string
+		relationshipFilters []*v1.RelationshipFilter
+		startCursor         *v1.ZedToken
+		mutations           []*v1.RelationshipUpdate
+		expectedCode        codes.Code
+		expectedUpdates     []*v1.RelationshipUpdate
 	}{
 		{
 			name:         "unfiltered watch",
@@ -86,6 +87,76 @@ func TestWatch(t *testing.T) {
 				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document2", "viewer", "user", "user1"),
 			},
 		},
+		{
+			name:         "watch with relationship filters",
+			expectedCode: codes.OK,
+			relationshipFilters: []*v1.RelationshipFilter{
+				{
+					ResourceType: "document",
+				},
+			},
+			mutations: []*v1.RelationshipUpdate{
+				update(v1.RelationshipUpdate_OPERATION_CREATE, "document", "document1", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document2", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_DELETE, "folder", "auditors", "viewer", "user", "auditor"),
+			},
+			expectedUpdates: []*v1.RelationshipUpdate{
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document1", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document2", "viewer", "user", "user1"),
+			},
+		},
+		{
+			name:         "watch with modified relationship filters",
+			expectedCode: codes.OK,
+			relationshipFilters: []*v1.RelationshipFilter{
+				{
+					ResourceType: "folder",
+				},
+			},
+			mutations: []*v1.RelationshipUpdate{
+				update(v1.RelationshipUpdate_OPERATION_CREATE, "document", "document1", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document2", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_DELETE, "folder", "auditors", "viewer", "user", "auditor"),
+			},
+			expectedUpdates: []*v1.RelationshipUpdate{
+				update(v1.RelationshipUpdate_OPERATION_DELETE, "folder", "auditors", "viewer", "user", "auditor"),
+			},
+		},
+		{
+			name:         "watch with resource ID prefix",
+			expectedCode: codes.OK,
+			relationshipFilters: []*v1.RelationshipFilter{
+				{
+					OptionalResourceIdPrefix: "document1",
+				},
+			},
+			mutations: []*v1.RelationshipUpdate{
+				update(v1.RelationshipUpdate_OPERATION_CREATE, "document", "document1", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document2", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_DELETE, "folder", "auditors", "viewer", "user", "auditor"),
+			},
+			expectedUpdates: []*v1.RelationshipUpdate{
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document1", "viewer", "user", "user1"),
+			},
+		},
+		{
+			name:         "watch with shorter resource ID prefix",
+			expectedCode: codes.OK,
+			relationshipFilters: []*v1.RelationshipFilter{
+				{
+					OptionalResourceIdPrefix: "doc",
+				},
+			},
+			mutations: []*v1.RelationshipUpdate{
+				update(v1.RelationshipUpdate_OPERATION_CREATE, "document", "document1", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document2", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_DELETE, "folder", "auditors", "viewer", "user", "auditor"),
+			},
+			expectedUpdates: []*v1.RelationshipUpdate{
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document1", "viewer", "user", "user1"),
+				update(v1.RelationshipUpdate_OPERATION_TOUCH, "document", "document2", "viewer", "user", "user1"),
+			},
+		},
 		{
 			name:         "invalid zedtoken",
 			startCursor:  &v1.ZedToken{Token: "bad-token"},
@@ -116,8 +187,9 @@ func TestWatch(t *testing.T) {
 			defer cancel()
 
 			stream, err := client.Watch(ctx, &v1.WatchRequest{
-				OptionalObjectTypes: tc.objectTypesFilter,
-				OptionalStartCursor: cursor,
+				OptionalObjectTypes:         tc.objectTypesFilter,
+				OptionalRelationshipFilters: tc.relationshipFilters,
+				OptionalStartCursor:         cursor,
 			})
 			require.NoError(err)