From c14683ec0d96045186cb59980c87f26b94194dbf Mon Sep 17 00:00:00 2001
From: Cedric Verstraeten <cedric@verstraeten.io>
Date: Fri, 18 Oct 2024 15:48:04 +0200
Subject: [PATCH] update database client

---
 machinery/src/database/main.go | 65 +++++++++++++++++++++++-----------
 1 file changed, 45 insertions(+), 20 deletions(-)

diff --git a/machinery/src/database/main.go b/machinery/src/database/main.go
index 23b68d5..c9bf1ff 100644
--- a/machinery/src/database/main.go
+++ b/machinery/src/database/main.go
@@ -15,12 +15,19 @@ type DB struct {
 	Client *mongo.Client
 }
 
+var TIMEOUT = 10 * time.Second
 var _init_ctx sync.Once
 var _instance *DB
-var DatabaseName = "KerberosFactory"
 
-func New() *mongo.Client {
+var DatabaseName = os.Getenv("MONGODB_DATABASE_STORAGE")
 
+func New() *DB {
+
+	if DatabaseName == "" {
+		DatabaseName = "KerberosFactory"
+	}
+
+	mongodbURI := os.Getenv("MONGODB_URI")
 	host := os.Getenv("MONGODB_HOST")
 	databaseCredentials := os.Getenv("MONGODB_DATABASE_CREDENTIALS")
 	replicaset := os.Getenv("MONGODB_REPLICASET")
@@ -28,28 +35,46 @@ func New() *mongo.Client {
 	password := os.Getenv("MONGODB_PASSWORD")
 	authentication := "SCRAM-SHA-256"
 
-	_init_ctx.Do(func() {
-		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		defer cancel()
+	ctx, cancel := context.WithTimeout(context.Background(), TIMEOUT)
+	defer cancel()
 
+	_init_ctx.Do(func() {
 		_instance = new(DB)
-		mongodbURI := fmt.Sprintf("mongodb://%s:%s@%s", username, password, host)
-		if replicaset != "" {
-			mongodbURI = fmt.Sprintf("%s/?replicaSet=%s", mongodbURI, replicaset)
-		}
 
-		client, err := mongo.Connect(ctx, options.Client().ApplyURI(mongodbURI).SetAuth(options.Credential{
-			AuthMechanism: authentication,
-			AuthSource:    databaseCredentials,
-			Username:      username,
-			Password:      password,
-		}))
-		if err != nil {
-			fmt.Printf("Error setting up mongodb connection: %+v\n", err)
-			os.Exit(1)
+		// We can also apply the complete URI
+		// e.g. "mongodb+srv://<username>:<password>@kerberos-hub.shhng.mongodb.net/?retryWrites=true&w=majority&appName=kerberos-hub"
+		if mongodbURI != "" {
+			serverAPI := options.ServerAPI(options.ServerAPIVersion1)
+			opts := options.Client().ApplyURI(mongodbURI).SetServerAPIOptions(serverAPI)
+
+			// Create a new client and connect to the server
+			client, err := mongo.Connect(ctx, opts)
+			if err != nil {
+				fmt.Printf("Error setting up mongodb connection: %+v\n", err)
+				os.Exit(1)
+			}
+			_instance.Client = client
+
+		} else {
+
+			// New MongoDB driver
+			mongodbURI := fmt.Sprintf("mongodb://%s:%s@%s", username, password, host)
+			if replicaset != "" {
+				mongodbURI = fmt.Sprintf("%s/?replicaSet=%s", mongodbURI, replicaset)
+			}
+			client, err := mongo.Connect(ctx, options.Client().ApplyURI(mongodbURI).SetAuth(options.Credential{
+				AuthMechanism: authentication,
+				AuthSource:    databaseCredentials,
+				Username:      username,
+				Password:      password,
+			}))
+			if err != nil {
+				fmt.Printf("Error setting up mongodb connection: %+v\n", err)
+				os.Exit(1)
+			}
+			_instance.Client = client
 		}
-		_instance.Client = client
 	})
 
-	return _instance.Client
+	return _instance
 }