Define a context manager generator

generator_contextmanager.py

from contextlib import contextmanager

@contextmanager
def connect_to_db(address):
    db = CrappyDBConnection(address)
    try:
        yield db
    except ConnectionError:
        logging.exception('Connection dropped')
        db.cleanup('rollback')
    else:
        db.cleanup('commit')
    db.disconnect()

# USE LIKE
with connect_to_db(address) as db:
    ...