from utils.snowflake.snowflake_client import get_snowflake_session


# Get data from Snowflake
def get_data_from_snowflake(warehouse, database_name, schema, query_sql):
    session = get_snowflake_session(warehouse, database_name, schema)
    try:
        print(f"Executing Snowflake query for {query_sql}")
        snowflake_df = session.sql(query_sql)
        data = snowflake_df.collect()
        if not data:
            print(f"No data found for {query_sql}")
            raise Exception(f"No data found for {query_sql}")
        row_count = len(data)
        print(f"Retrieved {row_count:,} rows from Snowflake for {query_sql}")
        return data
    except Exception as e:
        print(f"Error retrieving data from Snowflake for {query_sql}: {e}")
        raise e
    finally:
        session.close()