"""Visualise and analyze cache-related Snowflake performance issues.""" from concurrent.futures import ThreadPoolExecutor import datetime import json import logging import os import time import boto3 import datadog import matplotlib.pyplot as plt import pandas as pd from snowflake_connector.etl_connector import SQLLoader, SnowflakeSQLExecutor from snowflake_connector.metadata_connector import SnowflakeBadResponse, \ SnowflakeMetadataConnector import config import secret_manager logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) sqlloader = SQLLoader(__file__) def _get_s3_csv_filename(run_timestamp): """Get a full S3 key. Args: run_timestamp (datetime.datetime): A timestamp which contains a given day. Returns: str: A full S3 key, e.g. 's3://dev-cucumbers/cache_stats/2019-11-11_cache_stats.csv'. """ run_date = run_timestamp.date() csv_filename = '{}_cache_stats.csv'.format(run_date) return config.S3_PATH + csv_filename def populate_cache_stats_for_a_day(run_timestamp): """Populate CACHE_STATS_BY_QUERYID table for a given day. Args: run_timestamp (datetime.datetime): A timestamp which contains a given day. """ run_date = run_timestamp.date() start_timestamp = datetime.datetime.strptime( run_timestamp.strftime('%Y-%m-%d 00:00:00.000'), '%Y-%m-%d %H:%M:%S.000') timestamp_ranges = [ ( (start_timestamp + datetime.timedelta(hours=hour)).strftime( '%Y-%m-%d %H:%M:%S.000') + ' -0500', (start_timestamp + datetime.timedelta(hours=hour + 1)).strftime( '%Y-%m-%d %H:%M:%S.000') + ' -0500' ) for hour in range(0, 24) ] start_time = time.perf_counter() with SnowflakeSQLExecutor(config.SF_CONFIG) as executor: common_params = { 'date': run_date, 'db': config.SF_CONFIG['db'], 'schema': config.SF_CONFIG['schema'] } # clean up the table in order to avoid duplication executor.fetchone_query( sqlloader, 'delete_from_cache_stats_by_queryid', common_params) for start_timestamp, end_timestamp in timestamp_ranges: params = { **common_params, **{ 'start_timestamp': start_timestamp, 'end_timestamp': end_timestamp} } executor.fetchone_query( sqlloader, 'create_temp_stats', params) # populate CACHE_STATS_CACHE_STATS_BY_QUERYID for each warehouse for warehouse in config.SNOWFLAKE_WAREHOUSES_TO_MONITOR: insert_params = { **params, **{'users_to_monitor': config.SNOWFLAKE_USERS_TO_MONITOR, 'warehouse_to_monitor': warehouse, 'database_to_monitor': 'FACTS', 'schema_to_monitor': 'PROD' } } executor.fetchone_query( sqlloader, 'insert_into_cache_stats_by_queryid', insert_params) query_ids = [ query_id[0] for query_id in executor.fetchall_query( sqlloader, 'get_queryids_from_cache_stats_by_queryid', params) ] results_by_query_id = get_cache_stats(query_ids) query_ids_tuples = ','.join( "('" + query_id + "',{remote_bytes},{local_bytes})".format( local_bytes=results_by_query_id[query_id]['local_bytes'], remote_bytes=results_by_query_id[query_id]['remote_bytes'] ) for query_id in results_by_query_id) sql_template = sqlloader.load_query('insert_into_temp_cache_stats') sql = sql_template.format(values=query_ids_tuples) sql, non_identifier_params = executor.validator.format_identifiers( sql, common_params) executor.execute(sql, non_identifier_params) executor.execute_query( sqlloader, 'update_cache_stats_by_queryid', {**{'hour': start_timestamp[11:16], **common_params}}) end_time = time.perf_counter() logger.info('All done. Elapsed time: {}'.format(end_time - start_time)) def get_cache_stats(query_ids): """Get caching metadata for a list of query_ids. Args: query_ids (list): A list of query_ids. Returns: results (dict): {query_id: {'remote_bytes': 1, 'local_bytes': 0}} """ conn = SnowflakeMetadataConnector(config.SF_CONFIG) conn.authenticate() start_time = time.perf_counter() with ThreadPoolExecutor(max_workers=6) as executor: try: results = executor.map(conn.get_query_stats_by_sfqid, query_ids) except SnowflakeBadResponse as e: # TODO: move logic to the metadata connector logger.info(e) results = { item[0]: item[1] for item in zip(query_ids, results) } end_time = time.perf_counter() logger.info( 'Elapsed time for getting metadata: {}'.format(end_time - start_time)) return results def plot_data(run_timestamp): """Create a graph of daily cache stats and put an image on S3. Args: run_timestamp (datetime.datetime): timestamp """ s3_path = _get_s3_csv_filename(run_timestamp) # unload a file with aggregated stats for a day with SnowflakeSQLExecutor(config.SF_CONFIG) as executor: params = { **config.AWS_CONFIG, **{'db': config.SF_CONFIG['db'], 'schema': config.SF_CONFIG['schema']}, **{'s3_path': s3_path}} executor.execute_query(sqlloader, 'copy_into_s3', params) df = pd.read_csv(s3_path, sep='\t') df.rename( columns={'HOUR': 'Hour', 'NUMBER_OF_QUERIES': 'Number of queries', 'AVERAGE_PERCENTAGE_SCANNED_FROM_CACHE': 'Avg % scanned from cache', 'AVERAGE_EXECUTION_TIME': 'Avg execution time' }, inplace=True) plt.rc('legend', fontsize=30) plot = df.plot( kind='bar', x='Hour', y=['Number of queries', 'Avg % scanned from cache', 'Avg execution time'], figsize=(32, 18), fontsize=28) fig = plot.get_figure() graph_filename = s3_path.replace('csv', 'png').split('/')[-1] # save an image to local filesystem fig.savefig(graph_filename) s3_client = boto3.client('s3') s3_client.upload_file( os.path.join(os.getcwd(), graph_filename), config.S3_BUCKET, 'cache_stats/' + graph_filename) def send_event_to_datadog(): """Send an event to Datadog.""" # TODO: Send a daily event containing a link to a cache stats image, # and data about top slowest queries. creds = json.loads(secret_manager.get_secret('datadog')) datadog_options = { 'api_key': creds['DD_API_KEY'], 'app_key': creds['DD_APP_KEY'] } datadog.initialize(**datadog_options) if __name__ == '__main__': run_timestamp = datetime.datetime.today() - datetime.timedelta(days=1) # send_event_to_datadog() populate_cache_stats_for_a_day(run_timestamp) plot_data(run_timestamp)