"""Database utilities."""

from contextlib import contextmanager

import sqlalchemy

from asset_file_details import config


def create_engine(db_url):
    """Create engine based on configuration settings."""
    if config.POOL_CLASS == sqlalchemy.pool.QueuePool:
        return sqlalchemy.create_engine(
            db_url,
            pool_size=config.POOL_SIZE,
            max_overflow=config.POOL_MAX_OVERFLOW,
            pool_recycle=config.POOL_RECYCLE_MS,
            pool_pre_ping=config.POOL_PRE_PING)

    return sqlalchemy.create_engine(db_url, poolclass=config.POOL_CLASS)


@contextmanager
def create_session(session_maker):
    """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.

    Args:
        session_maker (sqlalchemy.orm.session.sessionmaker): session factory
    Usage:
        with create_session(session_maker) as session:
            session.execute(query)
    """
    session = session_maker()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()