feat(query): add support for Metric in aggregations

Signed-off-by: Charles Lariviere <charles@dribbble.com>

Author: Charles Lariviere

src/metabase/mbql/query.py

@@ -1,15 +1,16 @@
 from dataclasses import dataclass, field
-from typing import List
+from typing import List, Union
 
 from metabase.mbql.aggregations import Aggregation
 from metabase.mbql.filter import Filter
 from metabase.mbql.groupby import GroupBy
+from metabase.resources.metric import Metric
 
 
 @dataclass
 class Query:
     table_id: int
-    aggregations: List[Aggregation]
+    aggregations: List[Union[Aggregation, Metric]]
     group_by: List[GroupBy] = field(default_factory=list)
     filters: List[Filter] = field(default_factory=list)
 
@@ -23,7 +24,14 @@ def compile(self):
 
     @property
     def _aggregations(self):
-        return [aggregation.compile() for aggregation in self.aggregations]
+        aggregations = []
+        for aggregation in self.aggregations:
+            if isinstance(aggregation, Metric):
+                aggregations.append(["metric", aggregation.id])
+            else:
+                aggregations.append(aggregation.compile())
+
+        return aggregations
 
     @property
     def _group_by(self):

tests/mbql/test_query.py

@@ -1,5 +1,6 @@
 from unittest import TestCase
 
+from metabase import Metric
 from metabase.mbql.aggregations import Count, Max
 from metabase.mbql.filter import Equal
 from metabase.mbql.groupby import GroupBy
@@ -42,6 +43,12 @@ def test__aggregations(self):
         )
         self.assertEqual([["count"], ["max", ["field", 5, None]]], query._aggregations)
 
+        query = Query(
+            table_id=12,
+            aggregations=[Count(), Metric(id=4)],
+        )
+        self.assertEqual([["count"], ["metric", 4]], query._aggregations)
+
     def test__group_by(self):
         """Ensure Query._group_by returns a list of compiled GroupBy."""
         query = Query(table_id=12, aggregations=[Count()], group_by=[])