Add GroupByParam

Author: DevonFulcher

.changes/unreleased/Features-20250327-224851.yaml

@@ -0,0 +1,3 @@
+kind: Features
+body: Add support to pass in typed group by
+time: 2025-03-27T22:48:51.368818-05:00

dbtsl/api/shared/query_params.py

@@ -2,6 +2,14 @@
 from typing import List, Optional, TypedDict, Union
 
 
+@dataclass(frozen=True)
+class GroupByParam:
+    """Parameter for a group_by, i.e a dimension or an entity."""
+
+    name: str
+    grain: Optional[str]
+
+
 @dataclass(frozen=True)
 class OrderByMetric:
     """Spec for ordering by a metric."""
@@ -33,7 +41,7 @@ class QueryParameters(TypedDict, total=False):
 
     saved_query: str
     metrics: List[str]
-    group_by: List[str]
+    group_by: List[Union[GroupByParam, str]]
     limit: int
     order_by: List[Union[OrderBySpec, str]]
     where: List[str]
@@ -45,7 +53,7 @@ class AdhocQueryParametersStrict:
     """The parameters of an adhoc query, strictly validated."""
 
     metrics: Optional[List[str]]
-    group_by: Optional[List[str]]
+    group_by: Optional[List[Union[GroupByParam, str]]]
     limit: Optional[int]
     order_by: Optional[List[OrderBySpec]]
     where: Optional[List[str]]
@@ -64,7 +72,9 @@ class SavedQueryQueryParametersStrict:
 
 
 def validate_order_by(
-    known_metrics: List[str], known_group_bys: List[str], clause: Union[OrderBySpec, str]
+    known_metrics: List[str],
+    known_group_bys: List[Union[str, GroupByParam]],
+    clause: Union[OrderBySpec, str],
 ) -> OrderBySpec:
     """Validate an order by clause like `-metric_name`."""
     if isinstance(clause, OrderByMetric) or isinstance(clause, OrderByGroupBy):
@@ -77,7 +87,11 @@ def validate_order_by(
     if clause in known_metrics:
         return OrderByMetric(name=clause, descending=descending)
 
-    if clause in known_group_bys or clause == "metric_time":
+    normalized_known_group_bys = [
+        known_group_by.name if isinstance(known_group_by, GroupByParam) else known_group_by
+        for known_group_by in known_group_bys
+    ]
+    if clause in normalized_known_group_bys or clause == "metric_time":
         return OrderByGroupBy(name=clause, descending=descending, grain=None)
 
     # TODO: make this error less strict when server supports order_by type inference.

tests/api/adbc/test_protocol.py

@@ -1,5 +1,5 @@
 from dbtsl.api.adbc.protocol import ADBCProtocol
-from dbtsl.api.shared.query_params import OrderByGroupBy, OrderByMetric
+from dbtsl.api.shared.query_params import GroupByParam, OrderByGroupBy, OrderByMetric
 
 
 def test_serialize_val_basic_values() -> None:
@@ -74,6 +74,23 @@ def test_get_query_sql_simple_query() -> None:
     assert sql == expected
 
 
+def test_get_query_sql_group_by_param() -> None:
+    sql = ADBCProtocol.get_query_sql(
+        params={
+            "metrics": ["a", "b"],
+            "group_by": [
+                GroupByParam(name="c", grain="day"),
+                GroupByParam(name="d", grain="week"),
+            ],
+        }
+    )
+    expected = (
+        'SELECT * FROM {{ semantic_layer.query(metrics=["a","b"],'
+        'group_by=[Dimension("c").grain("day"),Dimension("d").grain("week")],read_cache=True) }}'
+    )
+    assert sql == expected
+
+
 def test_get_query_sql_dimension_values_query() -> None:
     sql = ADBCProtocol.get_dimension_values_sql(params={"metrics": ["a", "b"]})
     expected = 'SELECT * FROM {{ semantic_layer.dimension_values(metrics=["a","b"]) }}'

tests/query_test_cases.py

@@ -1,7 +1,7 @@
 from typing import List
 
 from dbtsl import OrderByGroupBy
-from dbtsl.api.shared.query_params import QueryParameters
+from dbtsl.api.shared.query_params import GroupByParam, QueryParameters
 
 TEST_QUERIES: List[QueryParameters] = [
     # ad hoc query, all parameters
@@ -38,4 +38,14 @@
     {
         "saved_query": "order_metrics",
     },
+    # group by param object
+    {"metrics": ["order_total"], "group_by": [GroupByParam(name="order_date", grain="month")]},
+    # multiple group by param objects
+    {
+        "metrics": ["order_total"],
+        "group_by": [
+            GroupByParam(name="order_date", grain="month"),
+            GroupByParam(name="fulfillment_date", grain="week"),
+        ],
+    },
 ]