"""Async Snowflake Connector wrapper.""" import asyncio from .threadpool import sf_executor from threading import Lock from typing import Any from enums import StrEnum import snowflake.connector as sf_connector from pandas import DataFrame from common.src import logger from . import auth, data_converters from ...enums import OutputFormat, SFClientSettings from ...typings import BoundParams, SqlQuery logger = logger.new_logger(__name__) class Client: """Snowflake convenience methods. This class is a wrapper around the Snowflake connector which provides convenience methods for fetching data from Snowflake. It also persists the connection to Snowflake so that it can be reused during the lifetime of the object. """ class _Actions(StrEnum): CONNECT = "connect" DISCONNECT = "disconnect" _data_converters = { None: data_converters.data_to_json, # Default to JSON OutputFormat.MSGPACK: data_converters.data_to_msgpack, OutputFormat.CSV: data_converters.data_to_csv_string, OutputFormat.JSON: data_converters.data_to_json, } def __init__( self, user: str, account: str, warehouse: str, private_key: str, private_key_password: str = None, database: str | None = None, schema: str | None = None, ): """ Args: user: Snowflake user. account: Snowflake account. warehouse: Snowflake warehouse. private_key: Accepts a string with any of the following: - A path to a private key file. - A base64-encoded private key. - The private key itself. private_key_password: Password for private key, if encrypted. database: Snowflake database. schema: Snowflake schema. """ self.user = user self.account = account self.warehouse = warehouse self.database = database self.schema = schema self._private_key_content = None self._private_key = private_key self.private_key_password = private_key_password self.connection = None self.echo_interval = 5 # To prevent race conditions: only allow one attempt to connect at # a time. This way, the connection will be established only once. self._locks = { self._Actions.CONNECT: asyncio.Lock(), self._Actions.DISCONNECT: Lock(), } @property def private_key(self) -> bytes: """Get private key content from the provided private key. Raises: ValueError: If private key has not been provided. """ if self._private_key_content: return self._private_key_content try: private_key_stripped = self._private_key.strip() except AttributeError: if self._private_key is None: raise ValueError("Private key has not been provided.") from None raise ValueError( "Private key must be a string, but got: " f"{type(self._private_key).__name__}" ) from None else: private_key_der: bytes = auth.get_private_key_der( private_key_stripped, password=self.private_key_password ) logger.debug("Snowflake private key loaded.") self._private_key_content = private_key_der return private_key_der async def connect(self) -> None: """Connect to Snowflake and persist the connection. https://docs.snowflake.com/en/user-guide/python-connector-example.html """ async with self._locks[self._Actions.CONNECT]: if self.connection is None: logger.debug("Establishing authenticated connection to Snowflake...") message_task = asyncio.create_task(self._log_periodic_message()) kwargs = ( SFClientSettings.USER, SFClientSettings.ACCOUNT, SFClientSettings.WAREHOUSE, SFClientSettings.DATABASE, SFClientSettings.SCHEMA, SFClientSettings.PRIVATE_KEY, ) self.connection = await asyncio.to_thread( sf_connector.connect, # Connection may occasionally have extreme latencies due to OCSP # issues. To avoid this, it is possible to set the insecure_mode # parameter to True. This will disable OCSP checks, so it is # discouraged to use this in production environments. **{**{k: getattr(self, k) for k in kwargs}, "insecure_mode": False}, ) logger.debug("Connection to Snowflake established.") # Once the connection is established, cancel the message task # and handle the cancellation exception. message_task.cancel() try: await message_task except asyncio.CancelledError: pass else: logger.debug("Reusing persisted connection to Snowflake.") def disconnect(self) -> None: """Disconnect from Snowflake and purge persisted connection, if any.""" with self._locks[self._Actions.DISCONNECT]: try: self.connection.close() except AttributeError: logger.debug("No connection to Snowflake to disconnect from.") else: self.connection = None logger.debug("Disconnected from Snowflake.") async def get_cursor(self): """Get a cursor to execute queries.""" if self.connection is None: await self.connect() return self.connection.cursor() async def afetch_one(self, query: SqlQuery, *, params: BoundParams = None) -> tuple: """Asynchronously fetch a query from Snowflake and return one result. Args: query: SQL query to execute. params: Parameters bound to the query. """ return await self._afetch(query, params=params, one=True) async def afetch_all( self, query: SqlQuery, *, params: BoundParams = None, output_format: OutputFormat = OutputFormat.JSON, ) -> DataFrame | list: """Asynchronously fetch a query from Snowflake and return all results. Args: query: SQL query to execute. params: Parameters bound to the query. output_format: The output format of the query results. Defaults to JSON. Returns: See `_afetch` for the return value. """ return await self._afetch( query, params=params, output_format=output_format, one=False ) async def _afetch( self, query: SqlQuery, params: BoundParams = None, output_format: OutputFormat = OutputFormat.JSON, one: bool = False, reauthenticate_if_needed: bool = True, ) -> DataFrame | list[dict] | str | tuple[Any]: """Asynchronously fetch a query from Snowflake and return one or all results, depending on the `one` parameter. Args: query: SQL query to execute. params: Parameters bound to the query. output_format: The output format of the query results. Defaults to JSON. one: Whether to return only one result, or all results. reauthenticate_if_needed: Whether to silently reauthenticate and retry the query if the auth token has expired. It should generally be set to True. In such case, a single authentication attempt will be made, and if the error persists, it will be raised. Returns: The result of the query, depending on the output format and the `one` parameter: * If `one` is True, a single result is returned, as a tuple. * If `one` is False and `output_format` is OutputFormat.DF, a DataFrame is returned. * If `one` is False and `output_format` is OutputFormat.CSV, a CSV string is returned. * If `one` is False and `output_format` is OutputFormat.JSON, a list of dictionaries is returned. """ try: cursor = await self.get_cursor() loop = asyncio.get_event_loop() def _fetch_async_nb(): query_id = cursor.execute_async(query, params=params)["queryId"] cursor.get_results_from_sfqid(query_id) if output_format == OutputFormat.DF: return cursor.fetch_pandas_all() data = cursor.fetchone() if one else cursor.fetchall() if one: return data col_names = tuple(col.name for col in cursor.description) return self._data_converters[output_format](col_names, data) return await loop.run_in_executor(sf_executor, _fetch_async_nb) except sf_connector.errors.ProgrammingError as e: if reauthenticate_if_needed and "expired" in str(e).lower(): logger.debug( "Silently authenticating and retrying query due to expired auth " "token..." ) self.disconnect() return await self._afetch( query, params=params, output_format=output_format, one=one, reauthenticate_if_needed=False, # Avoid infinite recursion ) raise async def _log_periodic_message(self) -> None: """Log a periodic message while establishing a connection to Snowflake.""" seconds_elapsed = 0 while True: await asyncio.sleep(self.echo_interval) logger.debug( "Still establishing connection to Snowflake... " "({} seconds elapsed)", seconds_elapsed, )