Merge pull request #2 from theplant/totalcount

refactor

README.md

@@ -17,23 +17,42 @@
 
 ```go
 p := relay.New(
-    true, // nodesOnly, true means only nodes are returned, otherwise only edges are returned
-    10, 10, // maxLimit / limitIfNotSet
     func(ctx context.Context, req *relay.ApplyCursorsRequest) (*relay.ApplyCursorsResponse[*User], error) {
         // Offset-based pagination
         // return gormrelay.NewOffsetAdapter[*User](db)(ctx, req)
+
         // Keyset-based pagination
         return gormrelay.NewKeysetAdapter[*User](db)(ctx, req)
     },
+    // maxLimit / limitIfNotSet
+    relay.EnsureLimits[*User](100, 10),
+    // Append primary sorting fields, if any are unspecified
+    relay.EnsurePrimaryOrderBy[*User](
+        relay.OrderBy{Field: "ID", Desc: false},
+        relay.OrderBy{Field: "Version", Desc: false},
+    ),
+)
+
+resp, err := p.Paginate(
+    // If you do not want to return edges
+    // relay.WithSkipEdges(context.Background()),
+
+    // If you do not want to return nodes
+    // relay.WithSkipNodes(context.Background()),
+
+    // If you want to skip the total count
+    relay.WithSkipTotalCount(context.Background()),
+
+    // Query first 10 records
+    &relay.PaginateRequest[*User]{
+        First: lo.ToPtr(10),
+    }
 )
-resp, err := p.Paginate(context.Background(), &relay.PaginateRequest[*User]{
-    First: lo.ToPtr(10), // query first 10 records
-})
 ```
 
-### Middleware
+### Cursor Encryption
 
-If you need to encrypt cursors, you can use `cursor.Base64` or `cursor.GCM` middlewares:
+If you need to encrypt cursors, you can use `cursor.Base64` or `cursor.GCM` wrappers:
 
 ```go
 // Encrypt cursors with Base64
@@ -45,67 +64,18 @@ require.NoError(t, err)
 cursor.GCM(gcm)(gormrelay.NewKeysetAdapter[*User](db))
 ```
 
