"""Tasks for the datadog_caching_monitor SWF workflow.""" from concurrent.futures import ThreadPoolExecutor import datetime import backoff import datadog from datadog.api import exceptions from garcon import task from snowflake.connector.errors import ProgrammingError from snowflake_connector.metadata_connector import SnowflakeBadResponse, \ SnowflakeMetadataConnector from swf_monitoring import registered_executors from swf_monitoring.flows.datadog_caching_monitor import config executor_class = registered_executors.get(config.FEED_NAME) def _create_metric(metric_name, event): return { 'metric': 'snowflake_caching_monitor.{}'.format(metric_name), 'points': [( event['start_time'].timestamp(), int(event[metric_name]))], 'tags': [ 'snowflake_caching_monitor', 'correlation_id:{}'.format( event['correlation_id']), 'query_id:{}'.format(event['query_id']), 'day:{}'.format(event['day']), 'hour:{}'.format(event['hour']), 'user_name:{}'.format(event['user_name']), 'warehouse_name:{}'.format(event['warehouse_name'])] } @backoff.on_exception( backoff.expo, SnowflakeBadResponse) def _get_cache_stats(activity, query_ids): """Get caching metadata for a list of query_ids. Args: activity: An activity object. 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() with ThreadPoolExecutor(max_workers=6) as executor: try: results = executor.map(conn.get_query_stats_by_sfqid, query_ids) except SnowflakeBadResponse as e: activity.logger.info( 'Snowflake API exception {}, backoff and retry...'.format( str(e))) raise e results = { item[0]: item[1] for item in zip(query_ids, results) } return results @backoff.on_exception( backoff.expo, (exceptions.ProxyError, exceptions.HTTPError, exceptions.HttpBackoff, exceptions.HttpTimeout)) def _send_metrics(metrics): """Send a list of metrics using backoff/retry decorator. Args: metrics (list): A list of metrics to send. """ res = datadog.api.Metric.send(metrics, compress_payload=True) if res.get('status') != 'ok': raise exceptions.HttpBackoff(5) # will be handled by backoff @task.decorate(timeout=60 * 15) def populate_cache_stats(activity): """Populate CACHE_STATS_BY_QUERYID table with the latest queries.""" activity.logger.info('Populating stats table...') run_date = datetime.datetime.now().date() common_params = { 'date': run_date, 'db': config.SF_CONFIG['db'], 'schema': config.SF_CONFIG['schema'] } with executor_class( config.SF_CONFIG, statement_timeout_in_seconds=60*10) as executor: # clean up the table from rows submitted to DD in order to # avoid duplication executor.fetchone_query( executor.sqlloader, 'delete_from_cache_stats_by_queryid', common_params) activity.logger.info( 'The table cleared from queries not send to Datadog...') executor.fetchone_query( executor_class.sqlloader, 'create_temp_stats', common_params) # populate CACHE_STATS_CACHE_STATS_BY_QUERYID for each warehouse for warehouse in config.get_snowflake_warehouses_to_monitor(): activity.logger.info( 'Populating table for {}...'.format(warehouse)) insert_params = { **common_params, **{'users_to_monitor': config.get_snowflake_users_to_monitor(), 'warehouse_to_monitor': warehouse, 'database_to_monitor': 'FACTS', 'schema_to_monitor': 'PROD' } } try: executor.fetchone_query( executor_class.sqlloader, 'merge_into_cache_stats_by_queryid', insert_params) except ProgrammingError as e: if 'Statement reached its statement or warehouse timeout' \ in str(e): return { 'stop': True, 'result': 'Statement reached its timeout'} else: raise e activity.logger.info('Stats table populated!') @task.decorate(timeout=60*5) def update_metadata_in_cache_stats_table(activity): """Get metadata from Snowflake API and update the main table.""" run_date = datetime.datetime.now().date() common_params = { 'date': run_date, 'db': config.SF_CONFIG['db'], 'schema': config.SF_CONFIG['schema'] } with executor_class(config.SF_CONFIG) as executor: query_ids = [ query_id[0] for query_id in executor.fetchall_query( executor_class.sqlloader, 'get_latest_queryids_from_cache_stats_by_queryid', common_params) if query_id[1] is True # get metadata only for queries with bytes_scanned > 0 ] activity.logger.info('Getting metadata from Snowflake API...') results_by_query_id = _get_cache_stats(activity, 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 = executor_class.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( executor_class.sqlloader, 'update_cache_stats_by_queryid', common_params) activity.logger.info('Metadata updated!') @task.decorate(timeout=60 * 5) def send_query_level_metrics_to_datadog(activity): """Send query-level metrics to Datadog.""" activity.logger.info('Sending metrics to Datadog...') run_date = datetime.datetime.today().date() query_ids_sent_to_datadog = [] datadog.initialize(**{'api_key': config.DATADOG_API_KEY}) with executor_class(config.SF_CONFIG) as executor: params = { 'date': run_date, 'db': config.SF_CONFIG['db'], 'schema': config.SF_CONFIG['schema'], 'day': run_date } result = executor.fetchall_query( executor_class.sqlloader, 'get_all_not_submitted_queries', params) with ThreadPoolExecutor(max_workers=5) as thread_executor: for row in result: day = str(run_date).replace('-', '') event = dict(zip( ['hour', 'database_name', 'schema_name', 'warehouse_name', 'user_name', 'query_id', 'start_time', 'compilation_time', 'execution_time', 'bytes_scanned', 'remote_bytes', 'local_bytes', 'correlation_id'], row)) if event['bytes_scanned']: percentage_scanned_from_cache = ( event['local_bytes'] / ( event['local_bytes'] + event['remote_bytes']) * 100.0) else: percentage_scanned_from_cache = 100.0 event.update( { 'percentage_scanned_from_cache': percentage_scanned_from_cache, 'day': day}) thread_executor.submit(_send_metrics, [ _create_metric('compilation_time', event), _create_metric('execution_time', event), _create_metric('remote_bytes', event), _create_metric('local_bytes', event), _create_metric('percentage_scanned_from_cache', event), _create_metric('bytes_scanned', event), ]) thread_executor.submit( query_ids_sent_to_datadog.append, event['query_id']) params = { 'date': run_date, 'db': config.SF_CONFIG['db'], 'schema': config.SF_CONFIG['schema'], 'query_ids': query_ids_sent_to_datadog } if query_ids_sent_to_datadog: executor.fetchall_query( executor_class.sqlloader, 'update_submitted_to_datadog_column', params) activity.logger.info('Metrics sent!') @task.decorate(timeout=60 * 5) def send_hits_to_kw_cache_percentage_to_datadog(activity): """Send daily aggregated metrics to Datadog.""" activity.logger.info('Sending daily aggregated metrics to Datadog...') run_date = datetime.datetime.today().date() datadog.initialize(**{'api_key': config.DATADOG_API_KEY}) with executor_class(config.SF_CONFIG) as executor: params = { 'date': run_date, 'db': config.SF_CONFIG['db'], 'schema': config.SF_CONFIG['schema'], 'day': run_date } result = executor.fetchone_query( executor_class.sqlloader, 'get_hits_to_kw_cache_percentage', params) _send_metrics([{ 'metric': 'snowflake_caching_monitor.hits_to_kw_cache_percentage', 'points': [result[0]], 'tags': [ 'snowflake_caching_monitor', 'day:{}'.format(str(run_date).replace('-', ''))]}]) activity.logger.info('Updated kv hits percentage for a day sent!')