From 06b8c824844bd543138f7d68ab47e8e7adfa1aac Mon Sep 17 00:00:00 2001
From: Shadow <hi@shadowing.dev>
Date: Wed, 18 Dec 2024 11:50:30 -0600
Subject: [PATCH] feat: add a toggle for replica sets to be used or not

---
 .../dashboard/project/add-database.tsx        | 27 ++++++++++
 packages/server/src/db/schema/mongo.ts        |  5 +-
 packages/server/src/utils/databases/mongo.ts  | 50 +++++++++++--------
 3 files changed, 60 insertions(+), 22 deletions(-)

diff --git a/apps/dokploy/components/dashboard/project/add-database.tsx b/apps/dokploy/components/dashboard/project/add-database.tsx
index aaf4940b..98e3019a 100644
--- a/apps/dokploy/components/dashboard/project/add-database.tsx
+++ b/apps/dokploy/components/dashboard/project/add-database.tsx
@@ -35,6 +35,7 @@ import {
 	SelectTrigger,
 	SelectValue,
 } from "@/components/ui/select";
+import { Switch } from "@/components/ui/switch";
 import { Textarea } from "@/components/ui/textarea";
 import { slugify } from "@/lib/slug";
 import { api } from "@/utils/api";
@@ -95,6 +96,7 @@ const mySchema = z.discriminatedUnion("type", [
 		.object({
 			type: z.literal("mongo"),
 			databaseUser: z.string().default("mongo"),
+			replicaSets: z.boolean().default(false),
 		})
 		.merge(baseDatabaseSchema),
 	z
@@ -216,6 +218,7 @@ export const AddDatabase = ({ projectId, projectName }: Props) => {
 				databaseUser:
 					data.databaseUser || databasesUserDefaultPlaceholder[data.type],
 				serverId: data.serverId,
+				replicaSets: data.replicaSets,
 			});
 		} else if (data.type === "redis") {
 			promise = redisMutation.mutateAsync({
@@ -540,6 +543,30 @@ export const AddDatabase = ({ projectId, projectName }: Props) => {
 										);
 									}}
 								/>
+
+								{type === "mongo" && (
+									<FormField
+										control={form.control}
+										name="replicaSets"
+										render={({ field }) => {
+											return (
+												<FormItem>
+													<FormLabel>Use Replica Sets</FormLabel>
+													<FormControl>
+														<Switch
+															checked={field.value}
+															onCheckedChange={field.onChange}
+															disabled
+															aria-readonly
+														/>
+													</FormControl>
+
+													<FormMessage />
+												</FormItem>
+											);
+										}}
+									/>
+								)}
 							</div>
 						</div>
 					</form>
diff --git a/packages/server/src/db/schema/mongo.ts b/packages/server/src/db/schema/mongo.ts
index 757ba9c9..73f297b7 100644
--- a/packages/server/src/db/schema/mongo.ts
+++ b/packages/server/src/db/schema/mongo.ts
@@ -1,5 +1,5 @@
 import { relations } from "drizzle-orm";
-import { integer, pgTable, text } from "drizzle-orm/pg-core";
+import { boolean, integer, pgTable, text } from "drizzle-orm/pg-core";
 import { createInsertSchema } from "drizzle-zod";
 import { nanoid } from "nanoid";
 import { z } from "zod";
@@ -43,6 +43,7 @@ export const mongo = pgTable("mongo", {
 	serverId: text("serverId").references(() => server.serverId, {
 		onDelete: "cascade",
 	}),
+	replicaSets: boolean("replicaSets").default(false),
 });
 
 export const mongoRelations = relations(mongo, ({ one, many }) => ({
@@ -77,6 +78,7 @@ const createSchema = createInsertSchema(mongo, {
 	externalPort: z.number(),
 	description: z.string().optional(),
 	serverId: z.string().optional(),
+	replicaSets: z.boolean().default(false),
 });
 
 export const apiCreateMongo = createSchema
@@ -89,6 +91,7 @@ export const apiCreateMongo = createSchema
 		databaseUser: true,
 		databasePassword: true,
 		serverId: true,
+		replicaSets: true,
 	})
 	.required();
 
diff --git a/packages/server/src/utils/databases/mongo.ts b/packages/server/src/utils/databases/mongo.ts
index 3c9556f9..2cce315c 100644
--- a/packages/server/src/utils/databases/mongo.ts
+++ b/packages/server/src/utils/databases/mongo.ts
@@ -28,48 +28,56 @@ export const buildMongo = async (mongo: MongoNested) => {
 		databasePassword,
 		command,
 		mounts,
+		replicaSets,
 	} = mongo;
 
-	const initReplicaSet = `
+	const startupScript = `
 #!/bin/bash
+${
+	replicaSets
+		? `
 mongod --port 27017 --replSet rs0 --bind_ip_all &
 MONGOD_PID=$!
 
 # Wait for MongoDB to be ready
 while ! mongosh --eval "db.adminCommand('ping')" > /dev/null 2>&1; do
-  sleep 2
+	sleep 2
 done
 
 # Check if replica set is already initialized
 REPLICA_STATUS=$(mongosh --quiet --eval "rs.status().ok || 0")
 
 if [ "$REPLICA_STATUS" != "1" ]; then
-  echo "Initializing replica set..."
-  mongosh --eval '
-    rs.initiate({
-      _id: "rs0",
-      members: [{ _id: 0, host: "localhost:27017", priority: 1 }]
-    });
+	echo "Initializing replica set..."
+	mongosh --eval '
+	rs.initiate({
+		_id: "rs0",
+		members: [{ _id: 0, host: "localhost:27017", priority: 1 }]
+	});
 
     // Wait for the replica set to initialize
-    while (!rs.isMaster().ismaster) {
-      sleep(1000);
-    }
+	while (!rs.isMaster().ismaster) {
+		sleep(1000);
+	}
 
     // Create root user after replica set is initialized and we are primary
-    db.getSiblingDB("admin").createUser({
-      user: "${databaseUser}",
-      pwd: "${databasePassword}",
-      roles: ["root"]
-    });
-  '
+	db.getSiblingDB("admin").createUser({
+		user: "${databaseUser}",
+		pwd: "${databasePassword}",
+		roles: ["root"]
+	});
+	'
+
 else
-  echo "Replica set already initialized."
+	echo "Replica set already initialized."
 fi
+`
+		: "mongod --port 27017 --bind_ip_all & MONGOD_PID=$!"
+}
 
-wait $MONGOD_PID`;
+${command ?? "wait $MONGOD_PID"}`;
 
-	const defaultMongoEnv = `MONGO_INITDB_DATABASE=admin\n${env ? `${env}` : ""}`;
+	const defaultMongoEnv = `MONGO_INITDB_ROOT_USERNAME=${databaseUser}\nMONGO_INITDB_ROOT_PASSWORD=${databasePassword}\nMONGO_INITDB_DATABASE=admin\n${env ? `${env}` : ""}`;
 
 	const resources = calculateResources({
 		memoryLimit,
@@ -96,7 +104,7 @@ wait $MONGOD_PID`;
 				Env: envVariables,
 				Mounts: [...volumesMount, ...bindsMount, ...filesMount],
 				Command: ["/bin/bash"],
-				Args: ["-c", command ?? initReplicaSet],
+				Args: ["-c", startupScript],
 			},
 			Networks: [{ Target: "dokploy-network" }],
 			Resources: {