"""MySQL Connector.

Manages interactions with MySQL.
"""

from contextlib import contextmanager

from sqlalchemy import create_engine, orm
from sqlalchemy.ext import declarative
from sqlalchemy.pool import QueuePool

from permissions import config

Base = declarative.declarative_base()

# please don't use the following private variables directly;
# use db_session
if config.POOL_CLASS == QueuePool:
    _transaction_db_engine = create_engine(
        config.AR_DB_URL,
        pool_size=config.POOL_SIZE,
        max_overflow=config.POOL_MAX_OVERFLOW,
        pool_recycle=config.POOL_RECYCLE_MS,
        connect_args=config.CONNECT_ARGS,
    )
else:
    _transaction_db_engine = create_engine(
        config.AR_DB_URL, connect_args=config.CONNECT_ARGS, poolclass=config.POOL_CLASS
    )
_db_transaction_session = orm.sessionmaker(bind=_transaction_db_engine)


if config.ENVIRONMENT == config.TEST_ENVIRONMENT:
    _read_db_engine = _transaction_db_engine
else:
    # Separate pool of DB connections for SELECT queries
    # These connections do not have the overhead of adding a commit & rollback
    # statement before returning the connection to the Pool. Because they don't
    # perform a commit/rollback to  release locks, we set autocommit=1 &
    # isolation_level='AUTOCOMMIT' to avoid creating any READ locks in the
    # first place as a guard against deadlocks
    _read_db_engine = create_engine(
        config.AR_DB_URL,
        pool_size=config.POOL_SIZE,
        max_overflow=config.POOL_MAX_OVERFLOW,
        pool_recycle=config.POOL_RECYCLE_MS,
        connect_args={**config.CONNECT_ARGS, 'autocommit': 1},
        isolation_level='AUTOCOMMIT',  # avoid READ locks
        pool_reset_on_return=None,
    )  # avoid rollback when connection returns
_db_read_session = orm.sessionmaker(bind=_read_db_engine, autoflush=False, expire_on_commit=False)


@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_transaction_session()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager
def db_read_session():
    """Provide a read scope around a series of operations.

    This handles closing of session to return connection to pool.

    No commit or rollback so use only for SELECT queries.

    Usage:
        with db_session() as session:
            session.execute(query)
    """
    session = _db_read_session()
    try:
        yield session
    except Exception:
        session.rollback()  # safeguard against accidental locks
        raise
    finally:
        session.close()