oc-lib/mongo/mongo.go

package mongo

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"time"

	"github.com/rs/zerolog"
	"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"

	"oc-lib/logs"
)

var (
	mngoClient 	*mongo.Client
	mngoDB     	*mongo.Database
	MngoCtx    	context.Context
	cancel 		context.CancelFunc
	logger		zerolog.Logger
	existingCollections	[]string

	ResourceMap map[string]interface{}

)



func init() {

	// var baseConfig string
	var err error
	var conf map[string]string
	var MongoURL string
	var DBname string

	ResourceMap = make(map[string]interface{})
	
	logger := logs.CreateLogger("oclib","")

	db_conf, err := os.ReadFile("tests/oclib_conf.json")
	if err != nil {
		logger.Fatal().Msg("Could not find configuration file")
	}
	json.Unmarshal(db_conf,&conf)
	
	if len(os.Getenv("DOCKER_ENVIRONMENT")) == 0 {
		MongoURL = conf["DB_URL_LOCAL"]
	} else {
		MongoURL = conf["DB_URL_DOCKER"]
	}

	DBname = conf["DCNAME"] + "-" + conf["DBPOINT"]


	logger.Info().Msg("Connecting to" + MongoURL)
	
	MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	createClient(MongoURL)

	logger.Info().Msg("Connecting mongo client to db " + DBname)
	prepareDB(conf["DCNAME"],conf["DBPOINT"])

	logger.Info().Msg("Database is READY")

}

func 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 {
		logger.Fatal().Msg("Mongodb NewClient " + MongoURL + ":"  + "err" )
		panic(err)
	}


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

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

}

func prepareDB(dc_name string, db_point string) {

	var err error
	DBname := dc_name + "-"  + db_point
	mngoDB = mngoClient.Database(DBname)
	
	list_collection := [...]string{"data","processing","storage","datacenter","workspace","schedule","workflow"}

	existingCollections, err = mngoDB.ListCollectionNames(MngoCtx,  bson.D{})
	if err != nil {
		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 {
			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 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 {
				logger.Fatal().Msg("Error creating indexes for " + collection_name + " collection : \n" + err.Error())
				panic(err)
			} else if !errors.As(err, &cmdErr) {
				logger.Fatal().Msg("Unexpected error: " + err.Error())
				panic(err)
			}
		}
	} else {
		mngoDB.CreateCollection(MngoCtx, collection_name)
	}

}

func StoreOne(obj interface{}, collection_name string) (string, error) {
	targetDBCollection := CollectionMap[collection_name]

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

	result, err := targetDBCollection.InsertOne(MngoCtx,obj)
	if err != nil {
		logger.Error().Msg("Couldn't insert resource: " + err.Error())
		return "", err
	}

	return result.InsertedID.(primitive.ObjectID).Hex(), nil
}

func LoadOne(id string, collection_name string) ( res *mongo.SingleResult , err error){
	
	// new_obj := ResourceMap[collection_name]
	// var doc bson.D
	
	filter := bson.M{"_id": GetObjIDFromString(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 {
		logger.Error().Msg("Couldn't find resource " + id + ". Error : " + res.Err().Error())
		err = res.Err()
		return nil, err
	}

	return res, nil
	
}

func GetObjIDFromString(id string) interface{} {
	objectID, err := primitive.ObjectIDFromHex(id)
	if err == nil {
		return objectID
	}

	return id
}