feat: add query.py

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

Author: Charles Larivier

src/metabase/mbql/query.py

@@ -1,21 +1,40 @@
 from dataclasses import dataclass, field
-from typing import List, Union
+from typing import List
 
 from metabase.mbql.aggregations import Aggregation
+from metabase.mbql.filter import Filter
 from metabase.mbql.groupby import GroupBy
 
 
 @dataclass
 class Query:
     table_id: int
-    aggregations: List[Union[Aggregation, Metric]]
+    aggregations: List[Aggregation]
     group_by: List[GroupBy] = field(default_factory=list)
     filters: List[Filter] = field(default_factory=list)
 
     def compile(self):
         return {
             "source-table": self.table_id,
-            "aggregation": [agg.compile() for agg in self.aggregations],
-            "breakout": [],
-            "filter": [],
+            "aggregation": self._aggregations,
+            "breakout": self._group_by,
+            "filter": self._filters,
         }
+
+    @property
+    def _aggregations(self):
+        return [aggregation.compile() for aggregation in self.aggregations]
+
+    @property
+    def _group_by(self):
+        return [group.compile() for group in self.group_by]
+
+    @property
+    def _filters(self):
+        if len(self.filters) == 0:
+            return self.filters
+
+        if len(self.filters) == 1:
+            return self.filters[0].compile()
+
+        return ["and"] + [filt.compile() for filt in self.filters]

tests/mbql/test_query.py

@@ -0,0 +1,76 @@
+from unittest import TestCase
+
+from metabase.mbql.aggregations import Count, Max
+from metabase.mbql.filter import Equal
+from metabase.mbql.groupby import GroupBy
+from metabase.mbql.query import Query
+
+
+class QueryTests(TestCase):
+    def test_compile(self):
+        """Ensure Query.compile() returns valid MBQL."""
+        query = Query(
+            table_id=14,
+            aggregations=[Count(), Max(5)],
+            group_by=[GroupBy(14)],
+            filters=[Equal(2, 5), Equal(5, "foo")],
+        )
+
+        self.assertEqual(
+            {
+                "source-table": 14,
+                "aggregation": [["count"], ["max", ["field", 5, None]]],
+                "breakout": [["field", 14, None]],
+                "filter": [
+                    "and",
+                    ["=", ["field", 2, None], 5],
+                    ["=", ["field", 5, None], "foo"],
+                ],
+            },
+            query.compile(),
+        )
+
+    def test__aggregations(self):
+        """Ensure Query._aggregations returns a list of compiled Aggregation."""
+        query = Query(table_id=12, aggregations=[])
+        self.assertEqual([], query._aggregations)
+
+        query = Query(
+            table_id=12,
+            aggregations=[Count(), Max(5)],
+        )
+        self.assertEqual([["count"], ["max", ["field", 5, None]]], 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=[])
+        self.assertEqual([], query._group_by)
+
+        query = Query(
+            table_id=12,
+            aggregations=[Count()],
+            group_by=[GroupBy(5)],
+        )
+        self.assertEqual([["field", 5, None]], query._group_by)
+
+    def test__filters(self):
+        """Ensure Query._filters returns a list of compiled Filter."""
+        query = Query(table_id=12, aggregations=[Count()], filters=[])
+        self.assertEqual([], query._filters)
+
+        query = Query(
+            table_id=12,
+            aggregations=[Count()],
+            filters=[Equal(5, 2)],
+        )
+        self.assertListEqual(["=", ["field", 5, None], 2], query._filters)
+
+        query = Query(
+            table_id=12,
+            aggregations=[Count()],
+            filters=[Equal(5, 2), Equal(6, 1)],
+        )
+        self.assertEqual(
+            ["and", ["=", ["field", 5, None], 2], ["=", ["field", 6, None], 1]],
+            query._filters,
+        )