github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/postgres/common/bulk.go (about)

     1  package common
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  
     7  	"github.com/jackc/pgx/v5"
     8  
     9  	"github.com/authzed/spicedb/pkg/datastore"
    10  	core "github.com/authzed/spicedb/pkg/proto/core/v1"
    11  )
    12  
    13  type tupleSourceAdapter struct {
    14  	source datastore.BulkWriteRelationshipSource
    15  	ctx    context.Context
    16  
    17  	current      *core.RelationTuple
    18  	err          error
    19  	valuesBuffer []any
    20  }
    21  
    22  // Next returns true if there is another row and makes the next row data
    23  // available to Values(). When there are no more rows available or an error
    24  // has occurred it returns false.
    25  func (tg *tupleSourceAdapter) Next() bool {
    26  	tg.current, tg.err = tg.source.Next(tg.ctx)
    27  	return tg.current != nil
    28  }
    29  
    30  // Values returns the values for the current row.
    31  func (tg *tupleSourceAdapter) Values() ([]any, error) {
    32  	var caveatName sql.NullString
    33  	var caveatContext map[string]any
    34  	if tg.current.Caveat != nil {
    35  		caveatName.String = tg.current.Caveat.CaveatName
    36  		caveatName.Valid = true
    37  		caveatContext = tg.current.Caveat.Context.AsMap()
    38  	}
    39  
    40  	tg.valuesBuffer[0] = tg.current.ResourceAndRelation.Namespace
    41  	tg.valuesBuffer[1] = tg.current.ResourceAndRelation.ObjectId
    42  	tg.valuesBuffer[2] = tg.current.ResourceAndRelation.Relation
    43  	tg.valuesBuffer[3] = tg.current.Subject.Namespace
    44  	tg.valuesBuffer[4] = tg.current.Subject.ObjectId
    45  	tg.valuesBuffer[5] = tg.current.Subject.Relation
    46  	tg.valuesBuffer[6] = caveatName
    47  	tg.valuesBuffer[7] = caveatContext
    48  
    49  	return tg.valuesBuffer, nil
    50  }
    51  
    52  // Err returns any error that has been encountered by the CopyFromSource. If
    53  // this is not nil *Conn.CopyFrom will abort the copy.
    54  func (tg *tupleSourceAdapter) Err() error {
    55  	return tg.err
    56  }
    57  
    58  func BulkLoad(
    59  	ctx context.Context,
    60  	tx pgx.Tx,
    61  	tupleTableName string,
    62  	colNames []string,
    63  	iter datastore.BulkWriteRelationshipSource,
    64  ) (uint64, error) {
    65  	adapter := &tupleSourceAdapter{
    66  		source:       iter,
    67  		ctx:          ctx,
    68  		valuesBuffer: make([]any, 8),
    69  	}
    70  	copied, err := tx.CopyFrom(ctx, pgx.Identifier{tupleTableName}, colNames, adapter)
    71  	return uint64(copied), err
    72  }