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" +}