From a6f340b5dce2fa68819cb35500485a0bc6a6110a Mon Sep 17 00:00:00 2001
From: sashayakovtseva <sashayakovtseva@gmail.com>
Date: Thu, 22 Feb 2024 20:57:35 +0800
Subject: [PATCH] Respect ydb select limit in bulk load test

Signed-off-by: sashayakovtseva <sashayakovtseva@gmail.com>
---
 pkg/datastore/test/bulk.go | 42 ++++++++++++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/pkg/datastore/test/bulk.go b/pkg/datastore/test/bulk.go
index 5a0b886d1d..8c15833e98 100644
--- a/pkg/datastore/test/bulk.go
+++ b/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
+			}
 		})
 	}
 }