""" Async Snowflake wrapper. """ import asyncio import base64 import os from concurrent.futures import ThreadPoolExecutor from threading import Lock 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 logger from .typings import BoundParams logger = logger.new_logger(__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 prioritize the user interface, 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 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. """ def __init__( self, user: str, account: str, warehouse: str, database: str, schema: str, private_key: str, private_key_password: str = None, ): """ Init Args: user: Snowflake user. account: Snowflake account. warehouse: Snowflake warehouse. database: Snowflake database. schema: Snowflake schema. 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. """ 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._connection_lock = asyncio.Lock() self._disconnect_lock = Lock() async def cursor(self): """Get a cursor to execute queries.""" if self.connection is None: await self.connect() return self.connection.cursor() async def _afetch( self, query: str, params: BoundParams = 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 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..." ) 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: 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: str, params: BoundParams = None, as_df: bool = False ) -> DataFrame | list: """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. """ return await self._afetch(query, params=params, as_df=as_df, one=False) async def connect(self) -> None: """Connect to Snowflake and persist the connection. https://docs.snowflake.com/en/user-guide/python-connector-example.html """ async def _log_periodic_message(): seconds_elapsed = 0 while True: await asyncio.sleep(self.echo_interval) logger.debug( "Still establishing connection to Snowflake... " "({} seconds elapsed)", seconds_elapsed, ) async with self._connection_lock: if self.connection is None: logger.debug("Establishing authenticated connection to Snowflake...") message_task = asyncio.create_task(_log_periodic_message()) kwargs = ( "user", "account", "warehouse", "database", "schema", "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._disconnect_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 = self._get_private_key_content( self._private_key.strip() ) p_key = serialization.load_pem_private_key( private_key_content, password=self.private_key_password or 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 @staticmethod def _get_private_key_content(provided_private_key: str) -> bytes: """Get private key content from the provided private key, which is a string representing any of the following: - A path to a private key file. - A base64-encoded private key. - The private key itself. Args: provided_private_key: Private key as a string. Any of the possible formats described above. """ # Handle private key as a string which contains the key itself. if provided_private_key.strip().startswith("-----BEGIN"): private_key_content = provided_private_key.encode("utf-8") # Handle private key as a file path. To distinguish between a file path # and base64-encoded private key, we check if the file exists. elif os.path.exists(provided_private_key): with open(provided_private_key, "rb") as key: private_key_content = key.read() # Handle private key as a base64-encoded string. else: private_key_content = base64.b64decode(provided_private_key) return private_key_content