From ed2bf8c2019e8ab49637fb447d013f507bbb462e Mon Sep 17 00:00:00 2001 From: Christophe Poucet Date: Fri, 26 Jun 2026 13:10:22 +0100 Subject: [PATCH] feat(datastore): idempotent (TOUCH-like) bulk relationship import ImportBulkRelationships routes through the datastore BulkLoad, which used strict CREATE semantics: importing a relationship that already existed aborted the load with an AlreadyExists error. Make BulkLoad idempotent (TOUCH-like) across all datastores, so already-existing relationships are silently skipped rather than failing the load, which makes re-imports idempotent. Per datastore: - postgres/crdb: replace the shared CopyFrom-based BulkLoad with batched INSERT ... ON CONFLICT DO NOTHING. Bare ON CONFLICT DO NOTHING works for both backends (Postgres infers the uq_relation_tuple_living_xid partial unique index; CockroachDB infers pk_relation_tuple). Rows are batched to stay under the 65535 bind-parameter wire limit. - mysql: use INSERT IGNORE. - memdb: skip relationships that already exist when scanning the source. - spanner: use an InsertOrUpdate mutation instead of Insert. Returned count semantics: BulkLoad now returns the number of relationships actually inserted (excluding skipped duplicates) for postgres, crdb, mysql, and memdb. Spanner cannot cheaply distinguish new from existing rows (mutations are applied blindly at commit), so it continues to report the number processed; this is documented at the call site. Testing: - Add BulkUploadIdempotentTest to the shared datastore suite, asserting the no-duplicate invariant via queries (holds for every datastore) plus the initial-load count. - Remove the now-obsolete BulkUploadAlreadyExistsError / BulkUploadAlreadyExistsSameCallError tests and the bulk-load portion of CreateAlreadyExistingTest, which asserted the removed error-on-duplicate behavior. --- internal/datastore/memdb/readwrite.go | 51 +++++- internal/datastore/mysql/readwrite.go | 29 +++- internal/datastore/postgres/common/bulk.go | 176 +++++++++++++++------ internal/datastore/spanner/readwrite.go | 8 +- pkg/datastore/test/bulk.go | 140 ++++++++-------- pkg/datastore/test/datastore.go | 3 +- pkg/datastore/test/relationships.go | 11 +- 7 files changed, 270 insertions(+), 148 deletions(-) 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.