Added bigquery support and bumped package version

Author: oravi

monitor/alerts.py

@@ -52,7 +52,7 @@ def __init__(self, alert_id, database_name, schema_name, table_name, detected_at
         super().__init__(alert_id)
         self.table_name = '.'.join([database_name, schema_name, table_name]).lower()
         self.detected_at = convert_utc_time_to_local_time(detected_at).strftime('%Y-%m-%d %H:%M:%S')
-        self.description = description[0].upper() + description[1:].lower()
+        self.description = description[0].upper() + description[1:].lower() if description else ''
         self.change_type = ' '.join([word[0].upper() + word[1:] for word in sub_type.split('_')])
 
     def to_slack_message(self) -> dict:
@@ -114,7 +114,7 @@ def __init__(self, alert_id, database_name, schema_name, table_name, detected_at
         super().__init__(alert_id)
         self.table_name = '.'.join([database_name, schema_name, table_name]).lower()
         self.detected_at = convert_utc_time_to_local_time(detected_at).strftime('%Y-%m-%d %H:%M:%S')
-        self.description = description[0].upper() + description[1:].lower()
+        self.description = description if description else ''
         self.anomaly_type = ' '.join([word[0].upper() + word[1:] for word in sub_type.split('_')])
 
     def to_slack_message(self) -> dict:

monitor/data_monitoring.py

@@ -1,12 +1,12 @@
 import os
 from typing import Any
-import snowflake.connector.errors
 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
+from utils.dbt import get_snowflake_client, get_bigquery_client
 from utils.log import get_logger
+from google.cloud import bigquery
 
 logger = get_logger(__name__)
 FILE_DIR = os.path.dirname(__file__)
@@ -20,20 +20,6 @@ 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)
 
-    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); 
-    """
-
-    COUNT_ROWS_QUERY = None
-
     def __init__(self, config: Config, db_connection: Any) -> None:
         self.config = config
         self.dbt_runner = DbtRunner(self.DBT_PROJECT_PATH, self.config.profiles_dir)
@@ -45,6 +31,9 @@ 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")
 
@@ -55,16 +44,10 @@ def _run_query(self, query: str, params: tuple = None) -> list:
         pass
 
     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)}")
+        pass
 
     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
+        pass
 
     def _send_to_slack(self, alerts: [Alert]) -> None:
         slack_webhook = self.config.slack_notification_webhook
@@ -150,6 +133,18 @@ def properties(self):
 
 
 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)
 
@@ -163,5 +158,56 @@ def _run_query(self, query: str, params: tuple = None) -> list:
             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/packages.yml

@@ -1,3 +1,3 @@
 packages:
   - package: elementary-data/elementary
-    version: 0.3.1
+    version: 0.3.3