go-mongo/app/database.go

package app

import (
	"context"
	"fmt"
	"time"

	"github.com/fatih/color"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

var (
	Mongo *mongo.Client
)

// Connect to MongoDB and return a cancel function
func Connect() (context.CancelFunc, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)

	var err error
	Mongo, err = mongo.Connect(ctx, options.Client().ApplyURI(Config.MongoURI))
	if err != nil {
		color.Red("❌ Database connection failed\n" + err.Error())
		return cancel, err
	}

	if err = Mongo.Ping(ctx, nil); err != nil {
		color.Red("❌ Connection Failed to Database")
		color.Red(err.Error())
		return cancel, err
	}

	color.Green("✅ Connected to MongoDB")
	return cancel, nil
}

// Disconnect cleanly from MongoDB
func Disconnect(ctx context.Context) error {
	return Mongo.Disconnect(ctx)
}

// Collection returns a MongoDB collection
func Collection(name string) *mongo.Collection {
	return Mongo.Database(Config.MongoDB).Collection(name)
}

// Find returns all matching documents from a collection
func Find[T any](ctx context.Context, coll string, filter any) ([]*T, error) {
	var results []*T

	cursor, err := Collection(coll).Find(ctx, filter)
	if err != nil {
		return nil, fmt.Errorf("error while finding documents in %s: %w", coll, err)
	}

	if err := cursor.All(ctx, &results); err != nil {
		return nil, fmt.Errorf("error while parsing documents in %s: %w", coll, err)
	}

	return results, nil
}

// FindOne finds a single document from a collection
func FindOne[T any](ctx context.Context, coll string, filter any) (*T, error) {
	var result T
	err := Collection(coll).FindOne(ctx, filter).Decode(&result)

	if err == mongo.ErrNoDocuments {
		return nil, fmt.Errorf("no data found")
	}
	if err != nil {
		return nil, fmt.Errorf("findOne error: %w", err)
	}

	return &result, nil
}

// FindById finds a document by its ObjectID
func FindById[T any](ctx context.Context, coll string, id primitive.ObjectID) (*T, error) {
	var result T
	err := Collection(coll).FindOne(ctx, bson.M{"_id": id}).Decode(&result)

	if err == mongo.ErrNoDocuments {
		return nil, fmt.Errorf("no data found")
	}
	if err != nil {
		return nil, fmt.Errorf("findById error: %w", err)
	}

	return &result, nil
}

// InsertOne inserts a new document and returns it
func InsertOne[T any](ctx context.Context, coll string, input any) (*T, error) {
	now := time.Now()

	// Marshal + Unmarshal to apply bson tags
	data, err := bson.Marshal(input)
	if err != nil {
		return nil, fmt.Errorf("marshal error: %w", err)
	}

	var doc bson.M
	if err := bson.Unmarshal(data, &doc); err != nil {
		return nil, fmt.Errorf("unmarshal error: %w", err)
	}

	doc["createdAt"] = now
	doc["updatedAt"] = now

	if user, err := CurrentUser(ctx); err == nil {
		if id, err := primitive.ObjectIDFromHex(user.ID); err == nil {
			doc["createdById"] = id
			doc["updatedById"] = id
			doc["ownerId"] = id
		}
	}

	res, err := Collection(coll).InsertOne(ctx, doc)
	if err != nil {
		return nil, fmt.Errorf("insert error: %w", err)
	}

	var inserted T
	err = Collection(coll).FindOne(ctx, bson.M{"_id": res.InsertedID}).Decode(&inserted)
	if err != nil {
		return nil, fmt.Errorf("fetch inserted document error: %w", err)
	}

	return &inserted, nil
}

// UpdateByID updates a document and returns the new version
func UpdateByID[T any](ctx context.Context, coll string, id string, input any) (*T, error) {
	docID, err := primitive.ObjectIDFromHex(id)
	if err != nil {
		return nil, fmt.Errorf("invalid ID format: %w", err)
	}

	now := time.Now()

	data, err := bson.Marshal(input)
	if err != nil {
		return nil, fmt.Errorf("marshal error: %w", err)
	}

	var updateData bson.M
	if err := bson.Unmarshal(data, &updateData); err != nil {
		return nil, fmt.Errorf("unmarshal error: %w", err)
	}

	updateData["updatedAt"] = now

	if user, err := CurrentUser(ctx); err == nil {
		if userID, err := primitive.ObjectIDFromHex(user.ID); err == nil {
			updateData["updatedById"] = userID
		}
	}

	res, err := Collection(coll).UpdateByID(ctx, docID, bson.M{"$set": updateData})
	if err != nil {
		return nil, fmt.Errorf("update error: %w", err)
	}

	if res.MatchedCount == 0 {
		return nil, fmt.Errorf("no document found with ID %s", id)
	}
	if res.ModifiedCount == 0 {
		color.Yellow("⚠️ Document with ID %s was found but not modified", id)
	}

	var updated T
	err = Collection(coll).FindOne(ctx, bson.M{"_id": docID}).Decode(&updated)
	if err != nil {
		return nil, fmt.Errorf("error fetching updated document: %w", err)
	}

	return &updated, nil
}