added support for group_by and Count

Author: dantownsend

docs/src/piccolo/query_clauses/group_by.rst

@@ -0,0 +1,37 @@
+.. _group_by:
+
+group_by
+========
+
+You can use ``group_by`` clauses with the following queries:
+
+* :ref:`Select`
+
+It is used in combination with aggregate functions - ``Count`` is currently
+supported.
+
+Count
+-----
+
+In the following query, we get a count of the number of bands per manager:
+
+.. code-block:: python
+
+    >>> from piccolo.query.methods.select import Count
+
+    >>> b = Band
+    >>> b.select(
+    >>>     b.manager.name,
+    >>>     Count(b.manager)
+    >>> ).group_by(
+    >>>     b.manager
+    >>> ).run_sync()
+
+    [
+        {"manager.name": "Graydon", "count": 1},
+        {"manager.name": "Guido", "count": 1}
+    ]
+
+.. currentmodule:: piccolo.query.methods.select
+
+.. autoclass:: Count

docs/src/piccolo/query_clauses/index.rst

@@ -10,6 +10,7 @@ by modifying the return values.
     :maxdepth: 0
 
     ./first
+    ./group_by
     ./limit
     ./offset
     ./order_by

docs/src/piccolo/query_types/select.rst

@@ -114,6 +114,11 @@ first
 
 See  :ref:`first`.
 
+group_by
+~~~~~~~~
+
+See  :ref:`group_by`.
+
 limit
 ~~~~~
 

piccolo/columns/base.py

