Adding paging method to query service store iterator

pkg/query/internal/sqliterator/sqliterator.go

@@ -65,6 +65,38 @@ func (i *iterator) All() ([]models.Object, error) {
 	return objects, nil
 }
 
+func (i *iterator) Page(count int, offset int) ([]models.Object, error) {
+	var objects []models.Object
+
+	index := -1
+
+	for i.rows.Next() {
+		index++
+
+		if index < offset {
+			continue
+		}
+
+		var object models.Object
+
+		if err := i.result.ScanRows(i.rows, &object); err != nil {
+			return nil, fmt.Errorf("failed to scan rows: %w", err)
+		}
+
+		objects = append(objects, object)
+
+		if index >= offset+count-1 {
+			break
+		}
+	}
+
+	if err := i.rows.Err(); err != nil {
+		return nil, fmt.Errorf("failed to get rows: %w", err)
+	}
+
+	return objects, nil
+}
+
 func (i *iterator) Close() error {
 	return i.rows.Close()
 }

pkg/query/store/indexer.go

@@ -328,10 +328,6 @@ func (i *indexerIterator) Row() (models.Object, error) {
 	return i.s.GetObjectByID(context.Background(), id)
 }
 
-func (i *indexerIterator) Close() error {
-	return nil
-}
-
 func (i *indexerIterator) All() ([]models.Object, error) {
 	i.mu.Lock()
 	defer i.mu.Unlock()
@@ -343,10 +339,41 @@ func (i *indexerIterator) All() ([]models.Object, error) {
 	}
 
 	iter, err := i.s.GetObjects(context.Background(), ids, i.opts)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get objects: %w", err)
+	}
+
+	return iter.All()
+}
+
+func (i *indexerIterator) Page(count int, offset int) ([]models.Object, error) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
 
+	ids := []string{}
+
+	numHits := i.result.Hits.Len()
+	upper := offset + count
+	if upper > numHits {
+		upper = numHits
+	}
+
+	for index := offset; index < upper; index++ {
+		ids = append(ids, i.result.Hits[index].ID)
+	}
+
+	if len(ids) == 0 {
+		return []models.Object{}, nil
+	}
+
+	iter, err := i.s.GetObjects(context.Background(), ids, i.opts)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get objects: %w", err)
 	}
 
 	return iter.All()
 }
+
+func (i *indexerIterator) Close() error {
+	return nil
+}

pkg/query/store/indexer_test.go

@@ -126,7 +126,6 @@ func TestIndexer_Metrics(t *testing.T) {
 		}
 		assertMetrics(g, metricsUrl, wantMetrics)
 	})
-
 }
 
 func TestIndexer_RemoveByQuery(t *testing.T) {
@@ -213,6 +212,161 @@ func TestIndexer_RemoveByQuery(t *testing.T) {
 	}
 }
 
