@@ -24,9 +24,11 @@ package dsx
2424import (
2525 "context"
2626 "errors"
27+ "fmt"
2728 "log"
2829
2930 "cloud.google.com/go/datastore"
31+ "cloud.google.com/go/datastore/apiv1/datastorepb"
3032 "google.golang.org/api/iterator"
3133 "google.golang.org/api/option"
3234)
@@ -363,6 +365,34 @@ func (qb *QueryBuilder[T]) KeysOnly() *QueryBuilder[T] {
363365 return qb
364366}
365367
368+ // Count returns the count of entities matching the current query filters.
369+ // It uses Datastore's aggregation query to efficiently count without loading entities into memory.
370+ //
371+ // Example:
372+ //
373+ // count, err := dsx.From[User](ctx, db).
374+ // Where("Status", "=", "active").
375+ // Count()
376+ //
377+ // Returns 0 and an error if the aggregation query fails or the count result is missing.
378+ // Note: Datastore count aggregations have a default limit of approximately 1 million entities.
379+ func (qb * QueryBuilder [T ]) Count () (int64 , error ) {
380+ aggQuery := qb .query .NewAggregationQuery ().WithCount ("total" )
381+ results , err := qb .db .client .RunAggregationQuery (qb .context , aggQuery )
382+ if err != nil {
383+ return 0 , err
384+ }
385+ count , ok := results ["total" ]
386+ if ! ok {
387+ return 0 , errors .New ("count result not found" )
388+ }
389+ val , ok := count .(* datastorepb.Value )
390+ if ! ok {
391+ return 0 , fmt .Errorf ("unexpected count type: %T" , count )
392+ }
393+ return val .GetIntegerValue (), nil
394+ }
395+
366396// SelectWithCursor executes the query and returns results with a cursor
367397// for pagination. The cursor can be passed to WithCursor in subsequent
368398// queries to fetch the next page.
0 commit comments