@@ -236,7 +236,7 @@ def _validate_default(
         """
         if getattr(self, "validated", None):
             # If it has previously been validated by a subclass, don't
-            # validated again.
+            # validate again.
             return True
         elif (
             default is None

piccolo/query/methods/select.py

@@ -9,6 +9,7 @@
 from piccolo.query.mixins import (
     ColumnsDelegate,
     DistinctDelegate,
+    GroupByDelegate,
     LimitDelegate,
     OffsetDelegate,
     OrderByDelegate,
@@ -22,13 +23,36 @@
     from piccolo.custom_types import Combinable
 
 
+class Count(Selectable):
+    """
+    Used in conjunction with the ``group_by`` clause in ``Select`` queries.
+
+    If a column is specified, the count is for non-null values in that
+    column. If no column is specified, the count is for all rows, whether
+    they have null values or not.
+    """
+
+    def __init__(self, column: t.Optional[Column] = None):
+        self.column = column
+
+    def get_select_string(self, engine_type: str, just_alias=False) -> str:
+        if self.column is None:
+            column_name = "*"
+        else:
+            column_name = self.column._meta.get_full_name(
+                just_alias=just_alias
+            )
+        return f"COUNT({column_name}) AS count"
+
+
 class Select(Query):
 
     __slots__ = (
         "columns_list",
         "exclude_secrets",
         "columns_delegate",
         "distinct_delegate",
+        "group_by_delegate",
         "limit_delegate",
         "offset_delegate",
         "order_by_delegate",
@@ -47,6 +71,7 @@ def __init__(
 
         self.columns_delegate = ColumnsDelegate()
         self.distinct_delegate = DistinctDelegate()
+        self.group_by_delegate = GroupByDelegate()
         self.limit_delegate = LimitDelegate()
         self.offset_delegate = OffsetDelegate()
         self.order_by_delegate = OrderByDelegate()
@@ -64,6 +89,11 @@ def distinct(self) -> Select:
         self.distinct_delegate.distinct()
         return self
 
+    def group_by(self, *columns: Column) -> Select:
+        columns = self.table._process_column_args(*columns)
+        self.group_by_delegate.group_by(*columns)
+        return self
+
     def limit(self, number: int) -> Select:
         self.limit_delegate.limit(number)
         return self
@@ -220,6 +250,10 @@ def querystrings(self) -> t.Sequence[QueryString]:
             query += " WHERE {}"
             args.append(self.where_delegate._where.querystring)
 
+        if self.group_by_delegate._group_by:
+            query += " {}"
+            args.append(self.group_by_delegate._group_by.querystring)
+
         if self.order_by_delegate._order_by:
             query += " {}"
             args.append(self.order_by_delegate._order_by.querystring)

piccolo/query/mixins.py

@@ -233,3 +233,34 @@ class OffsetDelegate:
 
     def offset(self, number: int = 0):
         self._offset = Offset(number)
+
+
+@dataclass
+class GroupBy:
+    __slots__ = ("columns",)
+
+    columns: t.Sequence[Column]
+
+    @property
+    def querystring(self) -> QueryString:
+        columns_names = ", ".join(
+            [i._meta.get_full_name(just_alias=True) for i in self.columns]
+        )
+        return QueryString(f" GROUP BY {columns_names}")
+
+    def __str__(self):
+        return self.querystring.__str__()
+
+
+@dataclass
+class GroupByDelegate:
+    """
+    Used to group results - needed when doing aggregation.
+
+    .group_by(Band.name)
+    """
+
+    _group_by: t.Optional[GroupBy] = None
+
+    def group_by(self, *columns: Column):
+        self._group_by = GroupBy(columns=columns)

tests/table/test_select.py

@@ -1,6 +1,7 @@
 from unittest import TestCase
 
 from piccolo.apps.user.tables import BaseUser
+from piccolo.query.methods.select import Count
 
 from ..base import DBTestCase, postgres_only, sqlite_only
 from ..example_project.tables import Band, Concert
@@ -293,6 +294,91 @@ def test_distinct(self):
 
         self.assertTrue(response == [{"name": "Pythonistas"}])
 
+    def test_count_group_by(self):
+        """
+        Test grouping and counting all rows.
+        """
+        self.insert_rows()
+        self.insert_rows()
+
+        response = (
+            Band.select(Band.name, Count())
+            .group_by(Band.name)
+            .order_by(Band.name)
+            .run_sync()
+        )
+
+        self.assertTrue(
+            response
+            == [
+                {"name": "CSharps", "count": 2},
+                {"name": "Pythonistas", "count": 2},
+                {"name": "Rustaceans", "count": 2},
+            ]
+        )
+
+    def test_count_column_group_by(self):
+        """
+        Test grouping and counting a specific column. Any null values in the
+        specified column will be omitted from the count.
+        """
+        self.insert_rows()
+        self.insert_rows()
+        self.run_sync(
+            """
+            INSERT INTO band (
+                name,
+                manager,
+                popularity
+            ) VALUES (
+                'SomeBand',
+                null,
+                1000
+            );"""
+        )
+
+        response = (
+            Band.select(Band.manager.name, Count(Band.manager))
+            .group_by(Band.manager.name)
+            .order_by(Band.manager.name)
+            .run_sync()
+        )
+
+        # We need to sort them, because SQLite and Postgres treat Null
+        # differently when sorting.
+        response = sorted(response, key=lambda x: x["manager.name"] or "")
+
+        self.assertTrue(
+            response
+            == [
+                {"manager.name": None, "count": 0},
+                {"manager.name": "Graydon", "count": 2},
+                {"manager.name": "Guido", "count": 2},
+                {"manager.name": "Mads", "count": 2},
+            ]
+        )
+
+        # This time the nulls should be counted, as we omit the column argument
+        # from Count:
+        response = (
+            Band.select(Band.manager.name, Count())
+            .group_by(Band.manager.name)
+            .order_by(Band.manager.name)
+            .run_sync()
+        )
+
+        response = sorted(response, key=lambda x: x["manager.name"] or "")
+
+        self.assertTrue(
+            response
+            == [
+                {"manager.name": None, "count": 1},
+                {"manager.name": "Graydon", "count": 2},
+                {"manager.name": "Guido", "count": 2},
+                {"manager.name": "Mads", "count": 2},
+            ]
+        )
+
     def test_columns(self):
         """
         Make sure the colums method can be used to specify which columns to