diff --git a/sidecar/app/query/base.py b/sidecar/app/query/base.py index d2509a2..992d3ee 100644 --- a/sidecar/app/query/base.py +++ b/sidecar/app/query/base.py @@ -201,8 +201,11 @@ def run_query(self, query: Query): ) self.running_queries[query.query_id] = query - query.start() - del self.running_queries[query.query_id] + try: + query.start() + finally: + # always remove this + del self.running_queries[query.query_id] @property def capacity_available(self): diff --git a/sidecar/tests/app/query/test_base.py b/sidecar/tests/app/query/test_base.py index 1c26bdd..9757a78 100644 --- a/sidecar/tests/app/query/test_base.py +++ b/sidecar/tests/app/query/test_base.py @@ -119,3 +119,22 @@ def fake_start(): ) as mock_start: query_manager.run_query(query) mock_start.assert_called_once() + + +def test_query_manager_run_query_exception(): + query_manager = QueryManager(max_parallel_queries=1) + query = Query(str(uuid4())) + + def mock_exception(): + raise Exception + + with mock.patch( + "sidecar.app.query.base.Query.start", side_effect=mock_exception + ) as mock_start: + with pytest.raises(Exception): + query_manager.run_query(query) + + mock_start.assert_called_once() + + assert query.query_id not in query_manager.running_queries + assert query_manager.capacity_available