add aggregate operation (#27)

* add aggregate operation Add support for aggregate operations on queryset, user can pass multiple aggregate operations, also all operations are stopped after aggregate. Added test cases for aggregate operation Signed-off-by: Priyank Singh <[email protected]>
vmware-archive · May 16, 2023 · 86b269b · 86b269b
1 parent 4b98811
commit 86b269b
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 3 deletions.
diff --git a/bridgeql/django/models.py b/bridgeql/django/models.py
@@ -9,7 +9,7 @@
     ValidationError,
     ObjectDoesNotExist
 )
-from django.db.models import QuerySet
+from django.db.models import QuerySet, aggregates
 from django.db.models.base import ModelBase
 from django.db.utils import IntegrityError
 
@@ -205,6 +205,7 @@ class ModelBuilder(object):
         ('order_by', 'order_by', list),
         ('offset', 'offset', int),
         ('limit', 'limit', int),
+        ('aggregate', 'aggregate', dict),
         ('fields', 'values', list),
         ('count', 'count', bool),
     ]
@@ -227,7 +228,7 @@ def _apply_opts(self):
             # offset and limit operation will return None
             func = getattr(self.qset, qset_opt, None)
             value = getattr(self.params, opt, None)
-            # do not execute operation if value is not passed
+            # do not execute operation if value is not passed,
             # or it does not have default value specified in
             # Parameters class such as [], {}, False
             if value is None:
@@ -237,7 +238,20 @@ def _apply_opts(self):
                                             ' expected %s'
                                             % (type(value), opt, opt_type))
             if isinstance(value, dict):
-                self.qset = func(**value)
+                if qset_opt == 'aggregate':
+                    aggr_arg = []
+                    for aggr_opt, aggr_field in value.items():
+                        aggr_func = getattr(aggregates, aggr_opt, None)
+                        if aggr_func is None:
+                            raise InvalidRequest('Invalid aggregate function %s' %
+                                                 aggr_opt)
+                        aggr_func_f = aggr_func(aggr_field)
+                        aggr_arg.append(aggr_func_f)
+                        self.qset = func(*aggr_arg)
+                    # stop all operations after aggregate
+                    break
+                else:
+                    self.qset = func(**value)
             elif qset_opt == 'offset':
                 self.qset = self.qset[self.params.offset:]
             elif qset_opt == 'limit':

diff --git a/tests/server/machine/tests/test_api_reader.py b/tests/server/machine/tests/test_api_reader.py
@@ -362,6 +362,53 @@ def test_invalid_query_type(self):
         self.assertEqual(resp.status_code, 400)
         self.assertFalse(resp.json()['success'])
 
+    def test_aggregate_query(self):
+        self.params = {
+            'app_name': 'machine',
+            'model_name': 'Machine',
+            'aggregate': {
+                'Max': 'cpu_count',
+                'Avg': 'memory'
+            }
+        }
+        resp = self.client.get(self.url, {'payload': json.dumps(self.params)})
+        self.assertEqual(resp.status_code, 200)
+        resp_json = resp.json()
+        self.assertEqual(14, resp_json['data']['cpu_count__max'])
+        self.assertEqual(3383.5, resp_json['data']['memory__avg'])
+
+    def test_aggregate_with_count_query(self):
+        self.params = {
+            'app_name': 'machine',
+            'model_name': 'Machine',
+            'aggregate': {
+                'Max': 'cpu_count',
+                'Avg': 'memory',
+                'Count': 'os'
+            },
+            'count': True
+        }
+        resp = self.client.get(self.url, {'payload': json.dumps(self.params)})
+        self.assertEqual(resp.status_code, 200)
+        resp_json = resp.json()
+        self.assertEqual(14, resp_json['data']['cpu_count__max'])
+        self.assertEqual(3383.5, resp_json['data']['memory__avg'])
+        self.assertEqual(100, resp_json['data']['os__count'])
+
+    def test_invalid_aggregate_query(self):
+        self.params = {
+            'app_name': 'machine',
+            'model_name': 'Machine',
+            'aggregate': {
+                'Mix': 'cpu_count',
+            }
+        }
+        resp = self.client.get(self.url, {'payload': json.dumps(self.params)})
+        self.assertEqual(resp.status_code, 400)
+        resp_json = resp.json()
+        self.assertFalse(resp_json['success'])
+        self.assertEqual('Invalid aggregate function Mix', resp_json['message'])
+
     @override_settings(BRIDGEQL_AUTHENTICATION_DECORATOR='server.auth.localtest')
     def test_custom_auth_decorator(self):
         self.params = {