"""DB Connector. Manages interactions with Snowflake."""

from contextlib import contextmanager

from sqlalchemy import create_engine
from sqlalchemy import event
from sqlalchemy import exc
from sqlalchemy import pool
from sqlalchemy import select
from sqlalchemy.orm import sessionmaker

from ows_accounting import config

role = config.SNOWFLAKE_ROLE
database = config.SNOWFLAKE_DATABASE
schema = config.SNOWFLAKE_SCHEMA
warehouse = config.SNOWFLAKE_WAREHOUSE


def _create_engine(db_url):
    """Create engine based on configuration settings."""
    if config.SNOWFLAKE_POOL_CLASS == pool.QueuePool:
        db_engine = create_engine(
            db_url, pool_size=config.SNOWFLAKE_POOL_SIZE,
            connect_args=config.SNOWFLAKE_CONNECT_ARGS,
            max_overflow=config.SNOWFLAKE_POOL_MAX_OVERFLOW,
            pool_recycle=config.SNOWFLAKE_POOL_RECYCLE)

        @event.listens_for(db_engine, 'connect')
        def set_snowflake_params(connection, connection_record):
            """Set session params.

            When using raw SQL with fully qualified table names
            (db.schema.table), USE DATABASE and USE SCHEMA are not required.
            """

            with connection.cursor() as cur:
                if role:
                    cur.execute('USE ROLE {};'.format(role))
                if database:
                    cur.execute('USE DATABASE {};'.format(database))
                if schema:
                    cur.execute('USE SCHEMA {};'.format(schema))
                if warehouse:
                    cur.execute('USE WAREHOUSE {};'.format(warehouse))

        if config.SNOWFLAKE_POOL_PRE_PING:
            @event.listens_for(db_engine, 'engine_connect')
            def ping_connection(connection, branch):
                _ping_connection(connection, branch)

        return db_engine
    else:
        return create_engine(
            db_url,
            connect_args=config.SNOWFLAKE_CONNECT_ARGS,
            poolclass=config.SNOWFLAKE_POOL_CLASS)


def _ping_connection(connection, branch):
    """Ping database connection after engine_connect event.

    This function is copied verbatim from
    http://docs.sqlalchemy.org/en/latest/core/pooling.html
    """
    if branch:
        return

    save_should_close_with_result = connection.should_close_with_result
    connection.should_close_with_result = False

    try:
        connection.scalar(select([1]))
    except exc.DBAPIError as err:
        if err.connection_invalidated:
            connection.scalar(select([1]))
        else:
            raise
    finally:
        connection.should_close_with_result = save_should_close_with_result


# Do not use these variables directly other than running unit tests
snowflake_db_engine = _create_engine(config.SNOWFLAKE_DB_URL)

# please don't use sessions directly; instead use db_session
sessions = {
    'snowflake': sessionmaker(bind=snowflake_db_engine),
}


@contextmanager
def db_session(db_name='snowflake'):
    """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:
        db_name (session): Name of the session
    Usage:
        with db_session() as session:
            session.execute(query)
    """
    session = sessions[db_name]()
    try:
        yield session
        session.commit()
    except: # noqa
        session.rollback()
        raise
    finally:
        session.close()