""" Async Snowflake wrapper. """ import asyncio import logging from concurrent.futures import ThreadPoolExecutor from typing import cast import snowflake.connector as sf_connector from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from pandas import DataFrame from . import typings, utils from .enums import SnowflakeClientArg logger = logging.getLogger(__name__) # This is the Snowflake thread pool executor. It prevents uncontrolled number # of threads from being spawned when executing queries in parallel, which will # block the UI in our single-CPU environment. This is a workaround to the # fact that the Snowflake connector does not support async operations and is # blocking. To keep resources available for other stuff, we limit the number of threads # to a relatively low number. sf_thread_pool = ThreadPoolExecutor(max_workers=8) class Client: """Snowflake convenience methods. This class is a wrapper around the Snowflake connector that 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. """ def __init__( self, user: typings.User, account: typings.Account, warehouse: typings.Warehouse, private_key: typings.PrivateKey, database: typings.Database | None = None, schema: typings.Schema | None = None, private_key_password: typings.PrivateKeyPassword | None = None, ): """ Init Args: user: Snowflake user. account: Snowflake account. warehouse: Snowflake warehouse. database: Snowflake database. If not provided, table names in queries should be fully qualified (e.g. "DATABASE.SCHEMA.TABLE"). schema: Snowflake schema. If not provided, table names in queries should be fully qualified (e.g. "DATABASE.SCHEMA.TABLE"). 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 the private key, if encrypted. """ self.user: typings.User = user self.account: typings.Account = account self.warehouse: typings.Warehouse = warehouse self.database: typings.Database | None = database self.schema: typings.Schema | None = schema self._private_key_content: bytes | None = None self._private_key: typings.PrivateKey = private_key self.private_key_password: typings.PrivateKeyPassword | None = ( private_key_password ) self.connection = None self.echo_interval: int = 5 # To prevent race conditions: only allow one attempt to connect at # a time. This way, the connection will be established only once. self._connection_lock = asyncio.Lock() async def cursor(self): """Get a cursor to execute queries.""" if self.connection is None: await self.connect() return self.connection.cursor(sf_connector.DictCursor) async def _afetch( self, query: str, params: typings.BoundParams | None = None, as_df: bool = False, one: bool = False, reauthenticate_if_needed: bool = True, ) -> DataFrame | list | tuple: """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. as_df: Whether to return the result as a DataFrame. 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 a case, a single reauthentication attempt will be made, and if the error persists, it will be raised. """ try: cursor = await self.cursor() def fetch_async_nb(): query_id = cursor.execute_async(query, params=params)["queryId"] cursor.get_results_from_sfqid(query_id) if as_df: return cursor.fetch_pandas_all() return cursor.fetchone() if one else cursor.fetchall() loop = asyncio.get_event_loop() return await loop.run_in_executor(sf_thread_pool, fetch_async_nb) except sf_connector.errors.ProgrammingError as e: if reauthenticate_if_needed and "expired" in str(e).lower(): logger.debug( "Silently reauthenticating and retrying query due to expired auth " "token..." ) await self.disconnect() return await self._afetch( query, params=params, as_df=as_df, one=one, reauthenticate_if_needed=False, # Avoid infinite recursion ) raise async def afetch_one( self, query: str, params: typings.BoundParams | None = None ) -> tuple[typings.PossibleValues]: """Asynchronously fetch a query from Snowflake and return one result. Args: query: SQL query to execute. params: Parameters bound to the query. """ result = await self._afetch(query, params=params, one=True) return cast(tuple[typings.PossibleValues], result) async def afetch_all( self, query: str, params: typings.BoundParams | None = None, as_df: bool = False ) -> DataFrame | list[dict[str, typings.PossibleValues]]: """Asynchronously fetch a query from Snowflake and return all results. Args: query: SQL query to execute. params: Parameters bound to the query. as_df: Whether to return the results as a DataFrame. """ result = await self._afetch(query, params=params, as_df=as_df, one=False) if as_df: return cast(DataFrame, result) return cast(list[dict[str, typings.PossibleValues]], result) async def connect(self, timeout: int = 60) -> None: """Connect to Snowflake and persist the connection. https://docs.snowflake.com/en/user-guide/python-connector-example.html """ async with self._connection_lock: if self.connection is None: logger.debug("Establishing authenticated connection to Snowflake...") logging_task = asyncio.create_task( _log_periodic_connecting_message(5, timeout) ) try: kwargs = [ SnowflakeClientArg.USER, SnowflakeClientArg.ACCOUNT, SnowflakeClientArg.WAREHOUSE, SnowflakeClientArg.PRIVATE_KEY, ] if sum(1 for x in (self.database, self.schema) if x) == 1: raise ValueError( "Both database and schema must be provided, or " "table names in queries should be fully qualified." ) if self.database is not None: kwargs.append(SnowflakeClientArg.DATABASE) if self.schema is not None: kwargs.append(SnowflakeClientArg.SCHEMA) # 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. connect_kwargs = {k: getattr(self, k, None) for k in kwargs} | { "insecure_mode": False, # Disable OCSP checks } self.connection = await asyncio.wait_for( asyncio.to_thread( lambda: sf_connector.connect(**connect_kwargs) # type: ignore[arg-type,func-returns-value] ), timeout=timeout, ) logger.debug("Connection to Snowflake established.") finally: # Once the connection is established, cancel the message task # and handle the cancellation exception. logging_task.cancel() try: await logging_task except asyncio.CancelledError: pass else: logger.debug("Reusing persisted connection to Snowflake.") async def disconnect(self) -> None: """Disconnect from Snowflake and purge persisted connection, if any.""" async with self._connection_lock: if self.connection is not None: self.connection.close() self.connection = None logger.debug("Disconnected from Snowflake.") else: logger.debug("No connection to Snowflake to disconnect from.") @property def private_key(self) -> bytes: """Load private key from file following the procedure described in the Snowflake documentation. """ if self._private_key_content: return self._private_key_content private_key_content: bytes = utils.get_private_key_content( self._private_key.strip() ) p_key = serialization.load_pem_private_key( private_key_content, password=( self.private_key_password.encode() if self.private_key_password else None ), backend=default_backend(), ) private_key = p_key.private_bytes( encoding=serialization.Encoding.DER, format=serialization.PrivateFormat.PKCS8, encryption_algorithm=serialization.NoEncryption(), ) logger.debug("Snowflake private key loaded.") self._private_key_content = private_key return private_key async def _log_periodic_connecting_message( interval: int = 5, timeout: int = 30 ) -> None: """Log a message every `interval` seconds while establishing a connection to Snowflake. """ seconds_elapsed: int = 0 while seconds_elapsed < timeout: await asyncio.sleep(interval) seconds_elapsed += interval logger.debug( "Still establishing connection to Snowflake... (%d seconds elapsed)", seconds_elapsed, )