Merge pull request #67 from elementary-data/redshift

Redshift

Author: oravi

monitor/alerts.py

@@ -2,6 +2,7 @@
 import json
 from exceptions.exceptions import InvalidAlertType
 from utils.time import convert_utc_time_to_local_time
+from datetime import datetime
 
 
 class Alert(object):
@@ -11,9 +12,10 @@ def __init__(self, alert_id) -> None:
         self.alert_id = alert_id
 
     @staticmethod
-    def create_alert_from_row(alert_row: list) -> 'Alert':
+    def create_alert_from_row(alert_row: dict) -> 'Alert':
         alert_id, detected_at, database_name, schema_name, table_name, column_name, alert_type, sub_type, \
-            alert_description = alert_row
+            alert_description = alert_row.values()
+        detected_at = datetime.fromisoformat(detected_at)
         if alert_type == 'schema_change':
             return SchemaChangeAlert(alert_id, database_name, schema_name, table_name, detected_at, sub_type,
                                      alert_description)

monitor/cli.py

@@ -70,7 +70,7 @@ def monitor(ctx, config_dir, profiles_dir, update_dbt_package, full_refresh_dbt_
     anonymous_tracking = AnonymousTracking(config)
     track_cli_start(anonymous_tracking, 'monitor', get_cli_properties(), ctx.command.name)
     try:
-        data_monitoring = DataMonitoring.create_data_monitoring(config)
+        data_monitoring = DataMonitoring(config)
         data_monitoring.run(update_dbt_package, full_refresh_dbt_package)
         track_cli_end(anonymous_tracking, 'monitor', data_monitoring.properties(), ctx.command.name)
     except Exception as exc:

monitor/data_monitoring.py

@@ -1,12 +1,10 @@
 import os
-from typing import Any
-from exceptions.exceptions import ConfigError
 from monitor.alerts import Alert
 from monitor.dbt_runner import DbtRunner
 from config.config import Config
-from utils.dbt import get_snowflake_client, get_bigquery_client
 from utils.log import get_logger
-from google.cloud import bigquery
+import json
+from alive_progress import alive_it
 
 logger = get_logger(__name__)
 FILE_DIR = os.path.dirname(__file__)
@@ -20,40 +18,43 @@ class DataMonitoring(object):
     DBT_PROJECT_MODULES_PATH = os.path.join(DBT_PROJECT_PATH, 'dbt_modules', DBT_PACKAGE_NAME)
     DBT_PROJECT_PACKAGES_PATH = os.path.join(DBT_PROJECT_PATH, 'dbt_packages', DBT_PACKAGE_NAME)
 
-    def __init__(self, config: Config, db_connection: Any) -> None:
+    def __init__(self, config: Config) -> None:
         self.config = config
         self.dbt_runner = DbtRunner(self.DBT_PROJECT_PATH, self.config.profiles_dir)
-        self.db_connection = db_connection
         self.execution_properties = {}
 
-    @staticmethod
-    def create_data_monitoring(config: Config) -> 'DataMonitoring':
-        if config.platform == 'snowflake':
-            snowflake_conn = get_snowflake_client(config.credentials, server_side_binding=False)
-            return SnowflakeDataMonitoring(config, snowflake_conn)
-        elif config.platform == 'bigquery':
-            bigquery_client = get_bigquery_client(config.credentials)
-            return BigQueryDataMonitoring(config, bigquery_client)
-        else:
-            raise ConfigError("Unsupported platform")
-
     def _dbt_package_exists(self) -> bool:
         return os.path.exists(self.DBT_PROJECT_PACKAGES_PATH) or os.path.exists(self.DBT_PROJECT_MODULES_PATH)
 
-    def _run_query(self, query: str, params: tuple = None) -> list:
-        pass
+    @staticmethod
+    def _split_list_to_chunks(items: list, chunk_size: int = 50) -> [list]:
+        chunk_list = []
+        for i in range(0, len(items), chunk_size):
+            chunk_list.append(items[i: i + chunk_size])
+        return chunk_list
 
     def _update_sent_alerts(self, alert_ids) -> None:
-        pass
+        alert_ids_chunks = self._split_list_to_chunks(alert_ids)
+        for alert_ids_chunk in alert_ids_chunks:
+            self.dbt_runner.run_operation(macro_name='update_sent_alerts',
+                                          macro_args={'alert_ids': alert_ids_chunk},
+                                          json_logs=False)
 
     def _query_alerts(self) -> list:
-        pass
+        json_alert_rows = self.dbt_runner.run_operation(macro_name='get_new_alerts')
+        self.execution_properties['alert_rows'] = len(json_alert_rows)
+        alerts = []
+        for json_alert_row in json_alert_rows:
+            alert_row = json.loads(json_alert_row)
+            alerts.append(Alert.create_alert_from_row(alert_row))
+        return alerts
 
     def _send_to_slack(self, alerts: [Alert]) -> None:
         slack_webhook = self.config.slack_notification_webhook
         if slack_webhook is not None:
             sent_alerts = []
-            for alert in alerts:
+            alerts_with_progress_bar = alive_it(alerts, title="Sending alerts")
+            for alert in alerts_with_progress_bar:
                 alert.send_to_slack(slack_webhook, self.config.is_slack_workflow)
                 sent_alerts.append(alert.id)
 
@@ -132,82 +133,4 @@ def properties(self):
         return data_monitoring_properties
 
 
-class SnowflakeDataMonitoring(DataMonitoring):
-    SELECT_NEW_ALERTS_QUERY = """
-        SELECT alert_id, detected_at, database_name, schema_name, table_name, column_name, alert_type, sub_type, 
-               alert_description
-            FROM ALERTS
-            WHERE alert_sent = FALSE;
-    """
-
-    UPDATE_SENT_ALERTS = """
-            UPDATE ALERTS set alert_sent = TRUE
-                WHERE alert_id IN (%s); 
-     """
-
-    def __init__(self, config: 'Config', db_connection: Any):
-        super().__init__(config, db_connection)
-
-    def _run_query(self, query: str, params: tuple = None) -> list:
-        with self.db_connection.cursor() as cursor:
-            if params is not None:
-                cursor.execute(query, params)
-            else:
-                cursor.execute(query)
-
-            results = cursor.fetchall()
-            return results
-
-    def _update_sent_alerts(self, alert_ids) -> None:
-        results = self._run_query(self.UPDATE_SENT_ALERTS, (alert_ids,))
-        logger.debug(f"Updated sent alerts -\n{str(results)}")
-
-    def _query_alerts(self) -> list:
-        alert_rows = self._run_query(self.SELECT_NEW_ALERTS_QUERY)
-        self.execution_properties['alert_rows'] = len(alert_rows)
-        alerts = []
-        for alert_row in alert_rows:
-            alerts.append(Alert.create_alert_from_row(alert_row))
-        return alerts
-
-
-class BigQueryDataMonitoring(DataMonitoring):
-    SELECT_NEW_ALERTS_QUERY = """
-        SELECT alert_id, detected_at, database_name, schema_name, table_name, column_name, alert_type, sub_type, 
-               alert_description
-            FROM {dataset}.alerts
-            WHERE alert_sent = FALSE;
-    """
-
-    UPDATE_SENT_ALERTS = """
-            UPDATE {dataset}.alerts set alert_sent = TRUE
-                WHERE alert_id IN UNNEST(@alert_ids); 
-    """
-
-    def __init__(self, config: 'Config', db_connection: Any):
-        super().__init__(config, db_connection)
-
-    def _run_query(self, query: str, params: list = None) -> list:
-        if params is not None:
-            job_config = bigquery.QueryJobConfig(
-                query_parameters=params
-            )
-            job = self.db_connection.query(query, job_config=job_config)
-        else:
-            job = self.db_connection.query(query)
-
-        return list(job.result())
-
-    def _update_sent_alerts(self, alert_ids) -> None:
-        params = [bigquery.ArrayQueryParameter("alert_ids", "STRING", alert_ids)]
-        results = self._run_query(self.UPDATE_SENT_ALERTS.format(dataset=self.config.credentials.schema), params)
-        logger.debug(f"Updated sent alerts -\n{str(results)}")
-
-    def _query_alerts(self) -> list:
-        alert_rows = self._run_query(self.SELECT_NEW_ALERTS_QUERY.format(dataset=self.config.credentials.schema))
-        self.execution_properties['alert_rows'] = len(alert_rows)
-        alerts = []
-        for alert_row in alert_rows:
-            alerts.append(Alert.create_alert_from_row(alert_row))
-        return alerts
 

monitor/dbt_project/macros/get_alerts_time_limit.sql

@@ -0,0 +1,5 @@
+{% macro get_alerts_time_limit(days_back=31) %}
+    {% set today = dbt_utils.date_trunc('day', dbt_utils.current_timestamp()) %}
+    {% set datetime_limit = dbt_utils.dateadd('day', days_back * -1, today) %}
+    {{ return(elementary.cast_as_timestamp(datetime_limit)) }}
+{% endmacro %}

monitor/dbt_project/macros/get_new_alerts.sql

@@ -0,0 +1,25 @@
+{% macro get_new_alerts() %}
+    -- depends_on: {{ ref('alerts') }}
+    {% set current_date = dbt_utils.date_trunc('day', dbt_utils.current_timestamp()) %}
+    {% set select_new_alerts_query %}
+        SELECT alert_id, detected_at, database_name, schema_name, table_name, column_name, alert_type, sub_type,
+               alert_description
+        FROM {{ ref('alerts') }}
+        WHERE alert_sent = FALSE and detected_at >= {{ get_alerts_time_limit() }}
+    {% endset %}
+    {% set results = run_query(select_new_alerts_query) %}
+    {% set new_alerts = [] %}
+    {% for result in results %}
+        {% set new_alert_dict = {'alert_id': result[0],
+                                 'detected_at': result[1].isoformat(),
+                                 'database_name': result[2],
+                                 'schema_name': result[3],
+                                 'table_name': result[4],
+                                 'column_name': result[5],
+                                 'alert_type': result[6],
+                                 'sub_type': result[7],
+                                 'alert_description': result[8]} %}
+        {% set new_alert_json = tojson(new_alert_dict) %}
+        {% do elementary.edr_log(new_alert_json) %}
+    {% endfor %}
+{% endmacro %}

monitor/dbt_project/macros/update_sent_alerts.sql

@@ -0,0 +1,12 @@
+{% macro update_sent_alerts(alert_ids) %}
+    -- depends_on: {{ ref('alerts') }}
+    {% if alert_ids %}
+        {% set update_sent_alerts_query %}
+            UPDATE {{ ref('alerts') }} set alert_sent = TRUE
+            WHERE alert_id IN {{ elementary.strings_list_to_tuple(alert_ids) }} and alert_sent = FALSE and
+                detected_at >= {{ get_alerts_time_limit() }}
+        {% endset %}
+        {% set results = dbt_utils.get_query_results_as_dict(update_sent_alerts_query) %}
+        {% do elementary.edr_log(results) %}
+    {% endif %}
+{% endmacro %}

monitor/dbt_project/packages.yml

@@ -1,3 +1,5 @@
 packages:
+  - package: dbt-labs/dbt_utils
+    version: 0.7.6
   - package: elementary-data/elementary
-    version: 0.3.5
+    version: 0.3.11

monitor/dbt_runner.py

@@ -50,8 +50,12 @@ def snapshot(self) -> bool:
         success, _ = self._run_command(['snapshot'])
         return success
 
-    def run_operation(self, macro_name, json_logs=True) -> list:
-        success, command_output = self._run_command(['run-operation', macro_name], json_logs)
+    def run_operation(self, macro_name: str, json_logs: bool = True, macro_args: dict = None) -> list:
+        command_args = ['run-operation', macro_name]
+        if macro_args is not None:
+            json_args = json.dumps(macro_args)
+            command_args.extend(['--args', json_args])
+        success, command_output = self._run_command(command_args, json_logs)
         run_operation_results = []
         if json_logs:
             json_messages = command_output.splitlines()