"""Connector for Snowflake Database.

Connects to a Snowflake database
"""

from typing import Any

from snowflake import connector
from snowflake.connector import SnowflakeConnection

from gridgen import config
from gridgen.connectors import helper


def get_connection() -> SnowflakeConnection:
    """Get an established snowflake connection.

    Returns:
        snowflake.connector.Connection: snowflake connection.
    """
    return connector.connect(
        user=config.SNOWFLAKE_CONFIG["user"],
        account=config.SNOWFLAKE_CONFIG["account"],
        database=config.SNOWFLAKE_CONFIG["database"],
        schema=config.SNOWFLAKE_CONFIG["schema"],
        warehouse=config.SNOWFLAKE_CONFIG["warehouse"],
        role=config.SNOWFLAKE_CONFIG["role"],
        private_key=config.SNOWFLAKE_CONFIG["private_key"],
    )


def execute_query(query: str, params: tuple[Any, ...] = ()) -> Any:
    """Execute a query on snowflake connection and return results.

    Args:
        query (str): query to execute on database
        params (tuple): parameters passed in to query

    Returns:
        tuple: query results
    """
    conn = get_connection()
    try:
        return helper.execute_query(conn, query, params)
    finally:
        conn.close()