"""MySQL Connector."""

from contextlib import contextmanager

from sqlalchemy import create_engine
from sqlalchemy.orm import declarative_base, sessionmaker
from sqlalchemy.pool import QueuePool

from ows_product_physical import config

# please don't use the following private variables directly;
# use db_session, delivery_db_session
if config.POOL_CLASS == QueuePool:
    _db_engine = create_engine(
        config.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)
    _delivery_db_engine = create_engine(
        config.DELIVERY_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)
else:
    _db_engine = create_engine(
        config.DB_URL, poolclass=config.POOL_CLASS)
    _delivery_db_engine = create_engine(
        config.DELIVERY_DB_URL, poolclass=config.POOL_CLASS)

# session maker for each DB
_db_session = sessionmaker(bind=_db_engine)
_delivery_db_session = sessionmaker(bind=_delivery_db_engine)

BaseModel = declarative_base()
DeliveryBaseModel = declarative_base()


@contextmanager
def db_session():
    """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 db_session() as session:
        session.execute(query)
    """
    session = _db_session()
    try:
        yield session
        session.commit()
    except:  # noqa
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager
def delivery_db_session():
    """Provide a transactional scope around delivery db operations.

    with delivery_db_session() as session:
        session.execute(query)
    """
    session = _delivery_db_session()
    try:
        yield session
        session.commit()
    except Exception as err:
        session.rollback()
        raise err
    finally:
        session.close()


def get_delivery_db_engine():
    """Get delivery DB engine."""
    return _delivery_db_engine