pasty/internal/storage/mongodb/mongodb_driver.go

package mongodb

import (
	"context"
	"time"

	"github.com/lus/pasty/internal/config"
	"github.com/lus/pasty/internal/shared"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
)

// MongoDBDriver represents the MongoDB storage driver
type MongoDBDriver struct {
	client     *mongo.Client
	database   string
	collection string
}

// Initialize initializes the MongoDB storage driver
func (driver *MongoDBDriver) Initialize() error {
	// Define the context for the following database operation
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Connect to the MongoDB host
	client, err := mongo.Connect(ctx, options.Client().ApplyURI(config.Current.MongoDB.DSN))
	if err != nil {
		return err
	}

	// Ping the MongoDB host
	err = client.Ping(ctx, readpref.Primary())
	if err != nil {
		return err
	}

	// Set the driver attributes
	driver.client = client
	driver.database = config.Current.MongoDB.Database
	driver.collection = config.Current.MongoDB.Collection
	return nil
}

// Terminate terminates the MongoDB storage driver
func (driver *MongoDBDriver) Terminate() error {
	return driver.client.Disconnect(context.TODO())
}

// ListIDs returns a list of all existing paste IDs
func (driver *MongoDBDriver) ListIDs() ([]string, error) {
	// Define the collection to use for this database operation
	collection := driver.client.Database(driver.database).Collection(driver.collection)

	// Define the context for the following database operation
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Retrieve all paste documents
	result, err := collection.Find(ctx, bson.M{})
	if err != nil {
		return nil, err
	}

	// Decode all paste documents
	var pasteSlice []shared.Paste
	err = result.All(ctx, &pasteSlice)
	if err != nil {
		return nil, err
	}

	// Read and return the IDs of all paste objects
	var ids []string
	for _, paste := range pasteSlice {
		ids = append(ids, paste.ID)
	}
	return ids, nil
}

// Get loads a paste
func (driver *MongoDBDriver) Get(id string) (*shared.Paste, error) {
	// Define the collection to use for this database operation
	collection := driver.client.Database(driver.database).Collection(driver.collection)

	// Define the context for the following database operation
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Try to retrieve the corresponding paste document
	filter := bson.M{"_id": id}
	result := collection.FindOne(ctx, filter)
	err := result.Err()
	if err != nil {
		if err == mongo.ErrNoDocuments {
			return nil, nil
		}
		return nil, err
	}

	// Return the retrieved paste object
	paste := new(shared.Paste)
	err = result.Decode(paste)
	if err != nil {
		return nil, err
	}
	return paste, nil
}

// Save saves a paste
func (driver *MongoDBDriver) Save(paste *shared.Paste) error {
	// Define the collection to use for this database operation
	collection := driver.client.Database(driver.database).Collection(driver.collection)

	// Define the context for the following database operation
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Upsert the paste object
	filter := bson.M{"_id": paste.ID}
	_, err := collection.UpdateOne(ctx, filter, paste, options.Update().SetUpsert(true))
	return err
}

// Delete deletes a paste
func (driver *MongoDBDriver) Delete(id string) error {
	// Define the collection to use for this database operation
	collection := driver.client.Database(driver.database).Collection(driver.collection)

	// Define the context for the following database operation
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Delete the document
	filter := bson.M{"_id": id}
	_, err := collection.DeleteOne(ctx, filter)
	return err
}

// Cleanup cleans up the expired pastes
func (driver *MongoDBDriver) Cleanup() (int, error) {
	// Retrieve all paste IDs
	ids, err := driver.ListIDs()
	if err != nil {
		return 0, err
	}

	// Define the amount of deleted items
	deleted := 0

	// Loop through all pastes
	for _, id := range ids {
		// Retrieve the paste object
		paste, err := driver.Get(id)
		if err != nil {
			return 0, err
		}

		// Delete the paste if it is expired
		lifetime := config.Current.AutoDelete.Lifetime
		if paste.AutoDelete && paste.Created+int64(lifetime.Seconds()) < time.Now().Unix() {
			err = driver.Delete(id)
			if err != nil {
				return 0, err
			}
			deleted++
		}
	}
	return deleted, nil
}