Capstone/ui/db/db.go

package db

import (
	"context"
	"errors"
	"time"

	"git.preston-baxter.com/Preston_PLB/capstone/frontend-service/config"
	"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"
	"go.mongodb.org/mongo-driver/mongo/readpref"
)

// Interface for any object that wants to take advantage of the DB package
type Model interface {

	//Should return the _id field of the object if it exits
	//if it is new it should generate a new objectId
	MongoId() primitive.ObjectID

	//It is expected that this will update the CommonFields part of the model
	UpdateObjectInfo()
}

type DB struct {
	client *mongo.Client
}

func NewClient(uri string) (*DB, error) {
	opts := options.Client().ApplyURI(uri).SetConnectTimeout(60 * time.Second)
	client, err := mongo.Connect(context.Background(), opts)
	if err != nil {
		return nil, err
	}

	if err := client.Ping(context.Background(), readpref.Primary()); err != nil {
		return nil, err
	}

	return &DB{client: client}, nil
}

// Upserts
func (db *DB) SaveModel(m Model) error {
	conf := config.Config()

	m.UpdateObjectInfo()
	opts := options.Update().SetUpsert(true)
	res, err := db.client.Database(conf.Mongo.EntDb).Collection(conf.Mongo.EntCol).UpdateOne(context.Background(), bson.M{"_id": m.MongoId()}, bson.M{"$set": m}, opts)
	if err != nil {
		return err
	}

	if res.MatchedCount == 0 && res.ModifiedCount == 0 && res.UpsertedCount == 0 {
		return errors.New("Failed to save model properly")
	}

	return nil
}

func (db *DB) SaveModels(m ...Model) error {
	conf := config.Config()

	writeEntry := make([]mongo.WriteModel, len(m))

	for index, model := range m {
		entry := mongo.NewUpdateOneModel()
		entry.SetFilter(bson.M{"_id": model.MongoId()})
		entry.SetUpsert(true)
		entry.SetUpdate(bson.M{"$set": model})
		model.UpdateObjectInfo()

		writeEntry[index] = entry
	}

	opts := options.BulkWrite()
	res, err := db.client.Database(conf.Mongo.EntDb).Collection(conf.Mongo.EntCol).BulkWrite(context.Background(), writeEntry, opts)
	if err != nil {
		return err
	}

	if res.MatchedCount == 0 && res.ModifiedCount == 0 && res.UpsertedCount == 0 {
		return errors.New("Failed to save models properly")
	}

	return nil
}

// For allowing more varidaic like things
func saveModels[T Model](db *DB, m ...T) error {
	conf := config.Config()

	writeEntry := make([]mongo.WriteModel, len(m))

	for index, model := range m {
		entry := mongo.NewUpdateOneModel()
		entry.SetFilter(bson.M{"_id": model.MongoId()})
		entry.SetUpsert(true)
		entry.SetUpdate(bson.M{"$set": model})
		model.UpdateObjectInfo()

		writeEntry[index] = entry
	}

	opts := options.BulkWrite()
	res, err := db.client.Database(conf.Mongo.EntDb).Collection(conf.Mongo.EntCol).BulkWrite(context.Background(), writeEntry, opts)
	if err != nil {
		return err
	}

	if res.MatchedCount == 0 && res.ModifiedCount == 0 && res.UpsertedCount == 0 {
		return errors.New("Failed to save models properly")
	}

	return nil
}

// Doesn't upsert
func (db *DB) InsertModel(m Model) error {
	conf := config.Config()

	m.UpdateObjectInfo()
	opts := options.InsertOne()
	_, err := db.client.Database(conf.Mongo.EntDb).Collection(conf.Mongo.EntCol).InsertOne(context.Background(), m, opts)
	if err != nil {
		return err
	}

	return nil
}

func (db *DB) DeleteModel(m Model) error {
	conf := config.Config()

	opts := options.Delete()
	res, err := db.client.Database(conf.Mongo.EntDb).Collection(conf.Mongo.EntCol).DeleteOne(context.Background(), bson.M{"_id": m.MongoId()}, opts)
	if err != nil {
		return err
	}

	if res.DeletedCount == 0 {
		return errors.New("There was no item to delete")
	}

	return nil
}