[WIP] Closes #415: Make data source health monitoring an extension.

redash/extensions.py

@@ -8,7 +8,19 @@ def init_extensions(app):
     if not hasattr(app, 'redash_extensions'):
         app.redash_extensions = {}
 
+    if not hasattr(app, 'task_extensions'):
+        app.task_extensions = {}
+
     for entry_point in iter_entry_points('redash.extensions'):
         app.logger.info('Loading Redash extension %s.', entry_point.name)
         extension = entry_point.load()
-        app.redash_extensions[entry_point.name] = extension(app)
+        extension_response = extension(app)
+        extension_repository = app.redash_extensions
+
+        # This is an interval task
+        if (type(extension_response) == dict and
+                "task" in extension_response and
+                "interval_in_seconds" in extension_response):
+            extension_repository = app.task_extensions
+
+        extension_repository[entry_point.name] = extension_response

redash/monitor.py

@@ -1,4 +1,3 @@
-import json
 from redash import redis_connection, models, __version__, settings
 
 
@@ -18,7 +17,6 @@ def get_status():
     status['workers'] = []
 
     status['manager'] = redis_connection.hgetall('redash:status')
-    status['data_sources'] = json.loads(redis_connection.get('data_sources:health') or '{}')
 
     queues = {}
     for ds in models.DataSource.query:

redash/query_runner/__init__.py

@@ -102,12 +102,10 @@ def get_data_source_version(self):
 
         return version
 
-    def test_connection(self, custom_query_text=None):
+    def test_connection(self):
         if self.noop_query is None:
             raise NotImplementedError()
-
-        query_text = custom_query_text or self.noop_query
-        data, error = self.run_query(query_text, None)
+        data, error = self.run_query(self.noop_query, None)
 
         if error is not None:
             raise Exception(error)

redash/settings/__init__.py

@@ -1,7 +1,7 @@
 import os
 from funcy import distinct, remove
 
-from .helpers import parse_db_url, fix_assets_path, array_from_string, parse_boolean, int_or_none, set_from_string, dict_from_string
+from .helpers import parse_db_url, fix_assets_path, array_from_string, parse_boolean, int_or_none, set_from_string
 
 
 def all_settings():
@@ -241,14 +241,3 @@ def all_settings():
 # Allow Parameters in Embeds
 # WARNING: With this option enabled, Redash reads query parameters from the request URL (risk of SQL injection!)
 ALLOW_PARAMETERS_IN_EMBEDS = parse_boolean(os.environ.get("REDASH_ALLOW_PARAMETERS_IN_EMBEDS", "false"))
-
-# Allow for a map of custom queries to test data source performance and availability.
-# A sample map may look like:
-# {
-#    "1": "select 1;",
-#    "5": "select 1;"
-# }
-CUSTOM_HEALTH_QUERIES = dict_from_string(os.environ.get("REDASH_CUSTOM_HEALTH_QUERIES", ""))
-
-# Frequency of health query runs in minutes (12 hours by default)
-HEALTH_QUERIES_REFRESH_SCHEDULE = int(os.environ.get("REDASH_HEALTH_QUERIES_REFRESH_SCHEDULE", 720))

redash/settings/helpers.py

@@ -33,12 +33,6 @@ def array_from_string(s):
 
     return [item.strip() for item in array]
 
-def dict_from_string(s):
-    try:
-        return json.loads(s)
-    except ValueError:
-        return {}
-
 
 def set_from_string(s):
     return set(array_from_string(s))

redash/tasks/__init__.py

@@ -1,4 +1,3 @@
 from .general import record_event, version_check, send_mail
-from .health import health_status
 from .queries import QueryTask, refresh_queries, refresh_schemas, cleanup_tasks, cleanup_query_results, execute_query
 from .alerts import check_alerts_for_query

redash/worker.py

@@ -16,10 +16,6 @@
                 include='redash.tasks')
 
 celery_schedule = {
-    'health_status': {
-        'task': 'redash.tasks.health_status',
-        'schedule': timedelta(minutes=settings.HEALTH_QUERIES_REFRESH_SCHEDULE)
-    },
     'refresh_queries': {
         'task': 'redash.tasks.refresh_queries',
         'schedule': timedelta(seconds=30)
@@ -82,3 +78,12 @@ def __call__(self, *args, **kwargs):
 def init_celery_flask_app(**kwargs):
     app = create_app()
     app.app_context().push()
+
+@celery.on_after_configure.connect
+def test(sender, **kwargs):
+    app = create_app()
+    for task in app.task_extensions.values():
+        interval = task["interval_in_seconds"]
+        task = task["task"]
+        task.delay()
+        sender.add_periodic_task(interval, task.s(), name=task.name)