github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/datastore/test/bulk.go (about)

     1  package test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"strconv"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  	"google.golang.org/grpc/codes"
    11  
    12  	"github.com/authzed/grpcutil"
    13  
    14  	"github.com/authzed/spicedb/internal/testfixtures"
    15  	"github.com/authzed/spicedb/pkg/datastore"
    16  	"github.com/authzed/spicedb/pkg/datastore/options"
    17  	core "github.com/authzed/spicedb/pkg/proto/core/v1"
    18  )
    19  
    20  func BulkUploadTest(t *testing.T, tester DatastoreTester) {
    21  	testCases := []int{0, 1, 10, 100, 1_000, 10_000}
    22  
    23  	for _, tc := range testCases {
    24  		t.Run(strconv.Itoa(tc), func(t *testing.T) {
    25  			require := require.New(t)
    26  			ctx := context.Background()
    27  
    28  			rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
    29  			require.NoError(err)
    30  
    31  			ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
    32  			bulkSource := testfixtures.NewBulkTupleGenerator(
    33  				testfixtures.DocumentNS.Name,
    34  				"viewer",
    35  				testfixtures.UserNS.Name,
    36  				tc,
    37  				t,
    38  			)
    39  
    40  			_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
    41  				loaded, err := rwt.BulkLoad(ctx, bulkSource)
    42  				require.NoError(err)
    43  				require.Equal(uint64(tc), loaded)
    44  				return err
    45  			})
    46  			require.NoError(err)
    47  
    48  			tRequire := testfixtures.TupleChecker{Require: require, DS: ds}
    49  
    50  			head, err := ds.HeadRevision(ctx)
    51  			require.NoError(err)
    52  
    53  			iter, err := ds.SnapshotReader(head).QueryRelationships(ctx, datastore.RelationshipsFilter{
    54  				OptionalResourceType: testfixtures.DocumentNS.Name,
    55  			})
    56  			require.NoError(err)
    57  			defer iter.Close()
    58  
    59  			tRequire.VerifyIteratorCount(iter, tc)
    60  		})
    61  	}
    62  }
    63  
    64  func BulkUploadErrorsTest(t *testing.T, tester DatastoreTester) {
    65  	require := require.New(t)
    66  	ctx := context.Background()
    67  
    68  	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
    69  	require.NoError(err)
    70  
    71  	ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
    72  
    73  	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
    74  		inserted, err := rwt.BulkLoad(ctx, &onlyErrorSource{})
    75  
    76  		// We can't check the specific error because pgx is not wrapping
    77  		require.Error(err)
    78  		require.Zero(inserted)
    79  		return err
    80  	})
    81  	require.Error(err)
    82  }
    83  
    84  func BulkUploadAlreadyExistsSameCallErrorTest(t *testing.T, tester DatastoreTester) {
    85  	require := require.New(t)
    86  	ctx := context.Background()
    87  
    88  	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
    89  	require.NoError(err)
    90  
    91  	ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
    92  
    93  	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
    94  		inserted, err := rwt.BulkLoad(ctx, testfixtures.NewBulkTupleGenerator(
    95  			testfixtures.DocumentNS.Name,
    96  			"viewer",
    97  			testfixtures.UserNS.Name,
    98  			1,
    99  			t,
   100  		))
   101  		require.NoError(err)
   102  		require.Equal(uint64(1), inserted)
   103  
   104  		_, serr := rwt.BulkLoad(ctx, testfixtures.NewBulkTupleGenerator(
   105  			testfixtures.DocumentNS.Name,
   106  			"viewer",
   107  			testfixtures.UserNS.Name,
   108  			1,
   109  			t,
   110  		))
   111  		return serr
   112  	}, options.WithDisableRetries(true))
   113  
   114  	// NOTE: spanner does not return an error for duplicates.
   115  	if err == nil {
   116  		return
   117  	}
   118  
   119  	grpcutil.RequireStatus(t, codes.AlreadyExists, err)
   120  }
   121  
   122  func BulkUploadAlreadyExistsErrorTest(t *testing.T, tester DatastoreTester) {
   123  	require := require.New(t)
   124  	ctx := context.Background()
   125  
   126  	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
   127  	require.NoError(err)
   128  
   129  	ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
   130  
   131  	// Bulk write a single relationship.
   132  	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
   133  		inserted, err := rwt.BulkLoad(ctx, testfixtures.NewBulkTupleGenerator(
   134  			testfixtures.DocumentNS.Name,
   135  			"viewer",
   136  			testfixtures.UserNS.Name,
   137  			1,
   138  			t,
   139  		))
   140  		require.NoError(err)
   141  		require.Equal(uint64(1), inserted)
   142  		return nil
   143  	}, options.WithDisableRetries(true))
   144  	require.NoError(err)
   145  
   146  	// Bulk write it again and ensure we get the expected error.
   147  	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
   148  		_, serr := rwt.BulkLoad(ctx, testfixtures.NewBulkTupleGenerator(
   149  			testfixtures.DocumentNS.Name,
   150  			"viewer",
   151  			testfixtures.UserNS.Name,
   152  			1,
   153  			t,
   154  		))
   155  		return serr
   156  	}, options.WithDisableRetries(true))
   157  
   158  	// NOTE: spanner does not return an error for duplicates.
   159  	if err == nil {
   160  		return
   161  	}
   162  
   163  	grpcutil.RequireStatus(t, codes.AlreadyExists, err)
   164  }
   165  
   166  type onlyErrorSource struct{}
   167  
   168  var errOnlyError = errors.New("source iterator error")
   169  
   170  func (oes onlyErrorSource) Next(_ context.Context) (*core.RelationTuple, error) {
   171  	return nil, errOnlyError
   172  }
   173  
   174  var _ datastore.BulkWriteRelationshipSource = onlyErrorSource{}