"""Snowflake related tasks."""
import contextlib

from garcon import task
from job import config
from snowflake import connector


@contextlib.contextmanager
def get_extended_wh_cursor(logger, sf_config):
    """Get cursor with warehouse set to medium.

    Args:
        logger: logger object.
        sf_config (dict): Snowflake connection params.

    Returns:
        SnowflakeCursor: cursor object.
    """
    def get_wh_size(cur, wh_name):
        """Get warehouse size."""
        cur.execute("SHOW WAREHOUSES LIKE '{}';".format(wh_name))
        res = cur.fetchone()
        return res['size']

    warehouse_name = sf_config['warehouse']
    conn = connector.connect(
        user=sf_config['user'],
        password=sf_config['password'],
        account=sf_config['account'],
        role=sf_config['role'])
    try:
        cur = conn.cursor(connector.DictCursor)
        cur.execute('USE WAREHOUSE {};'.format(warehouse_name))
        cur.execute('USE DATABASE {};'.format(sf_config['db']))
        cur.execute('USE SCHEMA {};'.format(sf_config['schema']))

        old_size = get_wh_size(cur, warehouse_name)
        logger.info('Current wh size: %s', old_size)
        try:
            cur.execute(
                "ALTER WAREHOUSE {} SET WAREHOUSE_SIZE='MEDIUM';".format(
                    warehouse_name))
            logger.info(
                'wh size set to: %s', get_wh_size(cur, warehouse_name))
            yield cur

        finally:
            cur.execute("ALTER WAREHOUSE {} SET WAREHOUSE_SIZE='{}';".format(
                warehouse_name, old_size))
            logger.info(
                'wh size set back to: %s', get_wh_size(cur, warehouse_name))
            cur.close()

    finally:
        conn.close()


@task.decorate(timeout=7200)
def run_unload(activity, query, dest):
    """Unload data from Snowflake onto s3.

    Args:
        activity (ActivityWorker): The swf activity worker.
        query (str): the unload query.
        dest (str): destination of the query results.
    """
    activity.logger.info(
        'Start Snowflake unload to s3 (query is not displayed as '
        'it may contain sensitive information.)')

    sql = """
        COPY INTO {s3_path}
        FROM ({query})
        CREDENTIALS=(
            AWS_KEY_ID='{aws_key_id}'
            AWS_SECRET_KEY='{aws_secret_key}'
        )
        FILE_FORMAT=(
            TYPE='CSV'
            FIELD_DELIMITER=','
            RECORD_DELIMITER='\\n'
            SKIP_HEADER=1
            COMPRESSION='NONE'
            ESCAPE='\\\\'
            ESCAPE_UNENCLOSED_FIELD='\\\\'
            FIELD_OPTIONALLY_ENCLOSED_BY='"'
            NULL_IF=('')
            EMPTY_FIELD_AS_NULL=TRUE
        )
        OVERWRITE = TRUE
        SINGLE = TRUE
        MAX_FILE_SIZE = {max_size}
    """.format(s3_path=dest, query=query,
               aws_secret_key=config.aws['aws_access_secret'],
               aws_key_id=config.aws['aws_access_key'],
               max_size=config.s3_upload_max_size_bytes)

    with get_extended_wh_cursor(activity.logger, config.sf_config) as cur:
        cur.execute(sql)

    activity.logger.info('Unload data from Snowflake was successful.')