-If you need to append `PrimaryOrderBys` to `PaginateRequest.OrderBys`
-
-```go
-// without middleware
-req.OrderBys = relay.AppendPrimaryOrderBy[*User](req.OrderBys, 
-    relay.OrderBy{Field: "ID", Desc: false},
-    relay.OrderBy{Field: "Version", Desc: false},
-)
-
-// use cursor middleware
-cursor.PrimaryOrderBy[*User](
-    relay.OrderBy{Field: "ID", Desc: true},
-    relay.OrderBy{Field: "Version", Desc: false},
-)(
-    gormrelay.NewKeysetAdapter[*User](db),
-)
-
-// use pagination middleware
-relay.PrimaryOrderBy[*User](
-    relay.OrderBy{Field: "ID", Desc: false},
-    relay.OrderBy{Field: "Version", Desc: false},
-)(p)
-```
-
-### Skipping `TotalCount` Query for Optimization
-
-To improve performance, you can skip querying `TotalCount`, especially useful for large datasets:
-
-```go
-// Keyset-based pagination without querying TotalCount
-// Note: The final PageInfo.TotalCount will be relay.InvalidTotalCount(-1)
-cursor.NewKeysetAdapter(gormrelay.NewKeysetFinder[any](db))
-
-// Offset-based pagination without querying TotalCount
-// Note: The final PageInfo.TotalCount will be relay.InvalidTotalCount(-1)
-// Note: Using `Last != nil && Before == nil` is not supported for this case.
-cursor.NewOffsetAdapter(gormrelay.NewOffsetFinder[any](db))
-
-// Compared to the version that queries TotalCount
-
-cursor.NewKeysetAdapter(gormrelay.NewKeysetCounter[any](db))
-// equals
-gormrelay.NewKeysetAdapter[any](db)
-
-cursor.NewOffsetAdapter(gormrelay.NewOffsetCounter[any](db))
-// equals
-gormrelay.NewOffsetAdapter[any](db)
-```
-
 ### Non-Generic Usage
 
 If you do not use generics, you can create a paginator with the `any` type and combine it with the `db.Model` method:
 
 ```go
 p := relay.New(
-    false, // nodesOnly
-    10, 10,
     func(ctx context.Context, req *relay.ApplyCursorsRequest) (*relay.ApplyCursorsResponse[any], error) {
         // Since this is a generic function (T: any), we must call db.Model(x)
         return gormrelay.NewKeysetAdapter[any](db.Model(&User{}))(ctx, req)
     },
+    relay.EnsureLimits[any](100, 10),
+    relay.EnsurePrimaryOrderBy[any](relay.OrderBy{Field: "ID", Desc: false}),
 )
 resp, err := p.Paginate(context.Background(), &relay.PaginateRequest[any]{
     First: lo.ToPtr(10), // query first 10 records

context.go

@@ -0,0 +1,36 @@
+package relay
+
+import "context"
+
+type ctxKeySkipTotalCount struct{}
+
+func WithSkipTotalCount(ctx context.Context) context.Context {
+	return context.WithValue(ctx, ctxKeySkipTotalCount{}, true)
+}
+
+func ShouldSkipTotalCount(ctx context.Context) bool {
+	b, _ := ctx.Value(ctxKeySkipTotalCount{}).(bool)
+	return b
+}
+
+type ctxKeySkipEdges struct{}
+
+func WithSkipEdges(ctx context.Context) context.Context {
+	return context.WithValue(ctx, ctxKeySkipEdges{}, true)
+}
+
+func ShouldSkipEdges(ctx context.Context) bool {
+	b, _ := ctx.Value(ctxKeySkipEdges{}).(bool)
+	return b
+}
+
+type ctxKeySkipNodes struct{}
+
+func WithSkipNodes(ctx context.Context) context.Context {
+	return context.WithValue(ctx, ctxKeySkipNodes{}, true)
+}
+
+func ShouldSkipNodes(ctx context.Context) bool {
+	b, _ := ctx.Value(ctxKeySkipNodes{}).(bool)
+	return b
+}

cursor/base64.go

@@ -33,8 +33,8 @@ func Base64[T any](next relay.ApplyCursorsFunc[T]) relay.ApplyCursorsFunc[T] {
 		}
 
 		// Encrypt the cursor
-		for i := range resp.Edges {
-			edge := &resp.Edges[i]
+		for i := range resp.LazyEdges {
+			edge := &resp.LazyEdges[i]
 			originalCursor := edge.Cursor
 			edge.Cursor = func(ctx context.Context, node T) (string, error) {
 				cursor, err := originalCursor(ctx, node)

cursor/gcm.go

@@ -82,8 +82,8 @@ func GCM[T any](gcm cipher.AEAD) relay.CursorMiddleware[T] {
 				return nil, err
 			}
 
-			for i := range resp.Edges {
-				edge := &resp.Edges[i]
+			for i := range resp.LazyEdges {
+				edge := &resp.LazyEdges[i]
 				originalCursor := edge.Cursor
 				edge.Cursor = func(ctx context.Context, node T) (string, error) {
 					cursor, err := originalCursor(ctx, node)

cursor/keyset.go

@@ -11,12 +11,7 @@ import (
 
 type KeysetFinder[T any] interface {
 	Find(ctx context.Context, after, before *map[string]any, orderBys []relay.OrderBy, limit int, fromLast bool) ([]T, error)
-}
-
-type KeysetFinderFunc[T any] func(ctx context.Context, after, before *map[string]any, orderBys []relay.OrderBy, limit int, fromLast bool) ([]T, error)
-
-func (f KeysetFinderFunc[T]) Find(ctx context.Context, after, before *map[string]any, orderBys []relay.OrderBy, limit int, fromLast bool) ([]T, error) {
-	return f(ctx, after, before, orderBys, limit, fromLast)
+	Count(ctx context.Context) (int, error)
 }
 
 func NewKeysetAdapter[T any](finder KeysetFinder[T]) relay.ApplyCursorsFunc[T] {
@@ -33,22 +28,21 @@ func NewKeysetAdapter[T any](finder KeysetFinder[T]) relay.ApplyCursorsFunc[T] {
 			return nil, err
 		}
 
-		totalCount := relay.InvalidTotalCount
-		counter, ok := finder.(Counter)
-		if ok {
-			var err error
-			totalCount, err = counter.Count(ctx)
+		var totalCount *int
+		if !relay.ShouldSkipTotalCount(ctx) {
+			count, err := finder.Count(ctx)
 			if err != nil {
 				return nil, err
 			}
+			totalCount = &count
 		}
 
 		cursorEncoder := func(_ context.Context, node T) (string, error) {
 			return EncodeKeysetCursor(node, keys)
 		}
 
 		var edges []relay.LazyEdge[T]
-		if req.Limit <= 0 || (counter != nil && totalCount <= 0) {
+		if req.Limit <= 0 || (totalCount != nil && *totalCount <= 0) {
 			edges = make([]relay.LazyEdge[T], 0)
 		} else {
 			nodes, err := finder.Find(ctx, after, before, req.OrderBys, req.Limit, req.FromLast)
@@ -65,7 +59,7 @@ func NewKeysetAdapter[T any](finder KeysetFinder[T]) relay.ApplyCursorsFunc[T] {
 		}
 
 		resp := &relay.ApplyCursorsResponse[T]{
-			Edges:      edges,
+			LazyEdges:  edges,
 			TotalCount: totalCount,
 			// It would be very costly to check whether after and before really exist,
 			// So it is usually not worth it. Normally, checking that it is not nil is sufficient.

cursor/offset.go

@@ -10,6 +10,7 @@ import (
 
 type OffsetFinder[T any] interface {
 	Find(ctx context.Context, orderBys []relay.OrderBy, skip, limit int) ([]T, error)
+	Count(ctx context.Context) (int, error)
 }
 
 type OffsetFinderFunc[T any] func(ctx context.Context, orderBys []relay.OrderBy, skip, limit int) ([]T, error)
@@ -19,29 +20,28 @@ func (f OffsetFinderFunc[T]) Find(ctx context.Context, orderBys []relay.OrderBy,
 }
 
 // NewOffsetAdapter creates a relay.ApplyCursorsFunc from an OffsetFinder.
-// If you want to use `last!=nil&&before==nil`, the finder must implement Counter.
+// If you want to use `last!=nil&&before==nil`, you can't skip totalCount.
 func NewOffsetAdapter[T any](finder OffsetFinder[T]) relay.ApplyCursorsFunc[T] {
 	return func(ctx context.Context, req *relay.ApplyCursorsRequest) (*relay.ApplyCursorsResponse[T], error) {
 		after, before, err := decodeOffsetCursors(req.After, req.Before)
 		if err != nil {
 			return nil, err
 		}
 
-		totalCount := relay.InvalidTotalCount
-		counter, hasCounter := finder.(Counter)
-		if hasCounter {
-			var err error
-			totalCount, err = counter.Count(ctx)
+		var totalCount *int
+		if !relay.ShouldSkipTotalCount(ctx) {
+			count, err := finder.Count(ctx)
 			if err != nil {
 				return nil, err
 			}
+			totalCount = &count
 		}
 
 		if req.FromLast && before == nil {
-			if !hasCounter {
-				return nil, errors.New("counter is required for fromLast and nil before")
+			if totalCount == nil {
+				return nil, errors.New("totalCount is required for fromLast and nil before")
 			}
-			before = &totalCount
+			before = totalCount
 		}
 
 		limit, skip := req.Limit, 0
@@ -67,7 +67,7 @@ func NewOffsetAdapter[T any](finder OffsetFinder[T]) relay.ApplyCursorsFunc[T] {
 		}
 
 		var edges []relay.LazyEdge[T]
-		if limit <= 0 || (hasCounter && (skip >= totalCount || totalCount <= 0)) {
+		if limit <= 0 || (totalCount != nil && (skip >= *totalCount || *totalCount <= 0)) {
 			edges = make([]relay.LazyEdge[T], 0)
 		} else {
 			nodes, err := finder.Find(ctx, req.OrderBys, skip, limit)
@@ -87,15 +87,15 @@ func NewOffsetAdapter[T any](finder OffsetFinder[T]) relay.ApplyCursorsFunc[T] {
 		}
 
 		resp := &relay.ApplyCursorsResponse[T]{
-			Edges:      edges,
+			LazyEdges:  edges,
 			TotalCount: totalCount,
 		}
 
-		if hasCounter {
-			resp.HasAfterOrPrevious = after != nil && *after < totalCount
-			resp.HasBeforeOrNext = before != nil && *before < totalCount
+		if totalCount != nil {
+			resp.HasAfterOrPrevious = after != nil && *after < *totalCount
+			resp.HasBeforeOrNext = before != nil && *before < *totalCount
 		} else {
-			// If we don't have a counter, it would be very costly to check whether after and before really exist,
+			// If we don't have totalCount, it would be very costly to check whether after and before really exist,
 			// So it is usually not worth it. Normally, checking that it is not nil is sufficient.
 			resp.HasAfterOrPrevious = after != nil
 			resp.HasBeforeOrNext = before != nil