defined database connect function

This commit is contained in:
Brian Bjarke Jensen
2024-03-31 12:17:59 +02:00
parent a49e55c159
commit e09188453e
+33
View File
@@ -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