34 lines
848 B
Python
34 lines
848 B
Python
"""
|
|
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
|