Propagate query execution errors from Celery tasks properly (re #290)

redash/tasks/queries.py

@@ -466,6 +466,8 @@ def run(self):
                 self.scheduled_query = models.db.session.merge(self.scheduled_query, load=False)
                 self.scheduled_query.schedule_failures += 1
                 models.db.session.add(self.scheduled_query)
+            models.db.session.commit()
+            raise result
         else:
             if (self.scheduled_query and self.scheduled_query.schedule_failures > 0):
                 self.scheduled_query = models.db.session.merge(self.scheduled_query, load=False)
@@ -482,8 +484,8 @@ def run(self):
             self._log_progress('finished')
 
             result = query_result.id
-        models.db.session.commit()
-        return result
+            models.db.session.commit()
+            return result
 
     def _annotate_query(self, query_runner):
         if query_runner.annotate_query():

tests/tasks/test_queries.py

@@ -7,7 +7,8 @@
 from tests import BaseTestCase
 from redash import redis_connection, models
 from redash.query_runner.pg import PostgreSQL
-from redash.tasks.queries import QueryTaskTracker, enqueue_query, execute_query
+from redash.tasks.queries import (QueryExecutionError, QueryTaskTracker,
+                                  enqueue_query, execute_query)
 
 
 class TestPrune(TestCase):
@@ -113,11 +114,15 @@ def test_failure_scheduled(self):
                         {'routing_key': 'test'})
         q = self.factory.create_query(query_text="SELECT 1, 2", schedule=300)
         with cm, mock.patch.object(PostgreSQL, "run_query") as qr:
-            qr.exception = ValueError("broken")
-            execute_query("SELECT 1, 2", self.factory.data_source.id, {}, scheduled_query_id=q.id)
+            qr.side_effect = ValueError("broken")
+            with self.assertRaises(QueryExecutionError):
+                execute_query("SELECT 1, 2", self.factory.data_source.id, {},
+                              scheduled_query_id=q.id)
             q = models.Query.get_by_id(q.id)
             self.assertEqual(q.schedule_failures, 1)
-            execute_query("SELECT 1, 2", self.factory.data_source.id, {}, scheduled_query_id=q.id)
+            with self.assertRaises(QueryExecutionError):
+                execute_query("SELECT 1, 2", self.factory.data_source.id, {},
+                              scheduled_query_id=q.id)
             q = models.Query.get_by_id(q.id)
             self.assertEqual(q.schedule_failures, 2)
 
@@ -129,10 +134,11 @@ def test_success_after_failure(self):
                         {'routing_key': 'test'})
         q = self.factory.create_query(query_text="SELECT 1, 2", schedule=300)
         with cm, mock.patch.object(PostgreSQL, "run_query") as qr:
-            qr.exception = ValueError("broken")
-            execute_query("SELECT 1, 2",
-                          self.factory.data_source.id, {},
-                          scheduled_query_id=q.id)
+            qr.side_effect = ValueError("broken")
+            with self.assertRaises(QueryExecutionError):
+                execute_query("SELECT 1, 2",
+                              self.factory.data_source.id, {},
+                              scheduled_query_id=q.id)
             q = models.Query.get_by_id(q.id)
             self.assertEqual(q.schedule_failures, 1)