From ec4219e5a54bea78ff19f23f1754a036e8d0eae3 Mon Sep 17 00:00:00 2001
From: zevaryx <zevaryx@gmail.com>
Date: Tue, 21 Mar 2023 12:52:12 -0600
Subject: [PATCH] Add support for replicated hosts

---
 jarvis_core/db/__init__.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/jarvis_core/db/__init__.py b/jarvis_core/db/__init__.py
index 5a1e219..7534ae5 100644
--- a/jarvis_core/db/__init__.py
+++ b/jarvis_core/db/__init__.py
@@ -14,7 +14,13 @@ CTC2_INST = MotorAsyncIOInstance()
 
 
 def connect(
-    host: str, username: str, password: str, port: int = 27017, testing: bool = False
+    username: str,
+    password: str,
+    port: int = 27017,
+    testing: bool = False,
+    host: str = None,
+    hosts: list[str] = None,
+    replicaset: str = None,
 ) -> None:
     """
     Connect to MongoDB.
@@ -26,9 +32,15 @@ def connect(
         port: Port
     """
     global CLIENT, JARVISDB, CTC2DB, JARVIS_INST, CTC2_INST
-    CLIENT = AsyncIOMotorClient(
-        host=host, username=username, password=password, port=port, tz_aware=True, tzinfo=utc
-    )
+
+    if not replicaset:
+        CLIENT = AsyncIOMotorClient(
+            host=host, username=username, password=password, port=port, tz_aware=True, tzinfo=utc
+        )
+    else:
+        CLIENT = AsyncIOMotorClient(
+            hosts, username=username, password=password, tz_aware=True, tzinfo=utc, replicaset=replicaset
+        )
     JARVISDB = CLIENT.narvis if testing else CLIENT.jarvis
     CTC2DB = CLIENT.ctc2