import contextlib
import os
import traceback

import pymongo
import pymongo.errors


@contextlib.contextmanager
def die_on_mongo_connection_errors():
    try:
        yield
    except pymongo.errors.ServerSelectionTimeoutError:
        traceback.print_exc()
        os.abort()


def open_mongo_or_die() -> pymongo.MongoClient:
    client = pymongo.MongoClient(host=os.environ['MONGODB_SERVICE_HOST'],
                                 tz_aware=True,
                                 connectTimeoutMS=3000,
                                 serverSelectionTimeoutMS=3000)
    with die_on_mongo_connection_errors():
        client.admin.command('ismaster')
    return client