"""Snowflake connector class for the python integration tests."""

from snowflake_connector.etl_connector import SQLLoader
from snowflake_connector.etl_connector import SnowflakeSQLExecutor


# Load SQL templates
sql_loader = SQLLoader(__file__)


class IntegrationTestSFExecutor(SnowflakeSQLExecutor):
    """Helper class to abstract Snowflake operations.

    This class inherits from SnowflakeSQLExecutor class, which provides basic
    set of methods. This class extends SnowflakeSQLExecutor with some specific
    methods, which are useful to encapsulate some flow specific operations.
    """

    def get_number_of_rows(self, table_name, **kwargs):
        """Load temp staging raw table.

        Args:
            table_name (str): A table name in Snowflake.
            kwargs (dict): Dictionary with additional filtering params.
        """
        filter_clause = ''
        for key in kwargs:
            filter_clause += ' AND {key} = %({key})s'.format(key=key)
        params = dict(
            db=self.sf_config['db'],
            schema=self.sf_config['schema'],
            table_name=table_name,
            **kwargs)

        sql_template = sql_loader.load_query('get_number_of_rows').format(
            filter_clause=filter_clause)
        sql, non_identifier_params = self.validator.format_identifiers(
            sql_template, params)
        return self.fetchall(sql, params=non_identifier_params)

    def load_table(self, table_name, aws, s3_path):
        """Load temp staging raw table.

        Args:
            table_name (str): A table name in Snowflake.
            aws (dict): AWS credentials to fill a template of COPY SQL
                statement.
            s3_path (str): S3 path for coping files.
        """
        self.execute_query(
            sql_loader, 'load_table', params=dict(
                db=self.sf_config['db'],
                schema=self.sf_config['schema'],
                table_name=table_name,
                s3_path=s3_path,
                aws_key_id=aws['access_key'],
                aws_secret_key=aws['access_secret']))