From d3253a854386393659f3e0f813383204e24e99c4 Mon Sep 17 00:00:00 2001 From: Erik Taubeneck Date: Thu, 26 Sep 2024 20:58:58 -0700 Subject: [PATCH] =?UTF-8?q?fix=20bug=20where=20an=20edge=20case=20exceptio?= =?UTF-8?q?n=20causes=20a=20query=20to=20not=20be=20removed=E2=80=A6=20(#8?= =?UTF-8?q?8)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix bug where an edge case exception causes a query to not be removed from active running queries * add test coverage for failure case * Update test_base.py Co-authored-by: Alex Koshelev --------- Co-authored-by: Alex Koshelev --- sidecar/app/query/base.py | 7 +++++-- sidecar/tests/app/query/test_base.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) 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