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