Merge pull request #9 from mstg/page-count

Allow fetching the count from GetPage

parser/gen.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package parser

pika.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika
@@ -155,7 +155,9 @@ type QuerySet[T any] interface {
 	AIP160(filter string, options AIPFilterOptions) (QuerySet[T], error)
 
 	// Page token functionality for gRPC
-	GetPage(paginatable Paginatable, options AIPFilterOptions) ([]*T, string, error)
+	// The count is optional and returns the total number of rows for the query.
+	// It is implemented as a variadic function to not break existing code.
+	GetPage(paginatable Paginatable, options AIPFilterOptions, count ...*int) ([]*T, string, error)
 
 	// Join table
 	InnerJoin(modelFirst, modelSecond interface{}, keyFirst, keySecond string) QuerySet[T]

pika_aip_filter.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika

pika_aip_filter_proto.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika

pika_aip_filter_proto_test.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika

pika_aip_filter_psql_test.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika

pika_page_token.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika
@@ -57,7 +57,7 @@ func (p *PageToken[T]) Encode() (string, error) {
 	return base64.URLEncoding.EncodeToString(data), nil
 }
 
-// ParsePageToken constructs a PageToken from a base64-encoded string
+// Decode constructs a PageToken from a base64-encoded string
 func (p *PageToken[T]) Decode(s string) error {
 	data, err := base64.URLEncoding.DecodeString(s)
 	if err != nil {

pika_psql.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika
@@ -542,7 +542,11 @@ func (b *basePsql[T]) AIP160(filter string, options AIPFilterOptions) (QuerySet[
 }
 
 // Page tokens for gRPC
-func (b *basePsql[T]) GetPage(paginatable Paginatable, options AIPFilterOptions) ([]*T, string, error) {
+func (b *basePsql[T]) GetPage(paginatable Paginatable, options AIPFilterOptions, countPointer ...*int) ([]*T, string, error) {
+	if len(countPointer) > 1 {
+		return nil, "", fmt.Errorf("too many arguments (count should be one pointer or none)")
+	}
+
 	if b.err != nil {
 		return nil, "", b.err
 	}
@@ -579,6 +583,10 @@ func (b *basePsql[T]) GetPage(paginatable Paginatable, options AIPFilterOptions)
 		return nil, "", fmt.Errorf("getting count: %w", err)
 	}
 
+	if len(countPointer) > 0 {
+		*countPointer[0] = count
+	}
+
 	// If no more results after this page, return empty page token
 	if b.PageToken.Offset >= uint(count) {
 		return result, "", nil

pika_psql_experimental.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika

pika_psql_test.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 package pika
@@ -16,6 +16,7 @@ import (
 	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
 	orderedmap "github.com/wk8/go-ordered-map/v2"
+	pikatestpb "go.ciq.dev/pika/testproto"
 )
 
 var pgInstance *embeddedpostgres.EmbeddedPostgres
@@ -1735,3 +1736,21 @@ func TestFilterMissingArgs(t *testing.T) {
 	_, err = qs.Count()
 	requireError(err)
 }
+
+func TestGetPageCount(t *testing.T) {
+	psql := newPsql(t)
+	createTestEntries(t, psql)
+	qs := Q[simpleModel1](psql)
+
+	aipOptions := ProtoReflect(&pikatestpb.SimpleModel1{})
+	req := &pikatestpb.TestRequest1{
+		PageSize: int32(1),
+	}
+	var count int
+	page, nt, err := qs.GetPage(req, aipOptions, &count)
+	require.Nil(t, err)
+
+	require.Equal(t, 1, len(page))
+	require.NotEmpty(t, nt)
+	require.Equal(t, 3, count)
+}

testproto/gen.go

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2023, Ctrl IQ, Inc. All rights reserved
+// SPDX-FileCopyrightText: Copyright (c) 2023-2024, Ctrl IQ, Inc. All rights reserved
 // SPDX-License-Identifier: Apache-2.0
 
 //go:generate protoc --go_opt=paths=source_relative --go_out=. test.proto