Add status argument to jobstep_aggregate_by_status

Summary: jobstep_aggregate_by_status can sometimes be slow (2-3 seconds per request), since we're loading thousands of items from the database. In some cases, like with autoscaler, we only want to get numbers for a particular type of status, such as just counts for jobsteps that are currently `pending_allocation`. This change lets us specify a specific `status` to get stats on, which significantly speeds up these kinds of requests.

Test Plan: Added a unit test for this new functionality.

Reviewers: kylec

Reviewed By: kylec

Subscribers: changesbot, treaster

Differential Revision: https://tails.corp.dropbox.com/D222177

changes/api/jobstep_aggregate_by_status.py

@@ -3,15 +3,32 @@
 from changes.api.base import APIView
 from changes.constants import Status
 from changes.models.jobstep import JobStep
+from flask_restful.reqparse import RequestParser
+
+
+def status_name(value, name):
+    # type: (str, str) -> List[Status]
+    try:
+        return [Status[value]]
+    except KeyError:
+        raise ValueError("The status '{}' is not a known status name.".format(value))
 
 
 class JobStepAggregateByStatusAPIView(APIView):
+    get_parser = RequestParser()
+    default_statuses = (Status.pending_allocation, Status.queued, Status.in_progress, Status.allocated)
+    get_parser.add_argument('status', type=status_name, location='args', default=default_statuses)
+
     def get(self):
         """GET method that returns aggregated data regarding jobsteps.
         Fetch pending, queued, allocated, and in-progress jobsteps from the database.
         Compute some aggregate metrics about them.
         Return the aggregated data in a JSON-friendly format.
 
+        Args (in the form of a query string):
+            status (Optional[str]): A specific status to look for.
+                                    If not specified, will search all statuses.
+
         Returns:
             {
                 'jobsteps': {
@@ -45,9 +62,9 @@ def process_row(agg, jobstep):
                 jobstep_id = jobstep.id
             agg[status] = (count, date_created, jobstep_id)
 
+        args = self.get_parser.parse_args()
         jobsteps = JobStep.query.filter(
-            JobStep.status.in_(
-                [Status.pending_allocation, Status.queued, Status.in_progress, Status.allocated]),
+            JobStep.status.in_(args.status),
         )
 
         by_cluster, by_project = defaultdict(dict), defaultdict(dict)

tests/changes/api/test_jobstep_aggregate_by_status.py

@@ -13,6 +13,52 @@ def get(self, **kwargs):
         query_string = '?' + urlencode(kwargs) if kwargs else ''
         return self.client.get(self.path + query_string)
 
+    def test_get_with_invalid_status(self):
+        raw_resp = self.get(status="meow")
+        assert raw_resp.status_code == 400
+
+    def test_get_with_specific_status(self):
+        project_1 = self.create_project(slug="project_1")
+        build_1 = self.create_build(project_1)
+        job_1 = self.create_job(build_1)
+        jobphase_1 = self.create_jobphase(job_1)
+        now = datetime.datetime.now()
+
+        jobstep_pending_allocation = self.create_jobstep(
+            jobphase_1,
+            status=Status.pending_allocation,
+            date_created=now,
+            cluster="cluster_c")
+        jobstep_in_progress = self.create_jobstep(
+            jobphase_1,
+            status=Status.in_progress,
+            date_created=now,
+            cluster="cluster_a")
+        raw_resp = self.get(status="pending_allocation")
+
+        now_iso = now.isoformat()
+        expected_output = {
+            'jobsteps': {
+                'by_cluster': {
+                    "cluster_c": {
+                        Status.pending_allocation.name:
+                            [1, now_iso, jobstep_pending_allocation.id.get_hex()],
+                    },
+                },
+                'by_project': {
+                    project_1.slug: {
+                        Status.pending_allocation.name:
+                            [1, now_iso, jobstep_pending_allocation.id.get_hex()],
+                    },
+                },
+                'global': {
+                    Status.pending_allocation.name:
+                        [1, now_iso, jobstep_pending_allocation.id.get_hex()],
+                },
+            }
+        }
+        assert self.unserialize(raw_resp) == expected_output
+
     def test_get(self):
         project_1 = self.create_project(slug="project_1")
         build_1 = self.create_build(project_1)