fix bug where an edge case exception causes a query to not be removed… (

#88) * 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 <[email protected]> --------- Co-authored-by: Alex Koshelev <[email protected]>
private-attribution · Sep 27, 2024 · d3253a8 · d3253a8
1 parent 40e7c9b
commit d3253a8
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git 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
@@ -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