package mongo import ( "context" "encoding/json" "errors" lib "oc-lib" "oc-lib/dbs" "os" "time" "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" ) var ( mngoClient *mongo.Client mngoDB *mongo.Database MngoCtx context.Context cancel context.CancelFunc 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{}) lib.Logger = lib.CreateLogger("oclib", "") db_conf, err := os.ReadFile("tests/oclib_conf.json") if err != nil { lib.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"] lib.Logger.Info().Msg("Connecting to" + MongoURL) MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second) defer cancel() createClient(MongoURL) lib.Logger.Info().Msg("Connecting mongo client to db " + DBname) prepareDB(conf["DCNAME"], conf["DBPOINT"]) lib.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 { lib.Logger.Fatal().Msg("Mongodb NewClient " + MongoURL + ":" + "err") panic(err) } // Ping the primary if mngoClient, err = mongo.Connect(MngoCtx, clientOptions); err != nil { lib.Logger.Fatal().Msg("Mongodb connect " + MongoURL + ":" + "err") panic(err) } if err = mngoClient.Ping(MngoCtx, nil); err != nil { lib.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 { lib.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 { lib.Logger.Fatal().Msg("Error creating indexes for " + collection_name + " collection : \n" + err.Error()) panic(err) } else if !errors.As(err, &cmdErr) { lib.Logger.Fatal().Msg("Unexpected error: " + err.Error()) panic(err) } } } else { mngoDB.CreateCollection(MngoCtx, collection_name) } } func DeleteOne(id string, collection_name string) (int64, error) { filter := bson.M{"_id": GetObjIDFromString(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 { lib.Logger.Error().Msg("Couldn't insert resource: " + err.Error()) return 0, err } return result.DeletedCount, nil } func UpdateOne(set map[string]interface{}, id string, collection_name string) (string, error) { filter := bson.M{"_id": GetObjIDFromString(id)} targetDBCollection := CollectionMap[collection_name] MngoCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second) defer cancel() result, err := targetDBCollection.UpdateOne(MngoCtx, filter, dbs.InputToBson(set, true)) if err != nil { lib.Logger.Error().Msg("Couldn't insert resource: " + err.Error()) return "", err } return result.UpsertedID.(primitive.ObjectID).Hex(), nil } 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 { lib.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) { 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 { lib.Logger.Error().Msg("Couldn't find resource " + id + ". Error : " + res.Err().Error()) err = res.Err() return nil, err } return res, nil }