Skip to content
This repository has been archived by the owner on Jul 12, 2024. It is now read-only.

Commit

Permalink
feat: implemented the mongodb connection feature
Browse files Browse the repository at this point in the history
  • Loading branch information
eureka320 committed Mar 27, 2024
1 parent 54bb18d commit f273952
Show file tree
Hide file tree
Showing 5 changed files with 124 additions and 2 deletions.
3 changes: 2 additions & 1 deletion .env.example
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
OPENAI_API_KEY=
MATHPIX_APP_ID=
MATHPIX_APP_KEY=
MATHPIX_APP_KEY=
Mongo_URI=
49 changes: 49 additions & 0 deletions mongodb.py
Original file line number Diff line number Diff line change
@@ -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)
4 changes: 3 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -21,4 +21,6 @@ datasets
docx2txt
cryptography
poppler-utils
PyMuPDF
PyMuPDF
pyjwt
pymongo
65 changes: 65 additions & 0 deletions src/mongodb/MongoDBClass.py
Original file line number Diff line number Diff line change
@@ -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."}
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
{
"mongo_uri": null,
"db_name": "oridosai",
"collection_name": "apis"
}

0 comments on commit f273952

Please sign in to comment.