Merge pull request #2176 from josephschorr/bulk-caveat-load

Add caveated bulk load test to datastore tests

internal/testfixtures/generator.go

@@ -7,7 +7,10 @@ import (
 	"testing"
 	"time"
 
+	"google.golang.org/protobuf/types/known/structpb"
+
 	"github.com/authzed/spicedb/pkg/datastore"
+	corev1 "github.com/authzed/spicedb/pkg/proto/core/v1"
 	"github.com/authzed/spicedb/pkg/tuple"
 )
 
@@ -39,6 +42,7 @@ func NewBulkRelationshipGenerator(objectType, relation, subjectType string, coun
 		relation,
 		subjectType,
 		false,
+		false,
 	}
 }
 
@@ -49,6 +53,7 @@ type BulkRelationshipGenerator struct {
 	relation       string
 	subjectType    string
 	WithExpiration bool
+	WithCaveat     bool
 }
 
 func (btg *BulkRelationshipGenerator) Next(_ context.Context) (*tuple.Relationship, error) {
@@ -63,6 +68,21 @@ func (btg *BulkRelationshipGenerator) Next(_ context.Context) (*tuple.Relationsh
 		expiration = &exp
 	}
 
+	var caveat *corev1.ContextualizedCaveat
+	if btg.WithCaveat {
+		c, err := structpb.NewStruct(map[string]interface{}{
+			"secret": "1235",
+		})
+		if err != nil {
+			btg.t.Fatalf("failed to create struct: %v", err)
+		}
+
+		caveat = &corev1.ContextualizedCaveat{
+			CaveatName: "test",
+			Context:    c,
+		}
+	}
+
 	return &tuple.Relationship{
 		RelationshipReference: tuple.RelationshipReference{
 			Resource: tuple.ObjectAndRelation{
@@ -76,6 +96,7 @@ func (btg *BulkRelationshipGenerator) Next(_ context.Context) (*tuple.Relationsh
 				Relation:   datastore.Ellipsis,
 			},
 		},
+		OptionalCaveat:     caveat,
 		OptionalExpiration: expiration,
 	}, nil
 }

pkg/datastore/test/bulk.go

@@ -122,6 +122,49 @@ func BulkUploadAlreadyExistsSameCallErrorTest(t *testing.T, tester DatastoreTest
 	grpcutil.RequireStatus(t, codes.AlreadyExists, err)
 }
 
+func BulkUploadWithCaveats(t *testing.T, tester DatastoreTester) {
+	tc := 10
+	require := require.New(t)
+	ctx := context.Background()
+
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(err)
+
+	ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
+	bulkSource := testfixtures.NewBulkRelationshipGenerator(
+		testfixtures.DocumentNS.Name,
+		"caveated_viewer",
+		testfixtures.UserNS.Name,
+		tc,
+		t,
+	)
+	bulkSource.WithCaveat = true
+
+	// This is statically defined so we can cast straight.
+	uintTc, _ := safecast.ToUint64(tc)
+	lastRevision, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+		loaded, err := rwt.BulkLoad(ctx, bulkSource)
+		require.NoError(err)
+		require.Equal(uintTc, loaded)
+		return err
+	})
+	require.NoError(err)
+
+	iter, err := ds.SnapshotReader(lastRevision).QueryRelationships(ctx, datastore.RelationshipsFilter{
+		OptionalResourceType: testfixtures.DocumentNS.Name,
+	})
+	require.NoError(err)
+
+	for found, err := range iter {
+		require.NoError(err)
+		require.Nil(found.OptionalExpiration)
+		require.NotNil(found.OptionalCaveat)
+		require.NotEmpty(found.OptionalCaveat.CaveatName)
+		require.NotNil(found.OptionalCaveat.Context)
+		require.Equal(map[string]any{"secret": "1235"}, found.OptionalCaveat.Context.AsMap())
+	}
+}
+
 func BulkUploadWithExpiration(t *testing.T, tester DatastoreTester) {
 	tc := 10
 	require := require.New(t)

pkg/datastore/test/datastore.go

@@ -161,8 +161,9 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories,
 	t.Run("TestBulkUploadErrors", runner(tester, BulkUploadErrorsTest))
 	t.Run("TestBulkUploadAlreadyExistsError", runner(tester, BulkUploadAlreadyExistsErrorTest))
 	t.Run("TestBulkUploadAlreadyExistsSameCallError", runner(tester, BulkUploadAlreadyExistsSameCallErrorTest))
-	t.Run("BulkUploadEditCaveat", runner(tester, BulkUploadEditCaveat))
-	t.Run("BulkUploadWithExpiration", runner(tester, BulkUploadWithExpiration))
+	t.Run("TestBulkUploadEditCaveat", runner(tester, BulkUploadEditCaveat))
+	t.Run("TestBulkUploadWithCaveats", runner(tester, BulkUploadWithCaveats))
+	t.Run("TestBulkUploadWithExpiration", runner(tester, BulkUploadWithExpiration))
 
 	if !except.Stats() {
 		t.Run("TestStats", runner(tester, StatsTest))