28 lines
704 B
Python
28 lines
704 B
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
|
|
from dotenv import load_dotenv
|
|
from pymongo import MongoClient
|
|
|
|
|
|
def connect():
|
|
"""Connect to MongoDB."""
|
|
# 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')]
|
|
collection = db[os.getenv('MONGO_COLLECTION')]
|
|
collection.create_index('name', unique=True)
|
|
logging.info('connected to database')
|
|
return collection, db, client
|