"""Snowflake logic."""
import glob

import snowflake.connector as sc
import sqlparse
from lambdacommon.common_config import logger
import config


def run_scripts(scripts_dir, username, private_key, private_key_pwd):
    """Run a set of SQL scripts from a directory.

    Args:
        db_credentials: credentials which identify the database to run the
            scripts against
        scripts_dir: the directory containing the scripts

    Returns:
        The number of scripts run
    """
    connection = get_connection(username, private_key, private_key_pwd)

    with connection.cursor() as cursor:
        scripts = sorted(glob.glob(f'{scripts_dir}/*.sql'))
        logger.info(f'Found {len(scripts)} scripts in {scripts_dir}')

        for script in scripts:
            logger.info(f'Running SQL script {script}')
            run_script(cursor, script)

    connection.close()
    return len(scripts)


def run_script(cursor, script_path):
    """Run a single SQL script in Snowflake."""
    with open(script_path) as script:
        script_body = script.read()
        commands = sqlparse.split(script_body)

        for command in commands:
            logger.info(f'Running command: {command}')
            cursor.execute(command)


def get_connection(username, private_key, private_key_pwd):
    """Establish a connection to Snowflake using key pair authentication."""
    return sc.connect(
        account=config.SNOWFLAKE_ACCOUNT,
        user=username,
        private_key_file=private_key,
        private_key_file_pwd=private_key_pwd,
    )