package mongo import ( "context" "encoding/base64" "errors" "fmt" "time" "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/mongo" "go.mongodb.org/mongo-driver/v2/mongo/options" ) // Common errors var ( ErrDocumentNotFound = errors.New("document not found") ErrInvalidCursor = errors.New("invalid cursor") ErrInvalidPagination = errors.New("invalid pagination parameters") ) // Pagination represents pagination parameters for queries type Pagination struct { Limit int `json:"limit"` // Number of documents to return Skip int `json:"skip"` // Number of documents to skip (for offset-based) 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 } // PaginatedResult represents a paginated result set type PaginatedResult[T any] struct { Items []T `json:"items"` // The actual documents TotalCount int64 `json:"totalCount"` // Total number of documents matching the filter NextCursor string `json:"nextCursor"` // Base64 encoded cursor for next page HasMore bool `json:"hasMore"` // Whether there are more documents } // Repository defines the interface for MongoDB operations type Repository[T any] interface { // FindByID retrieves a document by its ID FindByID(ctx context.Context, id string) (*T, error) // FindAll retrieves documents with optional filtering and pagination FindAll(ctx context.Context, filter bson.M, pagination Pagination) (*PaginatedResult[T], error) // Create inserts a new document Create(ctx context.Context, model *T) error // Update updates an existing document Update(ctx context.Context, model *T) error // Delete removes a document by ID Delete(ctx context.Context, id string) error // Count returns the number of documents matching the filter Count(ctx context.Context, filter bson.M) (int64, error) // Aggregate executes an aggregation pipeline Aggregate(ctx context.Context, pipeline mongo.Pipeline) ([]bson.M, error) // FindOne retrieves a single document matching the filter FindOne(ctx context.Context, filter bson.M) (*T, error) // FindMany retrieves multiple documents without pagination FindMany(ctx context.Context, filter bson.M, limit int) ([]T, error) } // repository implements the Repository interface type repository[T any] struct { collection *mongo.Collection logger Logger } // NewRepository creates a new repository instance func NewRepository[T any](collection *mongo.Collection, logger Logger) Repository[T] { return &repository[T]{ collection: collection, logger: logger, } } // FindByID retrieves a document by its ID func (r *repository[T]) FindByID(ctx context.Context, id string) (*T, error) { objectID, err := bson.ObjectIDFromHex(id) if err != nil { r.logger.Error("Invalid ObjectID format", map[string]interface{}{ "id": id, "error": err.Error(), }) return nil, fmt.Errorf("invalid id format: %w", err) } filter := bson.M{"_id": objectID} var result T err = r.collection.FindOne(ctx, filter).Decode(&result) if err != nil { if errors.Is(err, mongo.ErrNoDocuments) { return nil, ErrDocumentNotFound } r.logger.Error("Failed to find document by ID", map[string]interface{}{ "id": id, "error": err.Error(), }) return nil, fmt.Errorf("failed to find document: %w", err) } return &result, nil } // FindAll retrieves documents with optional filtering and pagination func (r *repository[T]) FindAll(ctx context.Context, filter bson.M, pagination Pagination) (*PaginatedResult[T], error) { // Validate pagination parameters if err := r.validatePagination(pagination); err != nil { return nil, err } // Set default values if pagination.Limit <= 0 { pagination.Limit = 10 } if pagination.SortField == "" { pagination.SortField = "_id" } if pagination.SortOrder == 0 { pagination.SortOrder = -1 // Default to descending } // Build sort options sort := bson.D{{Key: pagination.SortField, Value: pagination.SortOrder}} // Build find options findOptions := options.Find(). SetSort(sort). SetLimit(int64(pagination.Limit + 1)) // +1 to check if there are more documents // Handle cursor-based pagination if pagination.Cursor != "" { cursorFilter, err := r.buildCursorFilter(pagination.Cursor, pagination.SortField, pagination.SortOrder) if err != nil { return nil, err } // Merge cursor filter with existing filter for key, value := range cursorFilter { filter[key] = value } } else if pagination.Skip > 0 { // Use offset-based pagination findOptions.SetSkip(int64(pagination.Skip)) } // Execute query cursor, err := r.collection.Find(ctx, filter, findOptions) if err != nil { r.logger.Error("Failed to execute find query", map[string]interface{}{ "filter": filter, "pagination": pagination, "error": err.Error(), }) return nil, fmt.Errorf("failed to execute query: %w", err) } defer cursor.Close(ctx) // Decode results var items []T if err = cursor.All(ctx, &items); err != nil { r.logger.Error("Failed to decode query results", map[string]interface{}{ "error": err.Error(), }) return nil, fmt.Errorf("failed to decode results: %w", err) } // Check if there are more documents hasMore := false if len(items) > pagination.Limit { hasMore = true items = items[:pagination.Limit] // Remove the extra item } // Get total count totalCount, err := r.Count(ctx, filter) if err != nil { r.logger.Error("Failed to get total count", map[string]interface{}{ "error": err.Error(), }) // Don't fail the entire query, just set count to 0 totalCount = 0 } // Generate next cursor var nextCursor string if hasMore && len(items) > 0 { nextCursor, err = r.generateCursor(items[len(items)-1], pagination.SortField) if err != nil { r.logger.Error("Failed to generate next cursor", map[string]interface{}{ "error": err.Error(), }) // Don't fail the entire query, just leave cursor empty } } return &PaginatedResult[T]{ Items: items, TotalCount: totalCount, NextCursor: nextCursor, HasMore: hasMore, }, nil } // Create inserts a new document func (r *repository[T]) Create(ctx context.Context, model *T) error { // Set timestamps if the model implements Timestampable if timestampable, ok := any(model).(Timestampable); ok { now := time.Now().Unix() timestampable.SetCreatedAt(now) timestampable.SetUpdatedAt(now) } result, err := r.collection.InsertOne(ctx, model) if err != nil { r.logger.Error("Failed to create document", map[string]interface{}{ "error": err.Error(), }) return fmt.Errorf("failed to create document: %w", err) } // Set the ID if the model implements IDSetter if idSetter, ok := any(model).(IDSetter); ok { if oid, ok := result.InsertedID.(bson.ObjectID); ok { idSetter.SetID(oid.Hex()) } } r.logger.Info("Document created successfully", map[string]interface{}{ "id": result.InsertedID, }) return nil } // Update updates an existing document func (r *repository[T]) Update(ctx context.Context, model *T) error { // Set updated timestamp if the model implements Timestampable if timestampable, ok := any(model).(Timestampable); ok { timestampable.SetUpdatedAt(time.Now().Unix()) } // Get the ID from the model id, err := r.getModelID(model) if err != nil { return err } objectID, err := bson.ObjectIDFromHex(id) if err != nil { return fmt.Errorf("invalid id format: %w", err) } filter := bson.M{"_id": objectID} update := bson.M{"$set": model} result, err := r.collection.UpdateOne(ctx, filter, update) if err != nil { r.logger.Error("Failed to update document", map[string]interface{}{ "id": id, "error": err.Error(), }) return fmt.Errorf("failed to update document: %w", err) } if result.MatchedCount == 0 { return ErrDocumentNotFound } r.logger.Info("Document updated successfully", map[string]interface{}{ "id": id, }) return nil } // Delete removes a document by ID func (r *repository[T]) Delete(ctx context.Context, id string) error { objectID, err := bson.ObjectIDFromHex(id) if err != nil { return fmt.Errorf("invalid id format: %w", err) } filter := bson.M{"_id": objectID} result, err := r.collection.DeleteOne(ctx, filter) if err != nil { r.logger.Error("Failed to delete document", map[string]interface{}{ "id": id, "error": err.Error(), }) return fmt.Errorf("failed to delete document: %w", err) } if result.DeletedCount == 0 { return ErrDocumentNotFound } r.logger.Info("Document deleted successfully", map[string]interface{}{ "id": id, }) return nil } // Count returns the number of documents matching the filter func (r *repository[T]) Count(ctx context.Context, filter bson.M) (int64, error) { count, err := r.collection.CountDocuments(ctx, filter) if err != nil { r.logger.Error("Failed to count documents", map[string]interface{}{ "filter": filter, "error": err.Error(), }) return 0, fmt.Errorf("failed to count documents: %w", err) } return count, nil } // Aggregate executes an aggregation pipeline func (r *repository[T]) Aggregate(ctx context.Context, pipeline mongo.Pipeline) ([]bson.M, error) { cursor, err := r.collection.Aggregate(ctx, pipeline) if err != nil { r.logger.Error("Failed to execute aggregation", map[string]interface{}{ "pipeline": pipeline, "error": err.Error(), }) return nil, fmt.Errorf("failed to execute aggregation: %w", err) } defer cursor.Close(ctx) var results []bson.M if err = cursor.All(ctx, &results); err != nil { r.logger.Error("Failed to decode aggregation results", map[string]interface{}{ "error": err.Error(), }) return nil, fmt.Errorf("failed to decode aggregation results: %w", err) } return results, nil } // FindOne retrieves a single document matching the filter func (r *repository[T]) FindOne(ctx context.Context, filter bson.M) (*T, error) { var result T err := r.collection.FindOne(ctx, filter).Decode(&result) if err != nil { if errors.Is(err, mongo.ErrNoDocuments) { return nil, ErrDocumentNotFound } r.logger.Error("Failed to find document", map[string]interface{}{ "filter": filter, "error": err.Error(), }) return nil, fmt.Errorf("failed to find document: %w", err) } return &result, nil } // FindMany retrieves multiple documents without pagination func (r *repository[T]) FindMany(ctx context.Context, filter bson.M, limit int) ([]T, error) { findOptions := options.Find() if limit > 0 { findOptions.SetLimit(int64(limit)) } cursor, err := r.collection.Find(ctx, filter, findOptions) if err != nil { r.logger.Error("Failed to execute find query", map[string]interface{}{ "filter": filter, "limit": limit, "error": err.Error(), }) return nil, fmt.Errorf("failed to execute query: %w", err) } defer cursor.Close(ctx) var results []T if err = cursor.All(ctx, &results); err != nil { r.logger.Error("Failed to decode query results", map[string]interface{}{ "error": err.Error(), }) return nil, fmt.Errorf("failed to decode results: %w", err) } return results, nil } // Helper methods // validatePagination validates pagination parameters func (r *repository[T]) validatePagination(pagination Pagination) error { if pagination.Limit < 0 { return fmt.Errorf("%w: limit cannot be negative", ErrInvalidPagination) } if pagination.Skip < 0 { return fmt.Errorf("%w: skip cannot be negative", ErrInvalidPagination) } if pagination.SortOrder != 0 && pagination.SortOrder != 1 && pagination.SortOrder != -1 { return fmt.Errorf("%w: sort order must be 1, -1, or 0", ErrInvalidPagination) } return nil } // 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) if err != nil { return nil, fmt.Errorf("%w: invalid cursor encoding", ErrInvalidCursor) } 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) } // 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 } // generateCursor generates a cursor from a document func (r *repository[T]) generateCursor(doc T, sortField string) (string, error) { // Convert document to BSON docBytes, err := bson.Marshal(doc) if err != nil { return "", fmt.Errorf("failed to marshal document: %w", err) } var docMap bson.M if err = bson.Unmarshal(docBytes, &docMap); err != nil { 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) } // 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) } // Encode as base64 return base64.StdEncoding.EncodeToString([]byte(cursorStr)), nil } // getModelID extracts the ID from a model func (r *repository[T]) getModelID(model *T) (string, error) { if idGetter, ok := any(model).(IDGetter); ok { return idGetter.GetID(), nil } // Fallback: try to get ID from BSON docBytes, err := bson.Marshal(model) if err != nil { return "", fmt.Errorf("failed to marshal model: %w", err) } var docMap bson.M if err = bson.Unmarshal(docBytes, &docMap); err != nil { return "", fmt.Errorf("failed to unmarshal model: %w", err) } if id, exists := docMap["_id"]; exists { if oid, ok := id.(bson.ObjectID); ok { return oid.Hex(), nil } } return "", errors.New("model does not have an ID field") }