oc-lib/mongo/mongo.go
2024-07-17 18:02:30 +02:00

210 lines
5.1 KiB
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
}