from contextlib import contextmanager
from functools import wraps
import logging as logger

import snowflake.connector
from snowflake.connector import DictCursor

import config

connection = snowflake.connector.connect(
        user=config.SNOWFLAKE_USER,
        password=config.SNOWFLAKE_PASSWORD,
        account=config.SNOWFLAKE_ACCOUNT,
        database=config.SNOWFLAKE_CATALOG,
        schema=config.SNOWFLAKE_SCHEMA,
        **config.SNOWFLAKE_CONNECT_ARGS)


def get_cursor():
    """Get a snowflake cursor."""
    return connection.cursor()


def close_connection():
    """Close snowflake connection."""
    connection.close()


@contextmanager
def snow_session():
    """Provide a transactional scope around a series of operations.

    Usage:
        with snow_session() as session:
            session.execute(query)
    """
    cursor = connection.cursor(DictCursor)
    try:
        cursor.execute('BEGIN;')
        yield cursor
        cursor.execute('COMMIT;')
    except Exception as e:
        logger.error(
            '------------------------------------------------------------\n'
            '-------------------------  DB ERROR  -----------------------\n'
            '------------------------------------------------------------')
        logger.error(str(e))
        cursor.execute('ROLLBACK;')
        raise
    finally:
        cursor.close()


def snow_session_wrap(func):
    """DB Session Wrapper.

    Creates a new session if one isn't passed in.
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        cursor = kwargs.pop('cursor', None)
        try:
            if cursor:
                result = func(*args, cursor=cursor, **kwargs)
            else:
                with snow_session() as cursor:
                    result = func(*args, cursor=cursor, **kwargs)
            return result
        except Exception:
            raise

    return wrapper