From 41aba63ae94128552eb3c03ceb6d0892c95fc5f2 Mon Sep 17 00:00:00 2001 From: Allen Short Date: Tue, 27 Feb 2018 21:42:33 +0000 Subject: [PATCH] Propagate query execution errors from Celery tasks properly (re #290) --- redash/tasks/queries.py | 6 ++++-- tests/tasks/test_queries.py | 22 ++++++++++++++-------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/redash/tasks/queries.py b/redash/tasks/queries.py index 391a9b2fae..087fba06f1 100644 --- a/redash/tasks/queries.py +++ b/redash/tasks/queries.py @@ -464,6 +464,8 @@ def run(self): if self.scheduled_query: 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): @@ -479,8 +481,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(): diff --git a/tests/tasks/test_queries.py b/tests/tasks/test_queries.py index 06dddc1fcc..77217816c0 100644 --- a/tests/tasks/test_queries.py +++ b/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,10 +114,14 @@ 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) 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) @@ -128,10 +133,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)