oc-lib/dbs/mongo/mongo.go

package mongo

import (
	"context"
	"errors"
	"time"

	"cloud.o-forge.io/core/oc-lib/dbs"
	"cloud.o-forge.io/core/oc-lib/logs"
	"github.com/rs/zerolog"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

var (
	mngoClient *mongo.Client
	mngoDB     *mongo.Database
	MngoCtx    context.Context
	cancel     context.CancelFunc

	existingCollections []string

	ResourceMap map[string]interface{}
)

var MONGOService = MongoDB{}

type MongoConf interface {
	GetUrl() string
	GetDatabase() string
}

type MongoDB struct {
	Logger zerolog.Logger
}

func (m *MongoDB) Init(collections []string, config MongoConf) {
	// var baseConfig string
	m.Logger = logs.GetLogger()
	ResourceMap = make(map[string]interface{})

	m.Logger.Info().Msg("Connecting to" + config.GetUrl())
	MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	m.createClient(config.GetUrl())

	m.Logger.Info().Msg("Connecting mongo client to db " + config.GetDatabase())
	m.prepareDB(collections, config)

	m.Logger.Info().Msg("Database is READY")

}

func (m *MongoDB) createClient(MongoURL string) {

	var err error

	// Allows us to use marshal and unmarshall with results of FindOne() and others
	bsonOpts := &options.BSONOptions{
		UseJSONStructTags: true,
		NilSliceAsEmpty:   true,
	}

	clientOptions := options.Client().ApplyURI(MongoURL).SetBSONOptions(bsonOpts)
	mngoClient, err = mongo.Connect(MngoCtx, clientOptions)
	if err != nil {
		m.Logger.Fatal().Msg("Mongodb NewClient " + MongoURL + ":" + "err")
		panic(err)
	}

	// Ping the primary
	if mngoClient, err = mongo.Connect(MngoCtx, clientOptions); err != nil {
		m.Logger.Fatal().Msg("Mongodb connect " + MongoURL + ":" + "err")
		panic(err)
	}

	if err = mngoClient.Ping(MngoCtx, nil); err != nil {
		m.Logger.Fatal().Msg("Mongodb ping " + MongoURL + ":" + "err")
		panic(err)
	}

}

func (m *MongoDB) prepareDB(list_collection []string, config MongoConf) {
	var err error
	mngoDB = mngoClient.Database(config.GetDatabase())
	existingCollections, err = mngoDB.ListCollectionNames(MngoCtx, bson.D{})
	if err != nil {
		m.Logger.Fatal().Msg("Error contacting MongoDB\n" + err.Error())
	}
	collectionMap := make(map[string]bool)
	for _, name := range existingCollections {
		collectionMap[name] = true
	}
	// Only do the collection definition process if it doesn't already exists
	// we add the collection to the collection map from mongo/mongo_utils to provide faster access to the collection
	for _, collection_name := range list_collection {
		new_collection := mngoDB.Collection(collection_name)
		if _, exists := collectionMap[collection_name]; !exists {
			m.createCollection(collection_name, new_collection)
		} else {
			CollectionMap[collection_name] = new_collection
		}
	}
}

// Creates the collection with index specified in mongo/mongo_collections
// or use the basic collection creation function
func (m *MongoDB) createCollection(collection_name string, new_collection *mongo.Collection) {
	var err error
	CollectionMap[collection_name] = new_collection
	_, exists := IndexesMap[collection_name]
	if exists {
		if _, err = new_collection.Indexes().CreateMany(MngoCtx, IndexesMap[collection_name]); err != nil {
			var cmdErr mongo.CommandError
			if errors.As(err, &cmdErr) && cmdErr.Code != 85 {
				m.Logger.Fatal().Msg("Error creating indexes for " + collection_name + " collection : \n" + err.Error())
				panic(err)
			} else if !errors.As(err, &cmdErr) {
				m.Logger.Fatal().Msg("Unexpected error: " + err.Error())
				panic(err)
			}
		}
	} else {
		mngoDB.CreateCollection(MngoCtx, collection_name)
	}

}

func (m *MongoDB) DeleteOne(id string, collection_name string) (int64, int, error) {
	filter := bson.M{"_id": id}
	targetDBCollection := CollectionMap[collection_name]
	opts := options.Delete().SetHint(bson.D{{Key: "_id", Value: 1}})
	MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	result, err := targetDBCollection.DeleteOne(MngoCtx, filter, opts)
	if err != nil {
		m.Logger.Error().Msg("Couldn't insert resource: " + err.Error())
		return 0, 404, err
	}
	return result.DeletedCount, 200, nil
}

func (m *MongoDB) UpdateOne(set interface{}, id string, collection_name string) (string, int, error) {
	var doc map[string]interface{}
	b, _ := bson.Marshal(set)
	bson.Unmarshal(b, &doc)
	filter := bson.M{"_id": id}
	targetDBCollection := CollectionMap[collection_name]
	MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	_, err := targetDBCollection.UpdateOne(MngoCtx, filter, dbs.InputToBson(doc, true))
	if err != nil {
		m.Logger.Error().Msg("Couldn't update resource: " + err.Error())
		return "", 404, err
	}
	return id, 200, nil
}

func (m *MongoDB) StoreOne(obj interface{}, id string, collection_name string) (string, int, error) {
	var doc map[string]interface{}
	b, _ := bson.Marshal(obj)
	bson.Unmarshal(b, &doc)
	doc["_id"] = id
	targetDBCollection := CollectionMap[collection_name]

	MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	_, err := targetDBCollection.InsertOne(MngoCtx, doc)
	if err != nil {
		m.Logger.Error().Msg("Couldn't insert resource: " + err.Error())
		return "", 409, err
	}

	return id, 200, nil
}

func (m *MongoDB) LoadOne(id string, collection_name string) (*mongo.SingleResult, int, error) {
	filter := bson.M{"_id": id}
	targetDBCollection := CollectionMap[collection_name]

	MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	res := targetDBCollection.FindOne(MngoCtx, filter)
	if res.Err() != nil {
		m.Logger.Error().Msg("Couldn't find resource " + id + ". Error : " + res.Err().Error())
		err := res.Err()
		return nil, 404, err
	}
	return res, 200, nil
}