diff --git a/core/database/utils/connect.py b/core/database/utils/connect.py new file mode 100644 index 0000000..98d4386 --- /dev/null +++ b/core/database/utils/connect.py @@ -0,0 +1,33 @@ +""" +Definition of function to connect to database +using environment variables. +""" +from __future__ import annotations + +import logging +import os + +from dotenv import load_dotenv +from pymongo import MongoClient + + +def connect(): + """Connect to MongoDB using env vars.""" + # load env vars + load_dotenv() + necessary_env_vars = [ + 'MONGO_HOST', + 'MONGO_DB', + 'MONGO_COLLECTION', + ] + for env_var in necessary_env_vars: + assert env_var in os.environ, f"{env_var} not found" + # connect to database + client = MongoClient(os.getenv('MONGO_HOST')) + db = client[os.getenv('MONGO_DB')] + # extract collection + collection = db[os.getenv('MONGO_COLLECTION')] + # set unique index on "name" + collection.create_index('name', unique=True) + logging.debug('finished') + return collection, db, client