""" Snowflake connection layer for the HelpScout processing module. Adapted from processing_comments/SnowFlakeConnection.py. """ import os from pathlib import Path from snowflake.snowpark import Session from dotenv import load_dotenv import logging logger = logging.getLogger(__name__) # Load .env from the project root (two levels up from this file) _root_env = Path(__file__).resolve().parent.parent / ".env" load_dotenv(dotenv_path=_root_env) class SnowflakeConn: """Thin wrapper around Snowpark Session for running read queries.""" def __init__(self): self.session = self._connect() # ------------------------------------------------------------------ def _connect(self) -> Session: conn_params = dict( user=os.getenv("SNOWFLAKE_USER"), password=os.getenv("SNOWFLAKE_PASSWORD"), account=os.getenv("SNOWFLAKE_ACCOUNT"), role=os.getenv("SNOWFLAKE_ROLE"), warehouse=os.getenv("SNOWFLAKE_WAREHOUSE"), # No default database/schema — queries use fully-qualified names ) session = Session.builder.configs(conn_params).create() logger.info("Snowflake session created successfully.") return session # ------------------------------------------------------------------ def run_query(self, query: str, description: str = "query"): """Execute a SELECT query and return a pandas DataFrame.""" try: df = self.session.sql(query).to_pandas() df.columns = df.columns.str.lower() logger.info(f"Query '{description}' returned {len(df):,} rows.") return df except Exception as exc: logger.error(f"Error executing '{description}': {exc}") raise # ------------------------------------------------------------------ def run_query_from_file(self, sql_path: str, description: str = ""): """Read a .sql file and execute it, returning a pandas DataFrame.""" sql_path = Path(sql_path) if not sql_path.exists(): raise FileNotFoundError(f"SQL file not found: {sql_path}") query = sql_path.read_text(encoding="utf-8") return self.run_query(query, description or sql_path.name) # ------------------------------------------------------------------ def store_df_to_snowflake( self, table_name: str, dataframe, database: str, schema: str, overwrite: bool = False, ): """ Write a pandas DataFrame to a Snowflake table. Args: table_name: Target table name (without database/schema prefix) dataframe: pandas DataFrame to write database: Target Snowflake database schema: Target Snowflake schema overwrite: If True, truncate the table before inserting; if False (default), append rows """ if dataframe is None or len(dataframe) == 0: logger.warning(f"store_df_to_snowflake: empty DataFrame, skipping write to {table_name}") return mode = "overwrite" if overwrite else "append" try: self.session.write_pandas( df=dataframe, table_name=table_name, database=database, schema=schema, overwrite=overwrite, auto_create_table=False, # Table must be created via SQL first quote_identifiers=False, use_logical_type = True ) logger.info( f"Stored {len(dataframe):,} rows to {database}.{schema}.{table_name} " f"(mode={mode})" ) except Exception as exc: logger.error(f"Error storing to {database}.{schema}.{table_name}: {exc}") raise # ------------------------------------------------------------------ def close(self): self.session.close() logger.info("Snowflake session closed.")