"""MySQL connector."""

from contextlib import contextmanager
import functools

from sqlalchemy import create_engine
from sqlalchemy import exc
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import NullPool, StaticPool

import common_config as config

DB_POOLCLASS = NullPool
# Test
if config.ENVIRONMENT == config.ENVIRONMENT_TEST:
    # Override db connection parameters not to wait for network timeouts on
    # missing credentials. Such timeouts might be quite high in some
    # environments.
    DB_POOLCLASS = StaticPool


BaseModel = declarative_base()

ar_db_engine = create_engine(
    config.AR_DB_CONNECTION_STRING,
    connect_args=config.DB_CONNECT_ARGS,
    poolclass=DB_POOLCLASS)


_ar_db_session = sessionmaker(bind=ar_db_engine, expire_on_commit=False)


@contextmanager
def _db_session_manager(session_factory):
    session = session_factory()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def ar_db_session():
    """Provide a transactional scope around a series of operations for AR DB.

    Usage:
        with ar_db_session() as session:
            session.execute(query)
    """
    return _db_session_manager(_ar_db_session)


def wrap_db_errors(function):
    """Decorate the given function with logic to handle SQLAlchemy errors.

    If a SQLAlchemy exception is thrown, it will be caught and logged and the
    function will return a fatal response.

    Args:
        function (func): the function to decorate
    Returns:
        func: function decorated with error-handling logic
    """
    @functools.wraps(function)
    def call_function_with_error_handling(*args, **kwargs):
        try:
            function_return = function(*args, **kwargs)
        except exc.SQLAlchemyError as exception:
            config.logger.exception(str(exception))
            raise exception
        return function_return
    return call_function_with_error_handling