only search for idle and waiting task counts when doing queue_tasks

WIPACrepo · Oct 10, 2024 · 5c56c73 · 5c56c73
1 parent 3c5e0d6
commit 5c56c73
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 5 deletions.
diff --git a/iceprod/rest/handlers/tasks.py b/iceprod/rest/handlers/tasks.py
@@ -315,13 +315,28 @@ async def get(self):
         """
         Get the task counts for all tasks, group by status.
 
+        Params (optional):
+            status: | separated list of task status to filter by
+
         Returns:
             dict: {<status>: num}
         """
-        ret = {}
-        for status in TASK_STATUS:
-            ret[status] = await self.db.tasks.count_documents({"status": status})
+        match = {}
+        status = self.get_argument('status', None)
+        if status:
+            status_list = status.split('|')
+            if any(s not in TASK_STATUS for s in status_list):
+                raise tornado.web.HTTPError(400, reaosn='Unknown task status')
+            match['status'] = {'$in': status_list}
 
+        ret = {}
+        cursor = self.db.tasks.aggregate([
+            {'$match': match},
+            {'$group': {'_id': '$status', 'total': {'$sum': 1}}},
+        ])
+        ret = {}
+        async for row in cursor:
+            ret[row['_id']] = row['total']
         ret2 = {}
         for k in sorted(ret, key=task_status_sort):
             ret2[k] = ret[k]
@@ -552,8 +567,17 @@ async def get(self, dataset_id):
         Returns:
             dict: {<status>: num}
         """
+        match = {'dataset_id': dataset_id}
+        status = self.get_argument('status', None)
+        if status:
+            status_list = status.split('|')
+            if any(s not in TASK_STATUS for s in status_list):
+                raise tornado.web.HTTPError(400, reaosn='Unknown task status')
+            match['status'] = {'$in': status_list}
+
+        ret = {}
         cursor = self.db.tasks.aggregate([
-            {'$match': {'dataset_id': dataset_id}},
+            {'$match': match},
             {'$group': {'_id': '$status', 'total': {'$sum': 1}}},
         ])
         ret = {}

diff --git a/iceprod/scheduled_tasks/queue_tasks.py b/iceprod/scheduled_tasks/queue_tasks.py
@@ -33,7 +33,8 @@ async def run(rest_client, dataset_id=None, ntasks=NTASKS, ntasks_per_cycle=NTAS
             route = f'/datasets/{dataset_id}/task_counts/status'
         else:
             route = '/task_counts/status'
-        tasks = await rest_client.request('GET', route)
+        args = {'status': 'idle|waiting'}
+        tasks = await rest_client.request('GET', route, args)
         if 'idle' in tasks:
             num_tasks_waiting = tasks['idle']
         if 'waiting' in tasks:

diff --git a/tests/rest/tasks_test.py b/tests/rest/tasks_test.py
@@ -262,6 +262,9 @@ async def test_rest_tasks_dataset_counts_status(server):
     ret = await client.request('GET', f'/datasets/{data["dataset_id"]}/task_counts/status')
     assert ret == {states.TASK_STATUS_START: 1}
 
+    ret = await client.request('GET', f'/datasets/{data["dataset_id"]}/task_counts/status?status=complete')
+    assert ret == {}
+
 
 async def test_rest_tasks_dataset_counts_name_status(server):
     client = server(roles=['system'])