"""Snowflake connection adapter. Implements the DBConnection protocol for reference data loading using Snowflake's connector and private key authentication. """ from __future__ import annotations import logging from collections.abc import Callable, Generator from contextlib import contextmanager from enum import IntEnum from functools import wraps from typing import ParamSpec, TypeVar import snowflake.connector from pydantic import BaseModel, ConfigDict, Field from snowflake.connector import ( SnowflakeConnection as RawSnowflakeConnection, ) from snowflake.connector.cursor import SnowflakeCursor from snowflake.connector.errors import ( OperationalError as SnowflakeOperationalError, ) from snowflake.connector.s3_storage_client import SnowflakeS3RestClient from src.connectors.db import Connection from src.connectors.snowflake.utils import get_private_key from src.errors import TransientError # Disable Snowflake connector logging logging.getLogger('snowflake.connector.connection').propagate = False class SnowflakeConfig(BaseModel): """Configuration model for a Snowflake database connection.""" model_config = ConfigDict(frozen=True, populate_by_name=True) # Connection host: str | None = Field(default=None, description='Your Snowflake host.') account: str = Field( ..., description='Your Snowflake account ID (e.g., xy12345.region.aws).' ) user: str = Field(..., description='Username for authentication.') private_key: str | None = Field( default=None, description='The private key content (PEM).', ) private_key_path: str | None = Field( default=None, description='The absolute path to the .p8/.pem file.', ) private_key_passphrase: str | None = Field( default=None, description='The passphrase for the encrypted private key.' ) warehouse: str = Field(..., description='The name of the compute warehouse to use.') database: str = Field(..., description='The name of the database to connect to.') schema_name: str = Field( ..., description='The schema within the database to use.', alias='schema' ) role: str | None = Field( default=None, description='The access role to use for the connection.' ) connection_timeout: int = Field( default=10, description='Connection timeout in seconds.' ) # Settings autocommit: bool = Field( default=False, description='Whether transactions should commit automatically.' ) class SnowflakeConnection(Connection): """Snowflake database connection adapter.""" def __init__(self, conn: RawSnowflakeConnection): """Initialize the adapter with a raw Snowflake connection. Args: conn: An active snowflake connection. """ self._conn = conn def close(self) -> None: """Close the connection.""" self._conn.close() def commit(self) -> None: """Commit the current transaction to the database.""" self._conn.commit() def cursor(self) -> SnowflakeCursor: # type: ignore[override] """Return a raw Snowflake cursor object.""" return self._conn.cursor() # cursor.fetch_pandas_all() def rollback(self) -> None: """Roll back the current transaction, discarding changes.""" self._conn.rollback() class SnowflakeConnectionFactory: """Factory for creating Snowflake database connections with configuration management.""" def __init__(self, config: SnowflakeConfig) -> None: """Initialize the connection factory with Snowflake configuration. Args: config: SnowflakeConfig instance containing database connection parameters. """ self._config = config @contextmanager def connection(self) -> Generator[SnowflakeConnection, None, None]: """Create and yield a Snowflake database connection as a context manager. The connection is automatically closed when the context exits. Yields: SnowflakeConnection: Snowflake database adapter. Raises: TransientError: If connection fails due to transient database errors. SnowflakeOperationalError: If connection fails due to non-transient errors. """ with snowflake_connection(self._config) as conn: yield SnowflakeConnection(conn) class SnowflakeErrorCode(IntEnum): """Known Snowflake error codes, mostly related to transient errors.""" # Query issued in another transaction is running for too long QUERY_TIMEOUT = 57002 # Warehouse is suspended or cannot be found (often seen during warm-up) WAREHOUSE_SUSPENDED = 604 # Connection time-out or network issue CONNECTION_ERROR = 252001 # Statement cancelled by the user/system (can happen during resource contention) STATEMENT_CANCELLED = 607 # Transient network/service error INTERNAL_SERVICE_ERROR = 300001 P = ParamSpec('P') T = TypeVar('T') TRANSIENT_ERROR_CODES: frozenset[int] = frozenset( { SnowflakeErrorCode.QUERY_TIMEOUT, SnowflakeErrorCode.WAREHOUSE_SUSPENDED, SnowflakeErrorCode.CONNECTION_ERROR, SnowflakeErrorCode.STATEMENT_CANCELLED, SnowflakeErrorCode.INTERNAL_SERVICE_ERROR, } ) def _try_transient_error(e: SnowflakeOperationalError) -> None: """Check if an OperationalError is caused by a transient condition. If transient, a custom TransientError is raised. Args: e: The caught snowflake.connector.errors.OperationalError exception. Raises: TransientError: If the error code matches one in TRANSIENT_ERROR_CODES. """ if hasattr(e, 'errno') and e.errno in TRANSIENT_ERROR_CODES: raise TransientError(f'Snowflake operation failed: {e}') from e # Also check for connection-related errors in the message if 'failed to establish a connection' in str(e).lower(): raise TransientError(f'Snowflake connection failed: {e}') from e def handle_snowflake_errors(func: Callable[P, T]) -> Callable[P, T]: """Decorate database operations and handle transient errors. If a transient error occurs, it is converted into a TransientError to allow for retry logic. Non-transient errors are re-raised as is. """ @wraps(func) def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: try: return func(*args, **kwargs) except SnowflakeOperationalError as e: _try_transient_error(e) raise # Re-raise if it wasn't transient return wrapper @contextmanager def snowflake_connection( config: SnowflakeConfig, ) -> Generator[RawSnowflakeConnection, None, None]: """Context manager for connection objects. Args: config: Snowflake configuration object. Yields: RawSnowflakeConnection: An active Snowflake connection. Raises: TransientError: If the connection attempt fails due to a transient error. SnowflakeOperationalError: If the connection attempt fails due to a non-transient error. """ private_key_bytes = get_private_key( config.private_key, config.private_key_path, config.private_key_passphrase ) connection_timeout = config.connection_timeout if connection_timeout is None: connection_timeout = 10 try: conn = snowflake.connector.connect( host=config.host, account=config.account, user=config.user, private_key=private_key_bytes, warehouse=config.warehouse, database=config.database, schema=config.schema_name, role=config.role, connection_timeout=connection_timeout, autocommit=config.autocommit, ) except SnowflakeOperationalError as e: _try_transient_error(e) raise try: yield conn finally: if conn: conn.close() # MONKEY PATCHES def _get_bucket_accelerate_config_patch( self: SnowflakeS3RestClient, bucket_name: str ) -> bool: """Disable S3 transfer acceleration. Args: self: A SnowflakeS3RestClient instance. bucket_name: S3 bucket name. """ return False SnowflakeS3RestClient._get_bucket_accelerate_config = ( # type: ignore[method-assign] _get_bucket_accelerate_config_patch ) def _transfer_accelerate_config_patch( self: SnowflakeS3RestClient, use_accelerate_endpoint: bool | None = None ) -> bool: """Disable S3 transfer acceleration. Args: self: A SnowflakeS3RestClient instance. use_accelerate_endpoint: Boolean. """ self.endpoint = ( f'https://{self.s3location.bucket_name}.s3.{self.region_name}.amazonaws.com' ) return False SnowflakeS3RestClient.transfer_accelerate_config = _transfer_accelerate_config_patch # type: ignore[method-assign]