"""Connector for Snowflake Database."""

from typing import Any

import config
import snowflake.connector


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

    Returns:
        snowflake.connector.Connection: snowflake connection.
    """
    return snowflake.connector.connect(
        user=config.SNOWFLAKE_USER,
        account=config.SNOWFLAKE_ACCOUNT,
        database=config.SNOWFLAKE_DATABASE,
        schema=config.SNOWFLAKE_SCHEMA,
        warehouse=config.SNOWFLAKE_WAREHOUSE,
        role=config.SNOWFLAKE_ROLE,
        private_key=config.SNOWFLAKE_PRIVATE_KEY,
    )


def execute_query(query: str, params: tuple[Any, ...] = ()) -> Any:
    """Execute a query on a 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()
    db_cursor = None
    try:
        db_cursor = conn.cursor()
        db_cursor.execute(query, params)
        results = db_cursor.fetchall()
        conn.commit()
        return results
    except Exception:
        conn.rollback()
        raise
    finally:
        if db_cursor:
            db_cursor.close()
        conn.close()