feat(datastore): idempotent (TOUCH-like) bulk relationship import

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++
 	}

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)

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,75 +13,145 @@ 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,
 	tupleTableName string,
 	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 <table> (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
 }

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)
 		}