"""DB Connector. Manages interactions with Snowflake.""" from contextlib import contextmanager from functools import wraps import os import re from sqlalchemy import create_engine from sqlalchemy import event from sqlalchemy import exc from sqlalchemy import pool from sqlalchemy import select from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker from backend import config from backend.constants import db as db_consts BaseModel = declarative_base() role = os.environ.get('SNOWFLAKE_ROLE') database = os.environ.get('SNOWFLAKE_DATABASE') schema = os.environ.get('SNOWFLAKE_SCHEMA') warehouse = os.environ.get('SNOWFLAKE_WAREHOUSE') def _create_engine(db_url): """Create engine based on configuration settings.""" if config.POOL_CLASS == pool.QueuePool: db_engine = create_engine( db_url, pool_size=config.POOL_SIZE, max_overflow=config.POOL_MAX_OVERFLOW, pool_recycle=config.POOL_RECYCLE) @event.listens_for(db_engine, 'connect') def set_snowflake_params(connection, connection_record): """Set session params. When using raw SQL with fully qualified table names (db.schema.table), USE DATABASE and USE SCHEMA are not required. """ with connection.cursor() as cur: if role: cur.execute('USE ROLE {};'.format(role)) if database: cur.execute('USE DATABASE {};'.format(database)) if schema: cur.execute('USE SCHEMA {};'.format(schema)) if warehouse: cur.execute('USE WAREHOUSE {};'.format(warehouse)) # SQLAlchemy 1.2b supports pessimistic disconnect handling out of box. # Once version 1.2 is out of beta, it is recommended to use # the `pool_pre_ping` param for `create_engine` and remove any code # in this file related to `_ping_connection` # See http://docs.sqlalchemy.org/en/latest/core/pooling.html if config.POOL_PRE_PING: @event.listens_for(db_engine, 'engine_connect') def ping_connection(connection, branch): _ping_connection(connection, branch) return db_engine else: return create_engine(db_url, poolclass=config.POOL_CLASS) def _ping_connection(connection, branch): """Ping database connection after engine_connect event. This function is copied verbatim from http://docs.sqlalchemy.org/en/latest/core/pooling.html """ if branch: # "branch" refers to a sub-connection of a connection, # we don't want to bother pinging on these. return # turn off "close with result". This flag is only used with # "connectionless" execution, otherwise will be False in any case save_should_close_with_result = connection.should_close_with_result connection.should_close_with_result = False try: # run a SELECT 1. use a core select() so that # the SELECT of a scalar value without a table is # appropriately formatted for the backend connection.scalar(select([1])) except exc.DBAPIError as err: # catch SQLAlchemy's DBAPIError, which is a wrapper # for the DBAPI's exception. It includes a .connection_invalidated # attribute which specifies if this connection is a "disconnect" # condition, which is based on inspection of the original exception # by the dialect in use. if err.connection_invalidated: # run the same SELECT again - the connection will re-validate # itself and establish a new connection. The disconnect detection # here also causes the whole connection pool to be invalidated # so that all stale connections are discarded. connection.scalar(select([1])) else: raise finally: connection.should_close_with_result = save_should_close_with_result # Do not use these variables directly other than running unit tests snowflake_db_engine = _create_engine(config.SNOWFLAKE_DB_URL) # please don't use sessions directly; instead use db_session sessions = { db_consts.SNOWFLAKE: sessionmaker(bind=snowflake_db_engine), } @contextmanager def db_session(db_name=db_consts.SNOWFLAKE): """Provide a transactional scope around a series of operations. Taken from http://docs.sqlalchemy.org/en/latest/orm/session_basics.html. This handles rollback and closing of session, so there is no need to do that throughout the code. Args: db_name (session): Name of database to connect to Usage: with db_session() as session: session.execute(query) """ session = sessions[db_name]() try: yield session session.commit() except: session.rollback() raise finally: session.close() def db_session_wrap(func): """DB Session Wrappper. Creates a new session if one isn't passed in. """ @wraps(func) def wrapper(*args, **kwargs): session = kwargs.pop('session', None) if session: return func(*args, session=session, **kwargs) else: with db_session() as session: return func(*args, session=session, **kwargs) return wrapper class SQLLoader(object): """Loads SQL queries/templates from text files.""" def __init__(self, path_to_file, folder='/queries'): """Constructor for 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 BaseValidator(object): """Validate parameters of SQL queries..""" def __init__(self): """Constructor.""" pass def is_valid_identifier(self, identifier): """Check if passed identifier is valid. Args: identifier (str): An SQL identifier (e.g. schema or table name). Returns: bool: True if valid, False otherwise. """ # at least one non-digit char at the beginning if re.match('^[a-zA-Z_]+[a-zA-Z0-9_]*$', identifier): return True if re.match('^\"[^\"]*\"$', identifier): return True return False def format_identifiers(self, sql_template, params): """Format SQL template with all the identifiers. Args: sql_template (str): SQL template to format. params (dict): Params (identifiers and non-identifiers). Returns: tuple (sql_template, non_identifier_params): sql_template is a template which was formatted with identifiers, non_identifier_params is a dict with a rest of params, which are not identifiers and could be bound within execute(). """ id_pattern = re.compile('%\(([a-zA-Z_]+[a-zA-Z0-9_]*)\)i') ids = {k: params.pop(k) for k in list(set(id_pattern.findall( sql_template)))} for k, v in ids.items(): assert self.is_valid_identifier(v) sql_template = sql_template.replace('%({})i'.format(k), v) non_identifier_params = params return sql_template, non_identifier_params