oc-lib/dbs/mongo/mongo.go

package mongo

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"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
	mngoCollections     []string
	mngoConfig          MongoConf
	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()
	mngoCollections = collections
	mngoConfig = config
	if err := m.createClient(config.GetUrl()); err == nil {
		m.connect()
	}
}

func (m *MongoDB) connect() error {
	if mngoClient == nil {
		m.Logger.Info().Msg("Connecting mongo client to db " + mngoConfig.GetDatabase())
		m.prepareDB(mngoCollections, mngoConfig)

		m.Logger.Info().Msg("Database is READY")
		return nil
	} else {
		return m.createClient(mngoConfig.GetUrl())
	}
}

func (m *MongoDB) createClient(MongoURL string) error {
	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 {
		mngoClient = nil
		return errors.New("Mongodb NewClient " + MongoURL + ":" + "err")
	}

	// Ping the primary
	if mngoClient, err = mongo.Connect(MngoCtx, clientOptions); err != nil {
		mngoClient = nil
		return errors.New("Mongodb connect " + MongoURL + ":" + "err")
	}

	if err = mngoClient.Ping(MngoCtx, nil); err != nil {
		mngoClient = nil
		return errors.New("Mongodb connect " + MongoURL + ":" + "err")
	}
	return nil
}

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) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return 0, 503, err
	}
	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) DeleteMultiple(f map[string]interface{}, collection_name string) (int64, int, error) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return 0, 503, err
	}
	filter := bson.D{}
	for k, v := range f {
		filter = append(filter, bson.E{Key: k, Value: v})
	}
	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.DeleteMany(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) UpdateMultiple(set interface{}, filter map[string]interface{}, collection_name string) (int64, int, error) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return 0, 503, err
	}
	var doc map[string]interface{}
	b, _ := bson.Marshal(set)
	bson.Unmarshal(b, &doc)
	f := bson.D{}
	for k, v := range filter {
		f = append(f, bson.E{Key: k, Value: v})
	}
	targetDBCollection := CollectionMap[collection_name]
	MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	res, err := targetDBCollection.UpdateMany(MngoCtx, f, dbs.InputToBson(doc, true))
	if err != nil {
		m.Logger.Error().Msg("Couldn't update resource: " + err.Error())
		return 0, 404, err
	}
	return res.UpsertedCount, 200, nil
}

func (m *MongoDB) UpdateOne(set interface{}, id string, collection_name string) (string, int, error) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return "", 503, err
	}
	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) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return "", 503, err
	}
	var doc map[string]interface{}
	b, _ := bson.Marshal(obj)
	bson.Unmarshal(b, &doc)
	doc["_id"] = id
	targetDBCollection := CollectionMap[collection_name]
	fmt.Println("DB", collection_name, targetDBCollection, CollectionMap)
	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) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return nil, 503, err
	}
	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
}

func (m *MongoDB) Search(search string, filter []string, collection_name string) (*mongo.Cursor, int, error) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return nil, 503, err
	}
	opts := options.Find()
	opts.SetLimit(100)
	if strings.TrimSpace(search) == "*" {
		search = ""
	}
	search = ".*" + strings.TrimSpace(search) + ".*"
	targetDBCollection := CollectionMap[collection_name]
	list := []bson.M{}
	for _, k := range filter {
		list = append(list, bson.M{k: bson.M{"$regex": search, "$options": "i"}})
	}
	MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if cursor, err := targetDBCollection.Find(
		MngoCtx,
		bson.M{"$or": list},
		opts,
	); err != nil {
		return nil, 404, err
	} else {
		return cursor, 200, nil
	}
}

func (m *MongoDB) LoadFilter(filter map[string]interface{}, collection_name string) (*mongo.Cursor, int, error) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return nil, 503, err
	}
	f := bson.D{}
	for k, v := range filter {
		f = append(f, bson.E{Key: k, Value: v})
	}
	targetDBCollection := CollectionMap[collection_name]

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

	res, err := targetDBCollection.Find(MngoCtx, f)
	if err != nil {
		m.Logger.Error().Msg("Couldn't find any resources. Error : " + err.Error())
		return nil, 404, err
	}
	return res, 200, nil
}

func (m *MongoDB) LoadAll(collection_name string) (*mongo.Cursor, int, error) {
	if err := m.createClient(mngoConfig.GetUrl()); err != nil {
		return nil, 503, err
	}
	targetDBCollection := CollectionMap[collection_name]

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

	res, err := targetDBCollection.Find(MngoCtx, bson.D{})
	if err != nil {
		m.Logger.Error().Msg("Couldn't find any resources. Error : " + err.Error())
		return nil, 404, err
	}
	return res, 200, nil
}

func (m *MongoDB) toOperator(operator string) string {
	if operator == "like" {
		return "$regex"
	} else if operator == "exists" {
		return "$exists"
	} else if operator == "in" {
		return "$in"
	} else if operator == "gte" {
		return "$gte"
	} else if operator == "gt" {
		return "$gt"
	} else if operator == "lte" {
		return "$lte"
	} else if operator == "lt" {
		return "$lt"
	} else if operator == "eq" {
		return "$match"
	}
	return operator
}