149 lines
3.6 KiB
Go
149 lines
3.6 KiB
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
|
|
}
|