def _update_alert_count(self):
"""Transform Athena query results into alert counts for rules_engine
Args:
query (str): Athena query to run and wait for results
Returns:
dict: Representation of alert counts, where key is the rule name
and value is the alert count (int) since this rule was staged
"""
query = StagingStatistic.construct_compound_count_query(
list(self._staging_stats.values()))
LOGGER.debug('Running compound query for alert count: \'%s\'', query)
for page, results in enumerate(
self._athena_client.query_result_paginator(query)):
for i, row in enumerate(results['ResultSet']['Rows']):
if page == 0 and i == 0: # skip header row included in first page only
continue
row_values = [list(data.values())[0] for data in row['Data']]
rule_name, alert_count = row_values[0], int(row_values[1])
LOGGER.debug('Found %d alerts for rule \'%s\'', alert_count,
rule_name)
self._staging_stats[rule_name].alert_count = alert_count
def test_construct_compound_count_query(self):
"""StagingStatistic - Construct Compound Count Query"""
query = StagingStatistic.construct_compound_count_query(
[self.statistic, self.statistic])
expected_query = (
"SELECT rule_name, count(*) AS count "
"FROM alerts WHERE "
"(dt >= '2000-01-01-01' AND rule_name = 'test_rule') OR "
"(dt >= '2000-01-01-01' AND rule_name = 'test_rule') "
"GROUP BY rule_name")
assert_equal(query, expected_query)