Respect ydb select limit in bulk load test

Signed-off-by: sashayakovtseva <[email protected]>

pkg/datastore/test/bulk.go

@@ -6,14 +6,18 @@ import (
 	"strconv"
 	"testing"
 
+	"github.com/samber/lo"
 	"github.com/stretchr/testify/require"
 
 	"github.com/authzed/spicedb/internal/testfixtures"
 	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/datastore/options"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 )
 
 func BulkUploadTest(t *testing.T, tester DatastoreTester) {
+	const ydbSelectLimit = 1000
+
 	testCases := []int{0, 1, 10, 100, 1_000, 10_000}
 
 	for _, tc := range testCases {
@@ -47,13 +51,39 @@ func BulkUploadTest(t *testing.T, tester DatastoreTester) {
 			head, err := ds.HeadRevision(ctx)
 			require.NoError(err)
 
-			iter, err := ds.SnapshotReader(head).QueryRelationships(ctx, datastore.RelationshipsFilter{
-				ResourceType: testfixtures.DocumentNS.Name,
-			})
-			require.NoError(err)
-			defer iter.Close()
+			var (
+				after       *core.RelationTuple
+				isLastCheck bool
+			)
+			for left := tc; !isLastCheck; {
+				if left == 0 {
+					isLastCheck = true
+				}
+
+				iter, err := ds.SnapshotReader(head).QueryRelationships(ctx, datastore.RelationshipsFilter{
+					ResourceType: testfixtures.DocumentNS.Name,
+				},
+					options.WithLimit(lo.ToPtr(uint64(ydbSelectLimit))),
+					options.WithSort(options.ByResource),
+					options.WithAfter(after),
+				)
+				require.NoError(err)
+
+				expect := ydbSelectLimit
+				if left < ydbSelectLimit {
+					expect = left
+				}
+
+				tRequire.VerifyIteratorCount(iter, expect)
+
+				if expect > 0 {
+					after, err = iter.Cursor()
+					require.NoError(err)
+				}
 
-			tRequire.VerifyIteratorCount(iter, tc)
+				iter.Close()
+				left -= expect
+			}
 		})
 	}
 }