diff --git a/internal/datastore/memdb/readwrite.go b/internal/datastore/memdb/readwrite.go index eaf5d005a1..db894bbf66 100644 --- a/internal/datastore/memdb/readwrite.go +++ b/internal/datastore/memdb/readwrite.go @@ -346,18 +346,53 @@ func (rwt *memdbReadWriteTx) LegacyDeleteNamespaces(_ context.Context, nsNames [ } func (rwt *memdbReadWriteTx) BulkLoad(ctx context.Context, iter datastore.BulkWriteRelationshipSource) (uint64, error) { + rwt.mustLock() + defer rwt.Unlock() + + tx, err := rwt.txSource() + if err != nil { + return 0, err + } + + // BulkLoad has TOUCH-like (idempotent) semantics: relationships that already + // exist are silently skipped rather than causing the load to fail. The + // returned count reflects only the relationships actually inserted. var numCopied uint64 var next *tuple.Relationship - var err error + for next, err = iter.Next(ctx); next != nil && err == nil; next, err = iter.Next(ctx) { + mutation := tuple.RelationshipUpdate{Relationship: *next} + rel := &relationship{ + next.Resource.ObjectType, + next.Resource.ObjectID, + next.Resource.Relation, + next.Subject.ObjectType, + next.Subject.ObjectID, + next.Subject.Relation, + rwt.toCaveatReference(mutation), + rwt.toIntegrity(mutation), + next.OptionalExpiration, + } - updates := []tuple.RelationshipUpdate{{ - Operation: tuple.UpdateOperationCreate, - }} + found, ferr := tx.First( + tableRelationship, + indexID, + rel.namespace, + rel.resourceID, + rel.relation, + rel.subjectNamespace, + rel.subjectObjectID, + rel.subjectRelation, + ) + if ferr != nil { + return 0, fmt.Errorf("error loading existing relationship: %w", ferr) + } + if found != nil { + // Already exists; idempotently skip. + continue + } - for next, err = iter.Next(ctx); next != nil && err == nil; next, err = iter.Next(ctx) { - updates[0].Relationship = *next - if err := rwt.WriteRelationships(ctx, updates); err != nil { - return 0, err + if ierr := tx.Insert(tableRelationship, rel); ierr != nil { + return 0, fmt.Errorf("error inserting relationship: %w", ierr) } numCopied++ } diff --git a/internal/datastore/mysql/readwrite.go b/internal/datastore/mysql/readwrite.go index 9eafbdf9cf..c51cf451fa 100644 --- a/internal/datastore/mysql/readwrite.go +++ b/internal/datastore/mysql/readwrite.go @@ -21,7 +21,6 @@ import ( "github.com/authzed/spicedb/internal/datastore/common" "github.com/authzed/spicedb/internal/datastore/revisions" - log "github.com/authzed/spicedb/internal/logging" "github.com/authzed/spicedb/pkg/datastore" "github.com/authzed/spicedb/pkg/datastore/options" core "github.com/authzed/spicedb/pkg/proto/core/v1" @@ -517,7 +516,11 @@ func (rwt *mysqlReadWriteTXN) LegacyDeleteNamespaces(ctx context.Context, nsName func (rwt *mysqlReadWriteTXN) BulkLoad(ctx context.Context, iter datastore.BulkWriteRelationshipSource) (uint64, error) { var sqlStmt bytes.Buffer - sql, _, err := rwt.WriteRelsQuery.Values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10).ToSql() + // BulkLoad has TOUCH-like (idempotent) semantics: `INSERT IGNORE` silently + // skips relationships that already exist (conflicting on the + // uq_relation_tuple_living unique constraint) rather than failing the load. + // The returned count reflects only the relationships actually inserted. + sql, _, err := rwt.WriteRelsQuery.Options("IGNORE").Values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10).ToSql() if err != nil { return 0, err } @@ -564,13 +567,25 @@ func (rwt *mysqlReadWriteTXN) BulkLoad(ctx context.Context, iter datastore.BulkW } if batchLen > 0 { - log.Warn().Uint64("count", batchLen).Uint64("written", numWritten).Msg("writing batch") - if _, err := rwt.tx.Exec(sqlStmt.String(), args...); err != nil { - return 0, fmt.Errorf(errUnableToBulkWriteRelationships, fmt.Errorf("error writing batch: %w", err)) + result, execErr := rwt.tx.Exec(sqlStmt.String(), args...) + if execErr != nil { + return 0, fmt.Errorf(errUnableToBulkWriteRelationships, fmt.Errorf("error writing batch: %w", execErr)) + } + + // With INSERT IGNORE, RowsAffected counts only the rows actually + // inserted; relationships skipped because they already existed are + // excluded. + affected, affErr := result.RowsAffected() + if affErr != nil { + return 0, fmt.Errorf(errUnableToBulkWriteRelationships, affErr) } - } - numWritten += batchLen + inserted, castErr := safecast.Convert[uint64](affected) + if castErr != nil { + return 0, fmt.Errorf(errUnableToBulkWriteRelationships, castErr) + } + numWritten += inserted + } } if err != nil { return 0, fmt.Errorf(errUnableToBulkWriteRelationships, err) diff --git a/internal/datastore/postgres/common/bulk.go b/internal/datastore/postgres/common/bulk.go index 0fbbf7bb2a..0b69ed721a 100644 --- a/internal/datastore/postgres/common/bulk.go +++ b/internal/datastore/postgres/common/bulk.go @@ -2,6 +2,8 @@ package common import ( "context" + "strconv" + "strings" "github.com/ccoveille/go-safecast/v2" "github.com/jackc/pgx/v5" @@ -11,58 +13,60 @@ import ( "github.com/authzed/spicedb/pkg/tuple" ) -type tupleSourceAdapter struct { - source datastore.BulkWriteRelationshipSource - ctx context.Context +// maxParametersPerStatement is the maximum number of bind parameters that the +// PostgreSQL extended query protocol permits in a single statement. +const maxParametersPerStatement = 65535 - current *tuple.Relationship - err error - valuesBuffer []any - colNames []string -} - -// Next returns true if there is another row and makes the next row data -// available to Values(). When there are no more rows available or an error -// has occurred it returns false. -func (tg *tupleSourceAdapter) Next() bool { - tg.current, tg.err = tg.source.Next(tg.ctx) - return tg.current != nil -} +// numBaseColumns is the number of non-integrity columns written for each +// relationship. When more columns than this are requested, the additional +// columns carry the relationship integrity values. +const numBaseColumns = 9 -// Values returns the values for the current row. -func (tg *tupleSourceAdapter) Values() ([]any, error) { +// appendRelationshipValues appends the column values for a single relationship +// to args, in the same column order used by BulkLoad. When withIntegrity is +// true, the relationship integrity values are appended as well. +func appendRelationshipValues(args []any, rel *tuple.Relationship, withIntegrity bool) ([]any, error) { var caveatName string var caveatContext map[string]any - if tg.current.OptionalCaveat != nil { - caveatName = tg.current.OptionalCaveat.CaveatName - caveatContext = tg.current.OptionalCaveat.Context.AsMap() + if rel.OptionalCaveat != nil { + caveatName = rel.OptionalCaveat.CaveatName + caveatContext = rel.OptionalCaveat.Context.AsMap() } - tg.valuesBuffer[0] = tg.current.Resource.ObjectType - tg.valuesBuffer[1] = tg.current.Resource.ObjectID - tg.valuesBuffer[2] = tg.current.Resource.Relation - tg.valuesBuffer[3] = tg.current.Subject.ObjectType - tg.valuesBuffer[4] = tg.current.Subject.ObjectID - tg.valuesBuffer[5] = tg.current.Subject.Relation - tg.valuesBuffer[6] = caveatName - tg.valuesBuffer[7] = caveatContext - tg.valuesBuffer[8] = tg.current.OptionalExpiration - - if len(tg.colNames) > 9 && tg.current.OptionalIntegrity != nil { - tg.valuesBuffer[9] = tg.current.OptionalIntegrity.KeyId - tg.valuesBuffer[10] = tg.current.OptionalIntegrity.Hash - tg.valuesBuffer[11] = tg.current.OptionalIntegrity.HashedAt.AsTime() - } + args = append(args, + rel.Resource.ObjectType, + rel.Resource.ObjectID, + rel.Resource.Relation, + rel.Subject.ObjectType, + rel.Subject.ObjectID, + rel.Subject.Relation, + caveatName, + caveatContext, // PGX serializes map[string]any to JSONB columns. + rel.OptionalExpiration, + ) - return tg.valuesBuffer, nil -} + if withIntegrity { + if rel.OptionalIntegrity == nil { + return nil, spiceerrors.MustBugf("expected relationship integrity for bulk load") + } -// Err returns any error that has been encountered by the CopyFromSource. If -// this is not nil *Conn.CopyFrom will abort the copy. -func (tg *tupleSourceAdapter) Err() error { - return tg.err + args = append(args, + rel.OptionalIntegrity.KeyId, + rel.OptionalIntegrity.Hash, + rel.OptionalIntegrity.HashedAt.AsTime(), + ) + } + + return args, nil } +// BulkLoad writes all of the relationships produced by iter into tupleTableName +// using batched INSERT statements with `ON CONFLICT DO NOTHING`. The conflict +// clause gives the load TOUCH-like semantics: relationships that already exist +// are silently skipped rather than causing the load to fail, which makes +// re-importing the same data idempotent. The returned count reflects the number +// of relationships that were actually inserted (i.e. excludes ones skipped +// because they already existed). func BulkLoad( ctx context.Context, tx pgx.Tx, @@ -70,16 +74,84 @@ func BulkLoad( colNames []string, iter datastore.BulkWriteRelationshipSource, ) (uint64, error) { - adapter := &tupleSourceAdapter{ - source: iter, - ctx: ctx, - valuesBuffer: make([]any, len(colNames)), - colNames: colNames, + numCols := len(colNames) + if numCols == 0 { + return 0, spiceerrors.MustBugf("no columns provided to bulk load") + } + withIntegrity := numCols > numBaseColumns + + // The static prefix shared by every batch: `INSERT INTO (cols) VALUES `. + prefix := "INSERT INTO " + tupleTableName + " (" + strings.Join(colNames, ", ") + ") VALUES " + + // Cap the number of rows per statement so the bind-parameter count never + // exceeds the wire-protocol limit. + maxRowsPerBatch := maxParametersPerStatement / numCols + + var totalInserted uint64 + args := make([]any, 0, maxRowsPerBatch*numCols) + var sb strings.Builder + + flush := func() error { + rows := len(args) / numCols + if rows == 0 { + return nil + } + + sb.Reset() + sb.WriteString(prefix) + param := 1 + for row := 0; row < rows; row++ { + if row > 0 { + sb.WriteByte(',') + } + sb.WriteByte('(') + for col := 0; col < numCols; col++ { + if col > 0 { + sb.WriteByte(',') + } + sb.WriteByte('$') + sb.WriteString(strconv.Itoa(param)) + param++ + } + sb.WriteByte(')') + } + sb.WriteString(" ON CONFLICT DO NOTHING") + + tag, err := tx.Exec(ctx, sb.String(), args...) + if err != nil { + return err + } + + inserted, err := safecast.Convert[uint64](tag.RowsAffected()) + if err != nil { + return spiceerrors.MustBugf("number inserted was negative: %v", err) + } + totalInserted += inserted + + args = args[:0] + return nil } - copied, err := tx.CopyFrom(ctx, pgx.Identifier{tupleTableName}, colNames, adapter) - uintCopied, castErr := safecast.Convert[uint64](copied) - if castErr != nil { - return 0, spiceerrors.MustBugf("number copied was negative: %v", castErr) + + rel, err := iter.Next(ctx) + for ; err == nil && rel != nil; rel, err = iter.Next(ctx) { + args, err = appendRelationshipValues(args, rel, withIntegrity) + if err != nil { + return 0, err + } + + if len(args)/numCols >= maxRowsPerBatch { + if flushErr := flush(); flushErr != nil { + return 0, flushErr + } + } + } + if err != nil { + return 0, err } - return uintCopied, err + + if flushErr := flush(); flushErr != nil { + return 0, flushErr + } + + return totalInserted, nil } diff --git a/internal/datastore/spanner/readwrite.go b/internal/datastore/spanner/readwrite.go index 8773dd7fed..a9a6dc3803 100644 --- a/internal/datastore/spanner/readwrite.go +++ b/internal/datastore/spanner/readwrite.go @@ -416,11 +416,17 @@ func (rwt spannerReadWriteTXN) LegacyDeleteNamespaces(ctx context.Context, nsNam } func (rwt spannerReadWriteTXN) BulkLoad(ctx context.Context, iter datastore.BulkWriteRelationshipSource) (uint64, error) { + // BulkLoad has TOUCH-like (idempotent) semantics: an InsertOrUpdate mutation + // is used so that relationships which already exist do not fail the load. + // NOTE: unlike the other datastores, Spanner cannot cheaply distinguish a + // newly-inserted relationship from an existing one (mutations are applied + // blindly at commit), so the returned count reflects the number of + // relationships processed rather than only those newly inserted. var numLoaded uint64 var rel *tuple.Relationship var err error for rel, err = iter.Next(ctx); err == nil && rel != nil; rel, err = iter.Next(ctx) { - txnMut, _, err := spannerMutation(ctx, tuple.UpdateOperationCreate, *rel) + txnMut, _, err := spannerMutation(ctx, tuple.UpdateOperationTouch, *rel) if err != nil { return 0, fmt.Errorf(errUnableToBulkLoadRelationships, err) } diff --git a/pkg/datastore/test/bulk.go b/pkg/datastore/test/bulk.go index 351575db2e..eec046ea4b 100644 --- a/pkg/datastore/test/bulk.go +++ b/pkg/datastore/test/bulk.go @@ -8,9 +8,6 @@ import ( "github.com/ccoveille/go-safecast/v2" "github.com/stretchr/testify/require" - "google.golang.org/grpc/codes" - - "github.com/authzed/grpcutil" "github.com/authzed/spicedb/internal/testfixtures" "github.com/authzed/spicedb/pkg/datastore" @@ -84,44 +81,6 @@ func BulkUploadErrorsTest(t *testing.T, tester DatastoreTester) { require.Error(err) } -func BulkUploadAlreadyExistsSameCallErrorTest(t *testing.T, tester DatastoreTester) { - require := require.New(t) - ctx := t.Context() - - rawDS, err := tester.New(t, 0, veryLargeGCInterval, veryLargeGCWindow, 1) - require.NoError(err) - - ds, _ := testfixtures.StandardDatastoreWithSchema(t, rawDS) - - _, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error { - inserted, err := rwt.BulkLoad(ctx, testfixtures.NewBulkRelationshipGenerator( - testfixtures.DocumentNS.Name, - "viewer", - testfixtures.UserNS.Name, - 1, - t, - )) - require.NoError(err) - require.Equal(uint64(1), inserted) - - _, serr := rwt.BulkLoad(ctx, testfixtures.NewBulkRelationshipGenerator( - testfixtures.DocumentNS.Name, - "viewer", - testfixtures.UserNS.Name, - 1, - t, - )) - return serr - }, options.WithDisableRetries(true)) - - // NOTE: spanner does not return an error for duplicates. - if err == nil { - return - } - - grpcutil.RequireStatus(t, codes.AlreadyExists, err) -} - func BulkUploadWithCaveats(t *testing.T, tester DatastoreTester) { tc := 10 require := require.New(t) @@ -269,7 +228,16 @@ func BulkUploadEditCaveat(t *testing.T, tester DatastoreTester) { require.Equal(tc, foundChanged) } -func BulkUploadAlreadyExistsErrorTest(t *testing.T, tester DatastoreTester) { +// BulkUploadIdempotentTest verifies that BulkLoad has TOUCH-like (idempotent) +// semantics: re-loading relationships that already exist is a no-op rather than +// an error, and no duplicate relationships are created. +// +// NOTE: the count returned by BulkLoad reflects only newly-inserted rows for +// most datastores, but Spanner cannot cheaply distinguish new from existing +// relationships and reports the number processed instead. This test therefore +// asserts the no-duplicate invariant via queries (which holds for every +// datastore) and only checks the returned count for the initial load. +func BulkUploadIdempotentTest(t *testing.T, tester DatastoreTester) { require := require.New(t) ctx := t.Context() @@ -278,39 +246,71 @@ func BulkUploadAlreadyExistsErrorTest(t *testing.T, tester DatastoreTester) { ds, _ := testfixtures.StandardDatastoreWithSchema(t, rawDS) - // Bulk write a single relationship. - _, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error { - inserted, err := rwt.BulkLoad(ctx, testfixtures.NewBulkRelationshipGenerator( - testfixtures.DocumentNS.Name, - "viewer", - testfixtures.UserNS.Name, - 1, - t, - )) + // bulkLoad loads `count` deterministic relationships (object IDs count-1..0) + // in their own transaction and returns the count reported by BulkLoad. + bulkLoad := func(count int) uint64 { + var loaded uint64 + _, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error { + var ierr error + loaded, ierr = rwt.BulkLoad(ctx, testfixtures.NewBulkRelationshipGenerator( + testfixtures.DocumentNS.Name, + "viewer", + testfixtures.UserNS.Name, + count, + t, + )) + return ierr + }, options.WithDisableRetries(true)) require.NoError(err) - require.Equal(uint64(1), inserted) - return nil - }, options.WithDisableRetries(true)) - require.NoError(err) + return loaded + } - // Bulk write it again and ensure we get the expected error. - _, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error { - _, serr := rwt.BulkLoad(ctx, testfixtures.NewBulkRelationshipGenerator( - testfixtures.DocumentNS.Name, - "viewer", - testfixtures.UserNS.Name, - 1, - t, - )) - return serr - }, options.WithDisableRetries(true)) + countRelationships := func() int { + headResult, err := ds.HeadRevision(ctx) + require.NoError(err) + + iter, err := ds.SnapshotReader(headResult.Revision).QueryRelationships(ctx, datastore.RelationshipsFilter{ + OptionalResourceType: testfixtures.DocumentNS.Name, + }, options.WithQueryShape(queryshape.FindResourceOfType)) + require.NoError(err) - // NOTE: spanner does not return an error for duplicates. - if err == nil { - return + found := 0 + for _, qerr := range iter { + require.NoError(qerr) + found++ + } + return found } - grpcutil.RequireStatus(t, codes.AlreadyExists, err) + // Initial load inserts all relationships. + require.Equal(uint64(10), bulkLoad(10)) + require.Equal(10, countRelationships()) + + // Re-loading the exact same relationships is a no-op: it does not error and + // creates no duplicates (TOUCH semantics, rather than an AlreadyExists + // error). + bulkLoad(10) + require.Equal(10, countRelationships()) + + // Loading a superset silently skips the relationships that already exist and + // inserts only the previously-unseen ones. + bulkLoad(20) + require.Equal(20, countRelationships()) + + // Re-loading duplicates twice within a single transaction is also a no-op. + _, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error { + if _, ierr := rwt.BulkLoad(ctx, testfixtures.NewBulkRelationshipGenerator( + testfixtures.DocumentNS.Name, "viewer", testfixtures.UserNS.Name, 5, t, + )); ierr != nil { + return ierr + } + _, ierr := rwt.BulkLoad(ctx, testfixtures.NewBulkRelationshipGenerator( + testfixtures.DocumentNS.Name, "viewer", testfixtures.UserNS.Name, 5, t, + )) + return ierr + }, options.WithDisableRetries(true)) + require.NoError(err) + require.Equal(20, countRelationships()) } type onlyErrorSource struct{} diff --git a/pkg/datastore/test/datastore.go b/pkg/datastore/test/datastore.go index 6fecd070fd..f6076a5ad4 100644 --- a/pkg/datastore/test/datastore.go +++ b/pkg/datastore/test/datastore.go @@ -192,8 +192,7 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories) t.Run("TestBulkUpload", runner(tester, BulkUploadTest)) t.Run("TestBulkUploadErrors", runner(tester, BulkUploadErrorsTest)) - t.Run("TestBulkUploadAlreadyExistsError", runner(tester, BulkUploadAlreadyExistsErrorTest)) - t.Run("TestBulkUploadAlreadyExistsSameCallError", runner(tester, BulkUploadAlreadyExistsSameCallErrorTest)) + t.Run("TestBulkUploadIdempotent", runner(tester, BulkUploadIdempotentTest)) t.Run("TestBulkUploadEditCaveat", runner(tester, BulkUploadEditCaveat)) t.Run("TestBulkUploadWithCaveats", runner(tester, BulkUploadWithCaveats)) t.Run("TestBulkUploadWithExpiration", runner(tester, BulkUploadWithExpiration)) diff --git a/pkg/datastore/test/relationships.go b/pkg/datastore/test/relationships.go index 61da6f7024..25a845c7db 100644 --- a/pkg/datastore/test/relationships.go +++ b/pkg/datastore/test/relationships.go @@ -639,14 +639,9 @@ func CreateAlreadyExistingTest(t *testing.T, tester DatastoreTester) { require.Contains(err.Error(), "could not CREATE relationship ") grpcutil.RequireStatus(t, codes.AlreadyExists, err) - f := func(ctx context.Context, rwt datastore.ReadWriteTransaction) error { - _, err := rwt.BulkLoad(ctx, testfixtures.NewBulkRelationshipGenerator(testResourceNamespace, testReaderRelation, testUserNamespace, 1, t)) - return err - } - _, _ = ds.ReadWriteTx(ctx, f) - _, err = ds.ReadWriteTx(ctx, f) - require.Error(err) - grpcutil.RequireStatus(t, codes.AlreadyExists, err) + // NOTE: BulkLoad has TOUCH-like (idempotent) semantics and intentionally does + // not error on already-existing relationships; that behavior is covered by + // BulkUploadIdempotentTest. } // TouchAlreadyExistingTest tests touching a relationship twice.