+func TestIndexer_RemoveByQueryWithPagination(t *testing.T) {
+	g := NewWithT(t)
+	tests := []struct {
+		name     string
+		query    string
+		objects  []models.Object
+		expected []string
+	}{
+		{
+			name:  "removes by cluster",
+			query: "+cluster:management",
+			objects: []models.Object{
+				{
+					Cluster:    "management",
+					Kind:       "Namespace",
+					Name:       "name-1",
+					APIGroup:   "anyGroup",
+					APIVersion: "anyVersion",
+					Category:   "automation",
+					Namespace:  "anyNamespace",
+				},
+				{
+					Cluster:    "othercluster",
+					Kind:       "Namespace",
+					Name:       "name-2",
+					APIGroup:   "anyGroup",
+					APIVersion: "anyVersion",
+					Category:   "automation",
+					Namespace:  "anyNamespace",
+				},
+				{
+					Cluster:    "management",
+					Kind:       "Namespace",
+					Name:       "name-3",
+					APIGroup:   "anyGroup",
+					APIVersion: "anyVersion",
+					Category:   "automation",
+					Namespace:  "anyNamespace",
+				},
+				{
+					Cluster:    "management",
+					Kind:       "Namespace",
+					Name:       "name-4",
+					APIGroup:   "anyGroup",
+					APIVersion: "anyVersion",
+					Category:   "automation",
+					Namespace:  "anyNamespace",
+				},
+				{
+					Cluster:    "othercluster",
+					Kind:       "Namespace",
+					Name:       "name-5",
+					APIGroup:   "anyGroup",
+					APIVersion: "anyVersion",
+					Category:   "automation",
+					Namespace:  "anyNamespace",
+				},
+				{
+					Cluster:    "management",
+					Kind:       "Namespace",
+					Name:       "name-6",
+					APIGroup:   "anyGroup",
+					APIVersion: "anyVersion",
+					Category:   "automation",
+					Namespace:  "anyNamespace",
+				},
+				{
+					Cluster:    "othercluster",
+					Kind:       "Namespace",
+					Name:       "name-7",
+					APIGroup:   "anyGroup",
+					APIVersion: "anyVersion",
+					Category:   "automation",
+					Namespace:  "anyNamespace",
+				},
+			},
+			expected: []string{"name-2", "name-5", "name-7"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			idxFileLocation := filepath.Join(t.TempDir(), indexFile)
+			mapping := bleve.NewIndexMapping()
+
+			index, err := bleve.New(idxFileLocation, mapping)
+			g.Expect(err).NotTo(HaveOccurred())
+
+			s, err := NewStore(StorageBackendSQLite, t.TempDir(), logr.Discard())
+			g.Expect(err).NotTo(HaveOccurred())
+
+			idx := &bleveIndexer{
+				idx:   index,
+				store: s,
+			}
+
+			err = idx.Add(context.Background(), tt.objects)
+			g.Expect(err).NotTo(HaveOccurred())
+
+			err = s.StoreObjects(context.Background(), tt.objects)
+			g.Expect(err).NotTo(HaveOccurred())
+
+			// Ensure things got written to the index.
+			// Iterate through all pages without an initial offset.
+			iter, err := idx.Search(context.Background(), query{}, nil)
+			g.Expect(err).NotTo(HaveOccurred())
+
+			pageObjects, err := iter.Page(3, 0)
+			g.Expect(err).NotTo(HaveOccurred())
+			g.Expect(len(pageObjects)).To(Equal(3))
+
+			pageObjects, err = iter.Page(3, 3)
+			g.Expect(err).NotTo(HaveOccurred())
+			g.Expect(len(pageObjects)).To(Equal(3))
+
+			pageObjects, err = iter.Page(3, 6)
+			g.Expect(err).NotTo(HaveOccurred())
+			g.Expect(len(pageObjects)).To(Equal(1))
+
+			// Iterate through all pages with an initial offset.
+			iter, err = idx.Search(context.Background(), query{}, nil)
+			g.Expect(err).NotTo(HaveOccurred())
+
+			pageObjects, err = iter.Page(2, 2)
+			g.Expect(err).NotTo(HaveOccurred())
+			g.Expect(len(pageObjects)).To(Equal(2))
+
+			pageObjects, err = iter.Page(2, 4)
+			g.Expect(err).NotTo(HaveOccurred())
+			g.Expect(len(pageObjects)).To(Equal(2))
+
+			pageObjects, err = iter.Page(2, 6)
+			g.Expect(err).NotTo(HaveOccurred())
+			g.Expect(len(pageObjects)).To(Equal(1))
+
+			// Check that required objects were removed.
+			err = idx.RemoveByQuery(context.Background(), tt.query)
+			g.Expect(err).NotTo(HaveOccurred())
+
+			iter, err = idx.Search(context.Background(), query{}, nil)
+			g.Expect(err).NotTo(HaveOccurred())
+
+			all, err := iter.All()
+			g.Expect(err).NotTo(HaveOccurred())
+
+			names := []string{}
+			for _, obj := range all {
+				names = append(names, obj.Name)
+			}
+
+			g.Expect(names).To(Equal(tt.expected))
+
+		})
+	}
+}
+
 type query struct{}
 
 func (q query) GetTerms() string {

pkg/query/store/store.go

@@ -87,6 +87,8 @@ type Iterator interface {
 	Row() (models.Object, error)
 	// All returns all rows of the iterator
 	All() ([]models.Object, error)
+	// Page returns a specified number of rows of the iterator with a specified offset
+	Page(count int, offset int) ([]models.Object, error)
 
 	// Close closes the iterator
 	Close() error

pkg/query/store/store_test.go

@@ -40,11 +40,109 @@ func TestGetObjects(t *testing.T) {
 
 	g.Expect(len(objects) > 0).To(BeTrue())
 	g.Expect(objects[0].Name).To(Equal(obj.Name))
+}
+
+func TestGetObjectsWithPagination(t *testing.T) {
+	g := NewGomegaWithT(t)
+
+	store, db := createStore(t)
+
+	testObjects := []models.Object{
+		{
+			Cluster:   "test-cluster",
+			Name:      "someName-1",
+			Namespace: "namespace",
+			Kind:      "ValidKind",
+		},
+		{
+			Cluster:   "test-cluster",
+			Name:      "someName-2",
+			Namespace: "namespace",
+			Kind:      "ValidKind",
+		},
+		{
+			Cluster:   "test-cluster",
+			Name:      "someName-3",
+			Namespace: "namespace",
+			Kind:      "ValidKind",
+		},
+		{
+			Cluster:   "test-cluster",
+			Name:      "someName-4",
+			Namespace: "namespace",
+			Kind:      "ValidKind",
+		},
+		{
+			Cluster:   "test-cluster",
+			Name:      "someName-5",
+			Namespace: "namespace",
+			Kind:      "ValidKind",
+		},
+		{
+			Cluster:   "test-cluster",
+			Name:      "someName-6",
+			Namespace: "namespace",
+			Kind:      "ValidKind",
+		},
+		{
+			Cluster:   "test-cluster",
+			Name:      "someName-7",
+			Namespace: "namespace",
+			Kind:      "ValidKind",
+		},
+	}
+
+	g.Expect(SeedObjects(db, testObjects)).To(Succeed())
 
+	iter, err := store.GetObjects(context.Background(), nil, nil)
+	g.Expect(err).To(BeNil())
+
+	objects, err := iter.All()
+	g.Expect(err).To(BeNil())
+	g.Expect(len(objects)).To(Equal(len(testObjects)))
+	g.Expect(objects[0].Name).To(Equal(testObjects[0].Name))
+	g.Expect(objects[3].Name).To(Equal(testObjects[3].Name))
+
+	// With pagination and without an offset.
+	iter, err = store.GetObjects(context.Background(), nil, nil)
+	g.Expect(err).To(BeNil())
+
+	objects, err = iter.Page(3, 0)
+	g.Expect(err).To(BeNil())
+	g.Expect(len(objects)).To(Equal(3))
+	g.Expect(objects[0].Name).To(Equal(testObjects[0].Name))
+
+	objects, err = iter.Page(3, 0)
+	g.Expect(err).To(BeNil())
+	g.Expect(len(objects)).To(Equal(3))
+	g.Expect(objects[0].Name).To(Equal(testObjects[3].Name))
+
+	objects, err = iter.Page(3, 0)
+	g.Expect(err).To(BeNil())
+	g.Expect(len(objects)).To(Equal(1))
+	g.Expect(objects[0].Name).To(Equal(testObjects[6].Name))
+
+	// With pagination and with an initial offset.
+	iter, err = store.GetObjects(context.Background(), nil, nil)
+	g.Expect(err).To(BeNil())
+
+	objects, err = iter.Page(2, 2)
+	g.Expect(err).To(BeNil())
+	g.Expect(len(objects)).To(Equal(2))
+	g.Expect(objects[0].Name).To(Equal(testObjects[2].Name))
+
+	objects, err = iter.Page(2, 0)
+	g.Expect(err).To(BeNil())
+	g.Expect(len(objects)).To(Equal(2))
+	g.Expect(objects[0].Name).To(Equal(testObjects[4].Name))
+
+	objects, err = iter.Page(2, 0)
+	g.Expect(err).To(BeNil())
+	g.Expect(len(objects)).To(Equal(1))
+	g.Expect(objects[0].Name).To(Equal(testObjects[6].Name))
 }
 
 func TestDeleteObjects(t *testing.T) {
-
 	tests := []struct {
 		name     string
 		seed     []models.Object