from contextlib import contextmanager

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

import config

_masters_registry_engine = create_engine(
    config.TARGET_DB_URL, poolclass=config.POOLCLASS)
_masters_registry_session = sessionmaker(bind=_masters_registry_engine)


@contextmanager
def mr_session_scope():
    """Provide a transactional scope around a series of operations.
    Taken from http://docs.sqlalchemy.org/en/latest/orm/session_basics.html.
    This handles rollback and closing of session, so there is no need
    to do that throughout the code. Usage:

    with mr_session_scope() as session:
        session.execute(query)
    """

    session = _masters_registry_session()
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()