From f2739523827d4851a973d1d25c8d99dcb9611feb Mon Sep 17 00:00:00 2001
From: eureka320 <elias.h0320@gmail.com>
Date: Tue, 26 Mar 2024 22:28:41 -0400
Subject: [PATCH] feat: implemented the mongodb connection feature

---
 .env.example                                  |  3 +-
 mongodb.py                                    | 49 ++++++++++++++
 requirements.txt                              |  4 +-
 src/mongodb/MongoDBClass.py                   | 65 +++++++++++++++++++
 .../payload/mongodb_payload.json              |  5 ++
 5 files changed, 124 insertions(+), 2 deletions(-)
 create mode 100644 mongodb.py
 create mode 100644 src/mongodb/MongoDBClass.py
 create mode 100644 test/regression/regression_test013/payload/mongodb_payload.json

diff --git a/.env.example b/.env.example
index e759c6e..8819bff 100644
--- a/.env.example
+++ b/.env.example
@@ -1,3 +1,4 @@
 OPENAI_API_KEY=
 MATHPIX_APP_ID=
-MATHPIX_APP_KEY=
\ No newline at end of file
+MATHPIX_APP_KEY=
+Mongo_URI=
\ No newline at end of file
diff --git a/mongodb.py b/mongodb.py
new file mode 100644
index 0000000..e3158c5
--- /dev/null
+++ b/mongodb.py
@@ -0,0 +1,49 @@
+
+import os
+import gc
+import argparse
+from pathlib import Path
+
+from src.utils.read_json import read_json
+from src.mongodb.MongoDBClass import MongoDBClass
+
+def mongodb(args):
+    """
+    main entry point
+    """
+
+    # Payload
+    payload_data = read_json(args.payload_dir)
+
+    # Call class instance
+    mongodb = MongoDBClass(
+        db_name=payload_data["db_name"], 
+        collection_name=payload_data["collection_name"], 
+        mongo_uri=payload_data["mongo_uri"])
+    
+    mongodb.mongo_connect()
+
+    gc.collect()
+
+if __name__ == "__main__":
+    """
+    Form command lines
+    """
+    # Clean up buffer memory
+    gc.collect()
+
+    # Current directory
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+
+    # Payload directory
+    test_name    = "regression_test013"
+    payload_name = "mongodb_payload.json"
+    payload_dir  = os.path.join(current_dir, "test", "regression", test_name, "payload", payload_name)
+
+    # Add options
+    p = argparse.ArgumentParser()
+    p = argparse.ArgumentParser(description="Translate text within an image.")
+    p.add_argument("--payload_dir", type=Path, default=payload_dir, help="payload directory to the test example")
+    args = p.parse_args()
+
+    mongodb(args)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 1f2d5ab..3723113 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -21,4 +21,6 @@ datasets
 docx2txt
 cryptography
 poppler-utils
-PyMuPDF
\ No newline at end of file
+PyMuPDF
+pyjwt
+pymongo
\ No newline at end of file
diff --git a/src/mongodb/MongoDBClass.py b/src/mongodb/MongoDBClass.py
new file mode 100644
index 0000000..a3c8bb8
--- /dev/null
+++ b/src/mongodb/MongoDBClass.py
@@ -0,0 +1,65 @@
+import openai
+import re
+import os
+import datetime
+import jwt
+import json
+import pytz
+from dotenv import load_dotenv
+from bson.objectid import ObjectId
+from bson.json_util import dumps
+from pymongo import MongoClient
+
+class MongoDBClass():
+    def __init__(self, db_name, collection_name, mongo_uri=""):
+        self.db_name = db_name
+        self.collection_name = collection_name
+        self.set_mongo_uri(mongo_uri)
+
+    def set_mongo_uri(self, mongo_uri):
+        if mongo_uri:
+            self.mongo_uri = mongo_uri
+        else:
+            load_dotenv()
+            self.mongo_uri = os.getenv("Mongo_URI")
+
+        if self.mongo_uri is not None:
+            os.environ["Mongo_URI"] = self.mongo_uri
+            return True
+        else:
+            # Handle the absence of the environment variable
+            # You might want to log an error, raise an exception, or provide a default value
+            # For example, setting a default value
+            os.environ["Mongo_URI"] = mongo_uri
+            return False
+
+    def mongo_connect(self):
+        # mongo config
+        client = MongoClient(self.mongo_uri)
+        if client is None:
+            # no connection, exit early
+            print("MongoDB Server Connection Error")
+            return {"result": False ,"message": "MongoDB Server Connection Error."}
+        else:
+            print("MongoDB Server Connected.")
+            # confirm oridosai db exist
+            dblist = client.list_database_names()
+            if self.db_name in dblist:
+                print(f"The {self.db_name} database exists.")
+
+                # confirm apis collection exist
+                oridosai_db = client[self.db_name]
+                collist = oridosai_db.list_collection_names()
+                if self.collection_name in collist:
+                    print(f"The {self.collection_name} collection exists.")
+                    user_col = oridosai_db[self.collection_name]
+                    return {"result": True, "message": user_col}
+                else:
+                    print(f"The {self.collection_name} collection not exists.")
+                    return {"result": False ,"message": f"The {self.collection_name} collection not exists."}
+            else:
+                print(f"The {self.db_name} database does not exist.")
+                # create the database
+                client[self.db_name].create_collection(self.collection_name)
+                print(f"The {self.db_name} database has been created with {self.collection_name} collection.")
+                return {"result": True, "message": f"The {self.db_name} database has been created with {self.collection_name} collection."}
\ No newline at end of file
diff --git a/test/regression/regression_test013/payload/mongodb_payload.json b/test/regression/regression_test013/payload/mongodb_payload.json
new file mode 100644
index 0000000..dbae6ab
--- /dev/null
+++ b/test/regression/regression_test013/payload/mongodb_payload.json
@@ -0,0 +1,5 @@
+{
+    "mongo_uri": null,
+    "db_name": "oridosai",
+    "collection_name": "apis"
+}