"""Snowflake connector."""
import base64

import pandas as pd
import snowflake.connector
import config
from src.connectors.secrets_manager import SecretsManagerClient


class SnowflakeClient:
    """Snowflake client."""

    def __init__(self):
        """Initialize SnowflakeClient."""
        secrets_manager_client = SecretsManagerClient()
        private_key = secrets_manager_client.get_secret(
            config.SNOWFLAKE_PRIVATE_KEY_SECRET)
        self.conn = snowflake.connector.connect(
            user=config.SNOWFLAKE_USER,
            account=config.SNOWFLAKE_ACCOUNT,
            role=config.SNOWFLAKE_ROLE,
            private_key=base64.b64decode(
                bytes(private_key, encoding='utf-8')),
            warehouse=config.SNOWFLAKE_WAREHOUSE,
        )

    def execute_query(self, query):
        """Execute query."""
        cursor = self.conn.cursor()
        cursor.execute(query)
        column_names = [desc[0] for desc in cursor.description]
        df = pd.DataFrame(cursor, columns=column_names)
        cursor.close()
        return df

    def close_connection(self):
        """Close connection."""
        self.conn.close()