diff --git a/code/database.py b/code/database.py new file mode 100644 index 0000000..e49e705 --- /dev/null +++ b/code/database.py @@ -0,0 +1,31 @@ +from pymongo import MongoClient +from dotenv import load_dotenv +import logging +import os + +def connect( + ip_addr: str | None = None, + db_name: str | None = None, + collection_name: str | None = None +): + """ + Connect to MongoDB database and return collection. + """ + # ensure variables are set + if ip_addr is None: + ip_addr = os.getenv('DB_IP_ADDRESS') + if db_name is None: + db_name = os.getenv('DB_NAME') + if collection_name is None: + collection_name = os.getenv('DB_COLLECTION_NAME') + # connect to database + client = MongoClient(ip_addr) + db = client[db_name] + collection = db[collection_name] + return collection + + +if __name__ == '__main__': + load_dotenv() + collection = connect() + print(collection) \ No newline at end of file