Allow fetching the count from GetPage

This change adds an optional count argument to GetPage. To make it optional it's implemented as a variadic function but only one count pointer is accepted if set.

pika.go

@@ -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_page_token.go

@@ -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

@@ -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_test.go

@@ -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/test.proto

@@ -32,3 +32,16 @@ message Complete3 {
   Status status = 6;
   repeated string strs = 7;
 }
+
+message TestRequest1 {
+  int32 page_size = 1;
+  string page_token = 2;
+  string filter = 3;
+  string order_by = 4;
+}
+
+message SimpleModel1 {
+  int64 id = 1;
+  string title = 2;
+  string description = 3;
+}