Change CRDB datastore to use a combined deletion query, rather than a…

… query per relationship
authzed · Feb 29, 2024 · 3387b62 · 3387b62
1 parent 57206a1
commit 3387b62
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 9 deletions.
diff --git a/internal/datastore/crdb/readwrite.go b/internal/datastore/crdb/readwrite.go
@@ -96,6 +96,10 @@ func (rwt *crdbReadWriteTXN) WriteRelationships(ctx context.Context, mutations [
 	bulkTouch := queryTouchTuple
 	var bulkTouchCount int64
 
+	bulkDelete := queryDeleteTuples
+	bulkDeleteOr := sq.Or{}
+	var bulkDeleteCount int64
+
 	// Process the actual updates
 	for _, mutation := range mutations {
 		rel := mutation.Tuple
@@ -139,20 +143,27 @@ func (rwt *crdbReadWriteTXN) WriteRelationships(ctx context.Context, mutations [
 			bulkWriteCount++
 		case core.RelationTupleUpdate_DELETE:
 			rwt.relCountChange--
-			sql, args, err := queryDeleteTuples.Where(exactRelationshipClause(rel)).ToSql()
-			if err != nil {
-				return fmt.Errorf(errUnableToWriteRelationships, err)
-			}
+			bulkDeleteOr = append(bulkDeleteOr, exactRelationshipClause(rel))
+			bulkDeleteCount++
 
-			if _, err := rwt.tx.Exec(ctx, sql, args...); err != nil {
-				return fmt.Errorf(errUnableToWriteRelationships, err)
-			}
 		default:
 			log.Ctx(ctx).Error().Stringer("operation", mutation.Operation).Msg("unknown operation type")
 			return fmt.Errorf("unknown mutation operation: %s", mutation.Operation)
 		}
 	}
 
+	if bulkDeleteCount > 0 {
+		bulkDelete = bulkDelete.Where(bulkDeleteOr)
+		sql, args, err := bulkDelete.ToSql()
+		if err != nil {
+			return fmt.Errorf(errUnableToWriteRelationships, err)
+		}
+
+		if _, err := rwt.tx.Exec(ctx, sql, args...); err != nil {
+			return fmt.Errorf(errUnableToWriteRelationships, err)
+		}
+	}
+
 	bulkUpdateQueries := make([]sq.InsertBuilder, 0, 2)
 	if bulkWriteCount > 0 {
 		bulkUpdateQueries = append(bulkUpdateQueries, bulkWrite)

diff --git a/pkg/datastore/test/datastore.go b/pkg/datastore/test/datastore.go
@@ -98,10 +98,12 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories)
 	t.Run("TestWriteDeleteWrite", func(t *testing.T) { WriteDeleteWriteTest(t, tester) })
 	t.Run("TestCreateAlreadyExisting", func(t *testing.T) { CreateAlreadyExistingTest(t, tester) })
 	t.Run("TestTouchAlreadyExisting", func(t *testing.T) { TouchAlreadyExistingTest(t, tester) })
-	t.Run("TestCreateDeleteTouchTest", func(t *testing.T) { CreateDeleteTouchTest(t, tester) })
-	t.Run("TestCreateTouchDeleteTouchTest", func(t *testing.T) { CreateTouchDeleteTouchTest(t, tester) })
+	t.Run("TestCreateDeleteTouch", func(t *testing.T) { CreateDeleteTouchTest(t, tester) })
+	t.Run("TestDeleteOneThousandIndividualInOneCall", func(t *testing.T) { DeleteOneThousandIndividualInOneCallTest(t, tester) })
+	t.Run("TestCreateTouchDeleteTouch", func(t *testing.T) { CreateTouchDeleteTouchTest(t, tester) })
 	t.Run("TestTouchAlreadyExistingCaveated", func(t *testing.T) { TouchAlreadyExistingCaveatedTest(t, tester) })
 	t.Run("TestBulkDeleteRelationships", func(t *testing.T) { BulkDeleteRelationshipsTest(t, tester) })
+	t.Run("TestDeleteCaveatedTuple", func(t *testing.T) { DeleteCaveatedTupleTest(t, tester) })
 
 	t.Run("TestMultipleReadsInRWT", func(t *testing.T) { MultipleReadsInRWTTest(t, tester) })
 	t.Run("TestConcurrentWriteSerialization", func(t *testing.T) { ConcurrentWriteSerializationTest(t, tester) })

diff --git a/pkg/datastore/test/tuples.go b/pkg/datastore/test/tuples.go
@@ -636,6 +636,65 @@ func CreateDeleteTouchTest(t *testing.T, tester DatastoreTester) {
 	ensureTuples(ctx, require, ds, tpl1, tpl2)
 }
 
+// DeleteOneThousandIndividualInOneCallTest tests deleting 1000 relationships, individually.
+func DeleteOneThousandIndividualInOneCallTest(t *testing.T, tester DatastoreTester) {
+	require := require.New(t)
+
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(err)
+
+	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
+	ctx := context.Background()
+
+	// Write the 1000 relationships.
+	tuples := make([]*core.RelationTuple, 0, 1000)
+	for i := 0; i < 1000; i++ {
+		tpl := makeTestTuple("foo", fmt.Sprintf("user%d", i))
+		tuples = append(tuples, tpl)
+	}
+
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_CREATE, tuples...)
+	require.NoError(err)
+	ensureTuples(ctx, require, ds, tuples...)
+
+	// Add an extra tuple.
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_CREATE, makeTestTuple("foo", "extra"))
+	require.NoError(err)
+	ensureTuples(ctx, require, ds, makeTestTuple("foo", "extra"))
+
+	// Delete the first 1000 tuples.
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_DELETE, tuples...)
+	require.NoError(err)
+	ensureNotTuples(ctx, require, ds, tuples...)
+
+	// Ensure the extra tuple is still present.
+	ensureTuples(ctx, require, ds, makeTestTuple("foo", "extra"))
+}
+
+// DeleteCaveatedTupleTest tests deleting a relationship with a caveat.
+func DeleteCaveatedTupleTest(t *testing.T, tester DatastoreTester) {
+	require := require.New(t)
+
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(err)
+
+	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
+	ctx := context.Background()
+
+	tpl := tuple.Parse("test/resource:someresource#viewer@test/user:someuser[somecaveat]")
+
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_CREATE, tpl)
+	require.NoError(err)
+	ensureTuples(ctx, require, ds, tpl)
+
+	// Delete the tuple.
+	withoutCaveat := tuple.Parse("test/resource:someresource#viewer@test/user:someuser")
+
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_DELETE, withoutCaveat)
+	require.NoError(err)
+	ensureNotTuples(ctx, require, ds, tpl, withoutCaveat)
+}
+
 // CreateTouchDeleteTouchTest tests writing a relationship, touching it, deleting it, and then touching it.
 func CreateTouchDeleteTouchTest(t *testing.T, tester DatastoreTester) {
 	require := require.New(t)