[{'Text': 'import snowflake.connector\nfrom dotenv import load_dotenv\nimport os\nimport re\nimport time\n\n\ndef split_sql_statements(sql: str) -> list[str]:\n    """\n    Split SQL on semicolons that are not inside line comments (--), block comments (/* */),\n    single-quoted strings (\'\' escape), or double-quoted delimited identifiers ("" escape).\n    """\n    parts: list[str] = []\n    current: list[str] = []\n    i = 0\n    n = len(sql)\n    # normal | line_comment | block_comment | squote | dquote\n    state = "normal"\n\n    while i < n:\n        c = sql[i]\n        if state == "normal":\n            if c == "-" and i + 1 < n and sql[i + 1] == "-":\n                state = "line_comment"\n                current.append(c)\n                current.append(sql[i + 1])\n                i += 2\n                continue\n            if c == "/" and i + 1 < n and sql[i + 1] == "*":\n                state = "block_comment"\n                current.append(c)\n                current.append(sql[i + 1])\n                i += 2\n                continue\n            if c == "\'":\n                state = "squote"\n                current.append(c)\n                i += 1\n                continue\n            if c == \'"\':\n                state = "dquote"\n                current.append(c)\n                i += 1\n                continue\n            if c == ";":\n                parts.append("".join(current))\n                current = []\n                i += 1\n                continue\n            current.append(c)\n            i += 1\n            continue\n        if state == "line_comment":\n            current.append(c)\n            if c in "\\n\\r":\n                state = "normal"\n            i += 1\n            continue\n        if state == "block_comment":\n            current.append(c)\n            if c == "*" and i + 1 < n and sql[i + 1] == "/":\n                current.append(sql[i + 1])\n                i += 2\n                state = "normal"\n                continue\n            i += 1\n            continue\n        if state == "squote":\n            current.append(c)\n            if c == "\'":\n                if i + 1 < n and sql[i + 1] == "\'":\n                    current.append(sql[i + 1])\n                    i += 2\n                    continue\n                state = "normal"\n            i += 1\n            continue\n        if state == "dquote":\n            current.append(c)\n            if c == \'"\':\n                if i + 1 < n and sql[i + 1] == \'"\':\n                    current.append(sql[i + 1])\n                    i += 2\n                    continue\n                state = "normal"\n            i += 1\n            continue\n\n    parts.append("".join(current))\n    return parts\n\n\ndef non_empty_sql_statements(sql: str) -> list[str]:\n    """Statements from split_sql_statements with whitespace-only segments removed."""\n    return [p.strip() for p in split_sql_statements(sql) if p.strip()]\n\n\nclass SnowflakeDB:\n    """\n    Class for connecting to the Snowflake database.\n    \n    Example usage:\n        `db = SnowflakeDB()`\n        `db.connect()`\n        `result = db.execute_safe_query(sql_query)`\n        `db.close_connection()`\n\n    Attributes:\n        user (str): Snowflake user.\n        password (str): Snowflake password.\n        account (str): Snowflake account.\n        warehouse (str): Snowflake warehouse.\n        database (str): Snowflake database.\n        schema (str): Snowflake schema.\n        connection (snowflake.connector.connection.SnowflakeConnection): Snowflake connection.\n    """\n    def __init__(self):\n        load_dotenv()\n\n        self.user = os.environ[\'SNOWFLAKE_USER\']\n        self.password = os.environ[\'SNOWFLAKE_PASSWORD\']\n        \n        self.account = \'fz24086.eu-west-1\'\n        self.warehouse = \'COMPUTE_WH\'\n        self.database = \'UPC_CORE\'\n        self.schema = \'REPORTS\'\n\n        self.connection = None\n\n    def connect(self):\n        """Establishes a connection to the Snowflake database."""\n        self.connection = snowflake.connector.connect(\n            user=self.user,\n            password=self.password,\n            account=self.account,\n            warehouse=self.warehouse,\n            database=self.database,\n            schema=self.schema\n        )\n\n    def close_connection(self):\n        """Closes the Snowflake database connection."""\n        if self.connection is not None:\n            self.connection.close()\n\n    def _fetchall_to_csv(self, cursor):\n        """Convert fetchall results to CSV (no pandas required)."""\n        import io\n\n        columns = [desc[0] for desc in cursor.description]\n        rows = cursor.fetchall()\n        output = io.StringIO()\n        output.write(",".join(columns) + "\\n")\n        for row in rows:\n            output.write(\n                ",".join(str(cell) if cell is not None else "" for cell in row) + "\\n"\n            )\n        return output.getvalue()\n\n    def query_to_csv(self, query):\n        """Executes a SQL query and returns results as CSV. Uses fetchall (no pandas)."""\n        cursor = self.connection.cursor()\n        try:\n            cursor.execute(query)\n            return self._fetchall_to_csv(cursor)\n        finally:\n            cursor.close()\n\n    def execute_ddl(self, query):\n        """Executes a DDL statement such as CREATE or ALTER and returns True if successful."""\n        cursor = self.connection.cursor()\n        try:\n            cursor.execute(query)\n            return True\n        finally:\n            cursor.close()\n\n    def get_distinct_value_list(self, table, column):\n        """Gets the distinct list of values for a given column in a given table."""\n        query = f"select distinct {column} from {table}"\n        return self.query_to_csv(query)\n\n    def validate_safe_query(self, query: str) -> tuple[bool, str]:\n        """\n        Validates that a SQL query is safe to execute (read-only operations only).\n        \n        Returns:\n            tuple[bool, str]: (is_valid, error_message)\n        """\n        # Remove comments and normalize whitespace\n        clean_query = re.sub(r\'--.*$\', \'\', query, flags=re.MULTILINE)\n        clean_query = re.sub(r\'/\\*.*?\\*/\', \'\', clean_query, flags=re.DOTALL)\n        clean_query = clean_query.strip().upper()\n        \n        # Dangerous keywords that should never be allowed\n        dangerous_keywords = [\n            \'DROP\', \'DELETE\', \'INSERT\', \'UPDATE\', \'CREATE\', \'ALTER\', \n            \'TRUNCATE\', \'REPLACE\', \'MERGE\', \'COPY\', \'PUT\', \'GET\',\n            \'REMOVE\', \'GRANT\', \'REVOKE\', \'USE\', \'SET\', \'UNSET\',\n            \'CALL\', \'EXECUTE\', \'EXEC\'\n        ]\n        \n        # Check for dangerous keywords at word boundaries\n        for keyword in dangerous_keywords:\n            if re.search(rf\'\\b{keyword}\\b\', clean_query):\n                return False, f"Query contains prohibited keyword: {keyword}"\n        \n        # Must start with SELECT, WITH, SHOW, DESCRIBE, or EXPLAIN\n        allowed_start_keywords = [\'SELECT\', \'WITH\', \'SHOW\', \'DESCRIBE\', \'DESC\', \'EXPLAIN\']\n        \n        starts_with_allowed = any(clean_query.startswith(keyword) for keyword in allowed_start_keywords)\n        if not starts_with_allowed:\n            return False, f"Query must start with one of: {\', \'.join(allowed_start_keywords)}"\n\n        # Reject multiple statements; ignore semicolons inside comments/strings\n        if len(non_empty_sql_statements(query)) > 1:\n            return False, "Multiple statements not allowed"\n\n        return True, ""\n\n    def submit_async_query(self, query: str) -> tuple:\n        """\n        Validate and submit a query asynchronously.\n\n        Returns (cursor, query_id) immediately — before execution completes.\n        Raises ValueError on validation failure, Exception on submit failure.\n        """\n        is_valid, error_message = self.validate_safe_query(query)\n        if not is_valid:\n            raise ValueError(f"Query validation failed - {error_message}")\n        cursor = self.connection.cursor()\n        cursor.execute_async(query)\n        return cursor, cursor.sfqid\n\n    def poll_and_fetch(self, cursor, query_id: str) -> str:\n        """\n        Poll until a submitted async query settles, then return results as CSV.\n\n        Returns string starting with `query_id:<id>` followed by CSV data,\n        or an error/cancellation message.\n        """\n        running = {\'RUNNING\', \'QUEUED\', \'RESUMING_WAREHOUSE\', \'ABORTING\'}\n        while self.connection.get_query_status(query_id).name in running:\n            time.sleep(0.5)\n\n        try:\n            self.connection.get_query_status_throw_if_error(query_id)\n        except Exception as e:\n            if getattr(e, \'errno\', None) == 604:\n                return f"query_id:{query_id}\\nCANCELLED: query was cancelled"\n            return f"query_id:{query_id}\\nDATABASE ERROR: {str(e)}"\n\n        cursor.get_results_from_sfqid(query_id)\n        rows = cursor.fetchall()\n        # description is populated after fetchall when using async execution\n        columns = [desc[0] for desc in cursor.description]\n        import io\n        output = io.StringIO()\n        output.write(",".join(columns) + "\\n")\n        for row in rows:\n            output.write(",".join(str(cell) if cell is not None else "" for cell in row) + "\\n")\n        cursor.close()\n        return f"query_id:{query_id}\\n{output.getvalue()}"\n\n    def execute_safe_query(self, query: str) -> str:\n        """\n        Executes a validated safe SQL query and returns results as CSV.\n\n        Uses async execution internally so the query can be cancelled externally\n        via cancel_query(query_id) while it is running.\n\n        Returns string starting with `query_id:<snowflake_query_id>` followed\n        by CSV data, or an error message.\n        """\n        try:\n            cursor, query_id = self.submit_async_query(query)\n        except ValueError as e:\n            return f"ERROR: {e}"\n        except Exception as e:\n            return f"DATABASE ERROR: {str(e)}"\n        return self.poll_and_fetch(cursor, query_id)\n    \n    def execute_safe_queries(self, queries: str) -> str:\n        """\n        Executes multiple validated safe SQL queries separated by semicolons.\n        Each query is validated and executed separately.\n        \n        Args:\n            queries: Multiple SQL queries separated by semicolons (must be read-only)\n            \n        Returns:\n            str: Combined results with clear query separators\n        """\n        query_list = non_empty_sql_statements(queries)\n        \n        if len(query_list) == 1:\n            # Single query, use the regular method\n            return self.execute_safe_query(query_list[0])\n        \n        results = []\n        for i, query in enumerate(query_list, 1):\n            results.append(f"=== QUERY {i} ===")\n            results.append(f"SQL: {query}")\n            results.append("RESULTS:")\n            \n            # Execute each query individually\n            result = self.execute_safe_query(query)\n            results.append(result)\n            results.append("")  # Empty line for separation\n        \n        return "\\n".join(results)\n\n    def check_query_status(self, query_id: str) -> str:\n        """\n        Returns the current status of a Snowflake query by ID.\n\n        Returns one of: RUNNING, SUCCESS, CANCELLED, ERROR:<message>\n        """\n        raw = self.connection.get_query_status(query_id).name\n        if raw in (\'RUNNING\', \'QUEUED\', \'RESUMING_WAREHOUSE\', \'ABORTING\'):\n            return \'RUNNING\'\n        if raw == \'SUCCESS\':\n            return \'SUCCESS\'\n        try:\n            self.connection.get_query_status_throw_if_error(query_id)\n        except Exception as e:\n            if getattr(e, \'errno\', None) == 604:\n                return \'CANCELLED\'\n            return f\'ERROR: {e}\'\n        return raw\n\n    def cancel_query(self, query_id: str) -> str:\n        """Cancels a running Snowflake query. Works from any connection."""\n        self.connection.execute_string(f"select system$cancel_query(\'{query_id}\')")\n        return f"cancelled:{query_id}"\n\n    def describe_table(self, table_name: str) -> str:\n        """\n        Describe a table in the Snowflake database.\n\n        Args:\n            table_name: The name of the table to describe. \n            Use the full database and schema name.\n            e.g. "upc_core.reports.upc_output"\n\n        Returns:\n            The schema as CSV string\n        """\n        query = f"describe table {table_name}"\n        cursor = self.connection.cursor()\n        try:\n            cursor.execute(query)\n            columns = [desc[0] for desc in cursor.description]\n            rows = cursor.fetchall()\n            import io\n            output = io.StringIO()\n            output.write(\',\'.join(columns) + \'\\n\')\n            for row in rows:\n                output.write(\',\'.join(str(cell) if cell is not None else \'\' for cell in row) + \'\\n\')\n            return output.getvalue()\n        finally:\n            cursor.close()'}]