cursor pagination

pkg/mongo/cursor.go

@@ -0,0 +1,234 @@
+package mongo
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"strconv"
+
+	"go.mongodb.org/mongo-driver/v2/bson"
+)
+
+const pageCursorVersion = 1
+
+// pageCursor is an opaque pagination token for keyset (cursor) pagination.
+type pageCursor struct {
+	V         int    `json:"v"`
+	SortField string `json:"sort_field"`
+	SortOrder int    `json:"sort_order"`
+	ValueType string `json:"value_type"` // i64, f64, str, oid
+	SortValue string `json:"sort_value"`
+	ID        string `json:"id"`
+}
+
+func encodePageCursor(sortField string, sortOrder int, sortValue interface{}, id string) (string, error) {
+	if sortField == "" {
+		return "", fmt.Errorf("%w: sort field is required", ErrInvalidCursor)
+	}
+	if id == "" {
+		return "", fmt.Errorf("%w: document id is required", ErrInvalidCursor)
+	}
+
+	valueType, sortValueStr, err := formatCursorSortValue(sortField, sortValue)
+	if err != nil {
+		return "", err
+	}
+
+	payload, err := json.Marshal(pageCursor{
+		V:         pageCursorVersion,
+		SortField: sortField,
+		SortOrder: sortOrder,
+		ValueType: valueType,
+		SortValue: sortValueStr,
+		ID:        id,
+	})
+	if err != nil {
+		return "", fmt.Errorf("failed to marshal cursor: %w", err)
+	}
+
+	return base64.StdEncoding.EncodeToString(payload), nil
+}
+
+func decodePageCursor(encoded string) (*pageCursor, error) {
+	cursorBytes, err := base64.StdEncoding.DecodeString(encoded)
+	if err != nil {
+		return nil, fmt.Errorf("%w: invalid cursor encoding", ErrInvalidCursor)
+	}
+
+	var cursor pageCursor
+	if err := json.Unmarshal(cursorBytes, &cursor); err != nil {
+		return nil, fmt.Errorf("%w: invalid cursor payload", ErrInvalidCursor)
+	}
+	if cursor.V != pageCursorVersion || cursor.SortField == "" || cursor.ID == "" {
+		return nil, fmt.Errorf("%w: unsupported or incomplete cursor", ErrInvalidCursor)
+	}
+	if cursor.SortOrder != 1 && cursor.SortOrder != -1 {
+		return nil, fmt.Errorf("%w: invalid sort order in cursor", ErrInvalidCursor)
+	}
+
+	return &cursor, nil
+}
+
+func (c *pageCursor) matches(sortField string, sortOrder int) bool {
+	return c.SortField == sortField && c.SortOrder == sortOrder
+}
+
+func (c *pageCursor) mongoFilter() (bson.M, error) {
+	if c.SortField == "_id" {
+		objectID, err := bson.ObjectIDFromHex(c.ID)
+		if err != nil {
+			return nil, fmt.Errorf("%w: invalid ObjectID in cursor", ErrInvalidCursor)
+		}
+		operator := "$lt"
+		if c.SortOrder == 1 {
+			operator = "$gt"
+		}
+		return bson.M{"_id": bson.M{operator: objectID}}, nil
+	}
+
+	sortValue, err := c.parseSortValue()
+	if err != nil {
+		return nil, err
+	}
+
+	objectID, err := bson.ObjectIDFromHex(c.ID)
+	if err != nil {
+		return nil, fmt.Errorf("%w: invalid ObjectID in cursor", ErrInvalidCursor)
+	}
+
+	primaryOp := "$lt"
+	secondaryOp := "$lt"
+	if c.SortOrder == 1 {
+		primaryOp = "$gt"
+		secondaryOp = "$gt"
+	}
+
+	return bson.M{
+		"$or": []bson.M{
+			{c.SortField: bson.M{primaryOp: sortValue}},
+			{
+				c.SortField: sortValue,
+				"_id":       bson.M{secondaryOp: objectID},
+			},
+		},
+	}, nil
+}
+
+func (c *pageCursor) parseSortValue() (interface{}, error) {
+	switch c.ValueType {
+	case "i64":
+		v, err := strconv.ParseInt(c.SortValue, 10, 64)
+		if err != nil {
+			return nil, fmt.Errorf("%w: invalid int64 sort value", ErrInvalidCursor)
+		}
+		return v, nil
+	case "f64":
+		v, err := strconv.ParseFloat(c.SortValue, 64)
+		if err != nil {
+			return nil, fmt.Errorf("%w: invalid float sort value", ErrInvalidCursor)
+		}
+		return v, nil
+	case "str":
+		return c.SortValue, nil
+	case "oid":
+		return bson.ObjectIDFromHex(c.SortValue)
+	default:
+		return nil, fmt.Errorf("%w: unknown sort value type", ErrInvalidCursor)
+	}
+}
+
+func formatCursorSortValue(sortField string, sortValue interface{}) (string, string, error) {
+	if sortField == "_id" {
+		switch v := sortValue.(type) {
+		case bson.ObjectID:
+			return "oid", v.Hex(), nil
+		case string:
+			if _, err := bson.ObjectIDFromHex(v); err != nil {
+				return "", "", fmt.Errorf("%w: invalid ObjectID sort value", ErrInvalidCursor)
+			}
+			return "oid", v, nil
+		default:
+			return "", "", fmt.Errorf("%w: unsupported _id sort value", ErrInvalidCursor)
+		}
+	}
+
+	switch v := sortValue.(type) {
+	case int64:
+		return "i64", strconv.FormatInt(v, 10), nil
+	case int32:
+		return "i64", strconv.FormatInt(int64(v), 10), nil
+	case int:
+		return "i64", strconv.FormatInt(int64(v), 10), nil
+	case float64:
+		return "f64", strconv.FormatFloat(v, 'f', -1, 64), nil
+	case float32:
+		return "f64", strconv.FormatFloat(float64(v), 'f', -1, 64), nil
+	case string:
+		return "str", v, nil
+	case bson.ObjectID:
+		return "oid", v.Hex(), nil
+	default:
+		return "", "", fmt.Errorf("%w: unsupported sort value type for field %s", ErrInvalidCursor, sortField)
+	}
+}
+
+func documentSortValue(docMap bson.M, sortField string) (interface{}, error) {
+	if sortField == "_id" {
+		id, exists := docMap["_id"]
+		if !exists {
+			return nil, fmt.Errorf("sort field _id not found in document")
+		}
+		return id, nil
+	}
+
+	sortValue, exists := docMap[sortField]
+	if !exists {
+		return nil, fmt.Errorf("sort field %s not found in document", sortField)
+	}
+	return sortValue, nil
+}
+
+func documentIDHex(docMap bson.M) (string, error) {
+	id, exists := docMap["_id"]
+	if !exists {
+		return "", fmt.Errorf("document _id not found")
+	}
+	switch v := id.(type) {
+	case bson.ObjectID:
+		return v.Hex(), nil
+	case string:
+		if _, err := bson.ObjectIDFromHex(v); err != nil {
+			return "", fmt.Errorf("invalid document _id")
+		}
+		return v, nil
+	default:
+		return "", fmt.Errorf("unsupported document _id type")
+	}
+}
+
+func cloneFilter(filter bson.M) (bson.M, error) {
+	if len(filter) == 0 {
+		return bson.M{}, nil
+	}
+	raw, err := bson.Marshal(filter)
+	if err != nil {
+		return nil, err
+	}
+	var cloned bson.M
+	if err := bson.Unmarshal(raw, &cloned); err != nil {
+		return nil, err
+	}
+	return cloned, nil
+}
+
+func mergeFilters(base, extra bson.M) bson.M {
+	if len(extra) == 0 {
+		return base
+	}
+	if len(base) == 0 {
+		return extra
+	}
+	return bson.M{
+		"$and": []bson.M{base, extra},
+	}
+}

