process_helpscout/snowflake_conn.py

"""
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.")