Merge branch 'qs-with-totals' of https://github.com/carrotquest/infi.clickhouse_orm into carrotquest-qs-with-totals

docs/querysets.md

@@ -210,6 +210,19 @@ This queryset is translated to:
 
 After calling `aggregate` you can still use most of the regular queryset methods, such as `count`, `order_by` and `paginate`. It is not possible, however, to call `only` or `aggregate`. It is also not possible to filter the queryset on calculated fields, only on fields that exist in the model.
 
+If you limit aggregation results, it might be useful to get total aggregation values for all rows.
+To achieve this, you can use `with_totals` method. It will return extra row (last) with
+values aggregated for all rows suitable for filters.
+
+    qs = Person.objects_in(database).aggregate('first_name' num='count()').with_totals().order_by('-count')[:3]
+    >>> print qs.count()
+    4
+    >>> for row in qs: 
+    >>>     print(row.first_name, row.count)
+    'Cassandra' 2
+    'Alexandra' 2
+    '' 100
+
 ---
 
 [<< Models and Databases](models_and_databases.md) | [Table of Contents](toc.md) | [Field Types >>](field_types.md)

src/infi/clickhouse_orm/query.py

@@ -290,6 +290,7 @@ class QuerySet(object):
         self._where_q = Q()
         self._prewhere_q = Q()
         self._grouping_fields = []
+        self._grouping_with_totals = False
         self._fields = model_cls.fields().keys()
         self._limits = None
         self._distinct = False
@@ -353,6 +354,9 @@ class QuerySet(object):
         if self._grouping_fields:
             sql += '\nGROUP BY %s' % comma_join('`%s`' % field for field in self._grouping_fields)
 
+            if self._grouping_with_totals:
+                sql += ' WITH TOTALS'
+
         if self._order_by:
             sql += '\nORDER BY ' + self.order_by_as_sql()
 
@@ -572,6 +576,9 @@ class AggregateQuerySet(QuerySet):
     def select_fields_as_sql(self):
         return comma_join(list(self._fields) + ['%s AS %s' % (v, k) for k, v in self._calculated_fields.items()])
 
+    def group_by_as_sql(self):
+        return 'GROUP BY'
+
     def __iter__(self):
         return self._database.select(self.as_sql()) # using an ad-hoc model
 
@@ -582,3 +589,13 @@ class AggregateQuerySet(QuerySet):
         sql = u'SELECT count() FROM (%s)' % self.as_sql()
         raw = self._database.raw(sql)
         return int(raw) if raw else 0
+
+    def with_totals(self):
+        """
+        Adds WITH TOTALS modifier ot GROUP BY, making query return extra row
+        with aggregate function calculated across all the rows. More information:
+        https://clickhouse.yandex/docs/en/query_language/select/#with-totals-modifier
+        """
+        qs = copy(self)
+        qs._grouping_with_totals = True
+        return qs

tests/test_querysets.py

@@ -411,6 +411,17 @@ class AggregateTestCase(TestCaseWithData):
         print(qs.as_sql())
         self.assertEqual(qs.count(), 1)
 
+    def test_aggregate_with_totals(self):
+        qs = Person.objects_in(self.database).aggregate('first_name', count='count()').\
+            with_totals().order_by('-count')[:5]
+        print(qs.as_sql())
+        result = list(qs)
+        self.assertEqual(len(result), 6)
+        for row in result[:-1]:
+            self.assertEqual(2, row.count)
+
+        self.assertEqual(100, result[-1].count)
+
     def test_double_underscore_field(self):
         class Mdl(Model):
             the__number = Int32Field()