"""DB Connector. Manages interactions with Snowflake.""" import base64 from contextlib import contextmanager from enum import Enum from functools import partial from itertools import chain import logging import os import re try: from flask import g except ImportError: g = {} # for unit test purposes try: from ddtrace import patch except ImportError: def _blank_fn(*args, **kwargs): pass patch = _blank_fn patch(sqlalchemy=True) # patch sqlalchemy with Datadog if ddtrace is available from snowflake.connector.errors import DatabaseError # noqa from snowflake.sqlalchemy import URL # noqa from snowflake.sqlalchemy.base import SnowflakeExecutionContext # noqa import sqlalchemy # noqa from sqlalchemy.engine import Engine # noqa from sqlalchemy.engine.interfaces import CreateEnginePlugin # noqa from sqlalchemy.event import listens_for # noqa from sqlalchemy.exc import DBAPIError # noqa from sqlalchemy.orm import sessionmaker # noqa from sqlalchemy import text # noqa from . import validator # noqa logger = logging.getLogger(__name__) validator = validator.BaseValidator() @sqlalchemy.event.listens_for(Engine, 'before_cursor_execute') def _add_query_tag_with_correlation_id( conn, cursor, statement, parameters, context, executemany): """Replace _statement_params with a dict with a query_tag.""" if not isinstance(context, SnowflakeExecutionContext): return correlation_id = None try: correlation_id = g.get('correlation_id') except Exception as e: logger.warning( 'Error while getting correlation_id from g: {}'.format(str(e))) if correlation_id: cursor.execute = partial( cursor.execute, _statement_params={'query_tag': correlation_id}) class Fetch(Enum): """Enum corresponds to methods: execute, fetchone, fetchall.""" NONE = 1 ONE = 2 ALL = 3 def _merge_configs(c1, c2): """Merge two flat configs. Values from c1 get overridden by values from c2 if the keys collide. Args: c1 (dict): First config. c2 (dict): Second config. Returns: dict: Result dict containing merged result. """ c1 = c1 or {} c2 = c2 or {} return {k: v for k, v in chain(c1.items(), c2.items()) if v} def _get_sf_config(sf_config=None): """Get Snowflake credentials from environment. Args: sf_config (dict): Full or partial (without credentials) Snowflake config. Returns: dict: Merged Snowflake config. """ env_sf_config = dict( role=os.environ.get('SNOWFLAKE_ROLE'), account=os.environ.get('SNOWFLAKE_ACCOUNT'), user=os.environ.get('SNOWFLAKE_USER'), password=os.environ.get('SNOWFLAKE_PASSWORD'), database=os.environ.get('SNOWFLAKE_DATABASE'), schema=os.environ.get('SNOWFLAKE_SCHEMA'), warehouse=os.environ.get('SNOWFLAKE_WAREHOUSE'), private_key=os.environ.get('SNOWFLAKE_KEY')) if sf_config: # rename db to database, if necessary if 'db' in sf_config: sf_config['database'] = sf_config['db'] del sf_config['db'] # this way we can pass partial sf_config (e.g., only database and schema) merged_sf_config = _merge_configs(env_sf_config, sf_config) assert merged_sf_config.get('role') assert merged_sf_config.get('account') assert merged_sf_config.get('user') assert (merged_sf_config.get('password') or merged_sf_config.get('private_key')) assert merged_sf_config.get('database') assert merged_sf_config.get('schema') assert merged_sf_config.get('warehouse') return merged_sf_config _sf_default_sessionmaker = None def set_default_sessionmaker( sf_config=None, pool_size=5, pool_recycle=60*40, pool_pre_ping=False, pool_reset_on_return=None, invalidate_connections=True, **kwargs): """Set default sessionmaker, if no custom Snowflake config specified. Args: sf_config (dict): Full or partial (without credentials) Snowflake config. pool_size (int): Number of connections in a pool. pool_recycle: (int): A lifespan of a connection in a pool. pool_pre_ping (bool): If True, ping every established connection. pool_reset_on_return (str or None): Set reset mode for connections before returning them to a pool. invalidate_connections (bool): whether invalidate connections on error. """ global _sf_default_sessionmaker conn_params = _get_sf_config(sf_config) if invalidate_connections: conn_params['plugin'] = 'sf_disconnect' _sf_default_sessionmaker = sessionmaker( bind=sqlalchemy.create_engine( URL(**conn_params), pool_size=pool_size, pool_recycle=pool_recycle, pool_pre_ping=pool_pre_ping, pool_reset_on_return=pool_reset_on_return, **kwargs)) def _get_sessionmaker( sf_config=None, pool_pre_ping=False, pool_reset_on_return=None): """Get sessionmaker. Args: sf_config (dict): Full or partial (without credentials) Snowflake config. pool_pre_ping (bool): If True, ping every established connection. pool_reset_on_return (str or None): Set reset mode for connections before returning them to a pool. If no custom Snowflake config specified, returns default sessionmaker. """ if sf_config is None and _sf_default_sessionmaker: return _sf_default_sessionmaker merged_sf_config = _get_sf_config(sf_config) engine_params = {} private_key = merged_sf_config.get('private_key', None) if private_key: decoded_key = base64.b64decode(bytes(private_key, encoding='utf-8')) engine_params['connect_args'] = { 'private_key': decoded_key } if sf_config: return sessionmaker( bind=sqlalchemy.create_engine(URL(**merged_sf_config), **engine_params)) if not _sf_default_sessionmaker: set_default_sessionmaker( pool_pre_ping=pool_pre_ping, pool_reset_on_return=pool_reset_on_return, **engine_params) return _sf_default_sessionmaker @contextmanager def get_session( sf_config=None, commit_before_close=True, pool_pre_ping=False, pool_reset_on_return=None): """Create connection to the Snowflake. Provide a transactional scope around a series of operations. Args: sf_config (dict): Full or partial (without credentials) Snowflake config. commit_before_close (bool): Commit before close session or not. pool_pre_ping (bool): If True, ping every established connection. pool_reset_on_return (str or None): Set reset mode for connections before returning them to a pool. Returns: Session: SQLAlchemy session. """ session = _get_sessionmaker( sf_config, pool_pre_ping=pool_pre_ping, pool_reset_on_return=pool_reset_on_return)() try: yield session if commit_before_close: session.commit() except Exception: session.rollback() raise finally: session.close() def _execute( sql, params=None, sf_config=None, fetch=Fetch.NONE, retry_on_disconnect=True): """Execute an SQL statement. Args: sql (str): parametrized sql statement. params (dict): A dict of params to bind to sql_template. sf_config (dict): Full or partial (without credentials) Snowflake config. retry_on_disconnect (bool): whether we should retry to execute the query on disconnect. """ if params: sql, params = validator.format_identifiers(sql, params) sql = text(sql) while True: session = _get_sessionmaker(sf_config)() try: res = session.execute(sql, params) except DBAPIError as dpe: if not (retry_on_disconnect and isinstance(dpe.orig, DatabaseError)): raise logger.warning( 'Trying to re-execute due to database error: %s', str(dpe)) retry_on_disconnect = False else: if fetch == Fetch.NONE: return None if fetch == Fetch.ONE: return res.cursor.fetchone() if fetch == Fetch.ALL: return list(res.cursor.fetchall()) finally: session.close() def execute( sql_template, params=None, sf_config=None, retry_on_disconnect=True, **kwargs): """Execute an SQL statement. Args: sql_template (str): A template to be parametrized. All the identifiers (db name, schema name, table name, column names) should be already present. params (dict): A dict of params to bind to sql_template. sf_config (dict): Full or partial (without credentials) Snowflake config. retry_on_disconnect (bool): whether we should retry to execute the query on disconnect. Returns: dict: A dict with results. """ return _execute( sql_template, params, sf_config, fetch=Fetch.NONE, retry_on_disconnect=retry_on_disconnect) def fetchone( sql_template, params=None, sf_config=None, retry_on_disconnect=True, **kwargs): """Execute SQL and fetch first tuple. Useful for SQL statements which are supposed to return just one row, like COUNT, etc.. Args: sql_template (str): A template to be parametrized. All the identifiers (db name, schema name, table name, column names) should be already present. params (dict): A dict of params to bind to sql_template. sf_config (dict): Full or partial (without credentials) Snowflake config. retry_on_disconnect (bool): whether we should retry to execute the query on disconnect. Returns: dict: A dict with results. """ return _execute( sql_template, params, sf_config, fetch=Fetch.ONE, retry_on_disconnect=retry_on_disconnect) def fetchall( sql_template, params=None, sf_config=None, retry_on_disconnect=True, **kwargs): """Execute query and fetch all results. Args: sql_template (str): A template to be parametrized. All the identifiers (db name, schema name, table name, column names) should be already present. params (dict): A dict of params to bind to sql_template. sf_config (dict): Full or partial (without credentials) Snowflake config. retry_on_disconnect (bool): whether we should retry to execute the query on disconnect. Returns: dict: A dict with results. """ return _execute( sql_template, params, sf_config, fetch=Fetch.ALL, retry_on_disconnect=retry_on_disconnect) class SQLLoader(object): """Loads SQL queries/templates from text files. Do not forget to include these files to setup.py of the repo: setup( package_data={ # include all the .sql files 'snowflake_etl.flows.rs2sf': ['queries/*.sql']}) """ def __init__(self, path_to_file, folder='/queries'): """Init SQL loader. Usage example: import os sql = SQLLoader(__file__) Args: path_to_file (str): absolute path to the file which imports SQLLoader. """ self.sql_files_root = os.path.realpath( os.path.dirname(path_to_file)) + folder self.query_cash = {} def _load_query(self, query_name): """Load query from the disk. Args: query_name (str): name of the query (filename without extension). """ path = '{root}/{query_name}.sql'.format( root=self.sql_files_root, query_name=query_name) with open(path, 'r') as q: query = q.read() # get rid of docstring in the beginning of the SQL file # this is not required, but useful when debugging query = re.sub(r'^/\*.*\*/\n+', '', query, flags=re.DOTALL) return query def load_query(self, item): """Access queries by name. This method uses cache to reduce I/O operations. Usage example: sql_queries = SQLLoader('query_path') sql_queries.get_query('query_name') Args: item (str): name of the query (filename without extension). Returns: str: SQL query template. """ if item not in self.query_cash: self.query_cash[item] = self._load_query(item) return self.query_cash[item] class SnowflakeDisconnectPlugin(CreateEnginePlugin): """Snowflake+sqlalchemy plugin.""" def handle_dialect_kwargs(self, dialect_cls, dialect_args): """Parse and modify dialect kwargs.""" def is_disconnect(sf, e, connection, cursor): """Signal sqlalchemy pool to invalidate the connection.""" return isinstance(e, DatabaseError) # Monkey-patch snowflake dialect to invalidate connections on error dialect_cls.is_disconnect = is_disconnect def update_url(self, url): """Necessary for SQLalchemy 2.0.""" pass @listens_for(Engine, 'checkout') def log_overflow(dbapi_connection, connection_record, connection_proxy): """Emit a log message when pool overflow is not empty. Args: dbapi_connection: Snowflake connection instance. connection_record: _ConnectionRecord instance. connection_proxy: _ConnectionFairy instance. """ try: of = connection_proxy._pool.overflow() except Exception as e: logger.exception( 'Could not get the number of overflow connections: %s', str(e)) else: if of > 0: logger.warning('Connection pool overflow on checkout: %s', of)