"""Snowflake connector."""

import snowflake.connector
from snowflake.connector import errors

from accounting import config


def get_snowflake_connection():
    """Connect to Snowflake database.

    Returns:
        SnowflakeConnection: Snowflake connection object

    """
    try:
        # TODO: Remove pw fallback once key is confirmed to be working
        if (config.private_key):
            con = snowflake.connector.connect(
                account=config.SNOWFLAKE_ACCOUNT,
                user=config.SNOWFLAKE_USER,
                private_key=config.SNOWFLAKE_PRIVATE_KEY,
                database=config.SNOWFLAKE_DB,
                schema=config.SNOWFLAKE_SCHEMA,
                role=config.SNOWFLAKE_ROLE,
                warehouse=config.SNOWFLAKE_WAREHOUSE
            )
        else:
            con = snowflake.connector.connect(
                account=config.SNOWFLAKE_ACCOUNT,
                user=config.SNOWFLAKE_USER,
                password=config.SNOWFLAKE_PASSWORD,
                database=config.SNOWFLAKE_DB,
                schema=config.SNOWFLAKE_SCHEMA,
                role=config.SNOWFLAKE_ROLE,
                warehouse=config.SNOWFLAKE_WAREHOUSE
            )
        return con
    except errors.Error as ex:
        config.log.exception('Failed to connect to Snowflake.')
        raise ex