A context manager to handle a MongoDB Collection

mongo_collection_context_manager.py

'''
A context manager for handling a PyMongo driver MongoDB connection
to access a collection.
'''
from contextlib import contextmanager
from pymongo import MongoClient
import traceback


@contextmanager
def mongo_collection(uri: str, timeout_ms: int, collection: str, db_name=None):
    '''Open up a MongoDB connection, grab a db, then get a collection.'''
    try:
        mongo_client = MongoClient(uri, timeout_ms)
        if db_name is not None:
            db = mongo_client[db_name]
        else:
            db = mongo_client.get_default_database()
        yield db.get_collection(collection)
    except BaseException as err:
        traceback.print_exc()
    finally:
        mongo_client.close()


# Example:
# with mongo_collection(uri, 1000, collection) as c:
#     results = c.find_all('{foo.bar}')