pkg/mongo/repository.go

@@ -2,7 +2,6 @@ package mongo
 
 import (
 	"context"
-	"encoding/base64"
 	"errors"
 	"fmt"
 	"time"
@@ -26,6 +25,8 @@ type Pagination struct {
 	Cursor    string `json:"cursor"`    // Base64 encoded cursor for cursor-based pagination
 	SortField string `json:"sortField"` // Field to sort by (default: "_id")
 	SortOrder int    `json:"sortOrder"` // Sort order: 1 for ascending, -1 for descending
+	// SkipCount skips CountDocuments (use with cursor or deep offset pagination).
+	SkipCount bool `json:"skipCount,omitempty"`
 	// Projection limits fields returned from Find (e.g. exclude large blobs on list endpoints).
 	Projection bson.M `json:"projection,omitempty"`
 }
@@ -132,8 +133,16 @@ func (r *repository[T]) FindAll(ctx context.Context, filter bson.M, pagination P
 		pagination.SortOrder = -1 // Default to descending
 	}
 
-	// Build sort options
+	// Build sort with _id tie-breaker for stable keyset pagination.
 	sort := bson.D{{Key: pagination.SortField, Value: pagination.SortOrder}}
+	if pagination.SortField != "_id" {
+		sort = append(sort, bson.E{Key: "_id", Value: pagination.SortOrder})
+	}
+
+	countFilter, err := cloneFilter(filter)
+	if err != nil {
+		return nil, fmt.Errorf("failed to clone filter for count: %w", err)
+	}
 
 	// Build find options
 	findOptions := options.Find().
@@ -150,10 +159,7 @@ func (r *repository[T]) FindAll(ctx context.Context, filter bson.M, pagination P
 		if err != nil {
 			return nil, err
 		}
-		// Merge cursor filter with existing filter
-		for key, value := range cursorFilter {
-			filter[key] = value
-		}
+		filter = mergeFilters(filter, cursorFilter)
 	} else if pagination.Skip > 0 {
 		// Use offset-based pagination
 		findOptions.SetSkip(int64(pagination.Skip))
@@ -187,21 +193,23 @@ func (r *repository[T]) FindAll(ctx context.Context, filter bson.M, pagination P
 		items = items[:pagination.Limit] // Remove the extra item
 	}
 
-	// Get total count for pagination metadata
-	var totalCount int64
-	var countErr error
-	totalCount, countErr = r.Count(ctx, filter)
-	if countErr != nil {
-		r.logger.Error("Failed to get total count", map[string]interface{}{
-			"error": countErr.Error(),
-		})
-		totalCount = 0
+	// Get total count for pagination metadata (never include cursor constraints).
+	var totalCount int64 = -1
+	if !pagination.SkipCount {
+		var countErr error
+		totalCount, countErr = r.Count(ctx, countFilter)
+		if countErr != nil {
+			r.logger.Error("Failed to get total count", map[string]interface{}{
+				"error": countErr.Error(),
+			})
+			totalCount = 0
+		}
 	}
 
 	// Generate next cursor
 	var nextCursor string
 	if hasMore && len(items) > 0 {
-		nextCursor, err = r.generateCursor(items[len(items)-1], pagination.SortField)
+		nextCursor, err = r.generateCursor(items[len(items)-1], pagination.SortField, pagination.SortOrder)
 		if err != nil {
 			r.logger.Error("Failed to generate next cursor", map[string]interface{}{
 				"error": err.Error(),
@@ -445,43 +453,20 @@ func (r *repository[T]) validatePagination(pagination Pagination) error {
 	return nil
 }
 
-// buildCursorFilter builds a filter for cursor-based pagination
+// buildCursorFilter builds a filter for cursor-based pagination.
 func (r *repository[T]) buildCursorFilter(cursor, sortField string, sortOrder int) (bson.M, error) {
-	// Decode cursor
-	cursorBytes, err := base64.StdEncoding.DecodeString(cursor)
+	pageCursor, err := decodePageCursor(cursor)
 	if err != nil {
-		return nil, fmt.Errorf("%w: invalid cursor encoding", ErrInvalidCursor)
+		return nil, err
 	}
-
-	var cursorValue interface{}
-
-	// Try to parse as ObjectID first
-	if sortField == "_id" {
-		if objectID, err := bson.ObjectIDFromHex(string(cursorBytes)); err == nil {
-			cursorValue = objectID
-		} else {
-			return nil, fmt.Errorf("%w: invalid ObjectID in cursor", ErrInvalidCursor)
-		}
-	} else {
-		// For other fields, try to parse as appropriate type
-		// This is a simplified implementation - you might want to make this more sophisticated
-		cursorValue = string(cursorBytes)
+	if !pageCursor.matches(sortField, sortOrder) {
+		return nil, fmt.Errorf("%w: cursor does not match requested sort", ErrInvalidCursor)
 	}
-
-	// Build comparison operator based on sort order
-	var operator string
-	if sortOrder == 1 {
-		operator = "$gt"
-	} else {
-		operator = "$lt"
-	}
-
-	return bson.M{sortField: bson.M{operator: cursorValue}}, nil
+	return pageCursor.mongoFilter()
 }
 
-// generateCursor generates a cursor from a document
-func (r *repository[T]) generateCursor(doc T, sortField string) (string, error) {
-	// Convert document to BSON
+// generateCursor generates an opaque cursor from the last document in a page.
+func (r *repository[T]) generateCursor(doc T, sortField string, sortOrder int) (string, error) {
 	docBytes, err := bson.Marshal(doc)
 	if err != nil {
 		return "", fmt.Errorf("failed to marshal document: %w", err)
@@ -492,29 +477,17 @@ func (r *repository[T]) generateCursor(doc T, sortField string) (string, error)
 		return "", fmt.Errorf("failed to unmarshal document: %w", err)
 	}
 
-	// Get the sort field value
-	sortValue, exists := docMap[sortField]
-	if !exists {
-		return "", fmt.Errorf("sort field %s not found in document", sortField)
+	sortValue, err := documentSortValue(docMap, sortField)
+	if err != nil {
+		return "", err
 	}
 
-	// Convert to string for encoding
-	var cursorStr string
-	switch v := sortValue.(type) {
-	case bson.ObjectID:
-		cursorStr = v.Hex()
-	case string:
-		cursorStr = v
-	case int64:
-		cursorStr = fmt.Sprintf("%d", v)
-	case float64:
-		cursorStr = fmt.Sprintf("%f", v)
-	default:
-		cursorStr = fmt.Sprintf("%v", v)
+	idHex, err := documentIDHex(docMap)
+	if err != nil {
+		return "", err
 	}
 
-	// Encode as base64
-	return base64.StdEncoding.EncodeToString([]byte(cursorStr)), nil
+	return encodePageCursor(sortField, sortOrder, sortValue, idHex)
 }
 
 // getModelID extracts the ID from a model

pkg/mongo/utils.go

@@ -299,6 +299,12 @@ func (pb *PaginationBuilder) Projection(projection bson.M) *PaginationBuilder {
 	return pb
 }
 
+// SkipCount disables CountDocuments for this query.
+func (pb *PaginationBuilder) SkipCount(skip bool) *PaginationBuilder {
+	pb.pagination.SkipCount = skip
+	return pb
+}
+
 // Build returns the final pagination options
 func (pb *PaginationBuilder) Build() Pagination {
 	return pb.pagination