From 5c56c7339a67c66cde7e6568a68bae451f34413b Mon Sep 17 00:00:00 2001
From: David Schultz <david.schultz@icecube.wisc.edu>
Date: Thu, 10 Oct 2024 15:05:10 -0500
Subject: [PATCH] only search for idle and waiting task counts when doing
 queue_tasks

---
 iceprod/rest/handlers/tasks.py         | 32 ++++++++++++++++++++++----
 iceprod/scheduled_tasks/queue_tasks.py |  3 ++-
 tests/rest/tasks_test.py               |  3 +++
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/iceprod/rest/handlers/tasks.py b/iceprod/rest/handlers/tasks.py
index 5ff139125..fa318b414 100644
--- 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
index 71fb64867..7234e8882 100644
--- 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
index 4123c8279..65d2525fc 100644
--- 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'])