Add support to pass in typed metrics & group by

Author: DevonFulcher

.changes/unreleased/Features-20250320-155941.yaml

@@ -0,0 +1,3 @@
+kind: Features
+body: Add support to pass in typed metrics and group by for query function
+time: 2025-03-20T15:59:41.744532-05:00

dbtsl/api/graphql/protocol.py

@@ -7,6 +7,7 @@
 from dbtsl.api.graphql.util import render_query
 from dbtsl.api.shared.query_params import (
     AdhocQueryParametersStrict,
+    OrderByGroupBy,
     OrderByMetric,
     QueryParameters,
     validate_query_parameters,
@@ -201,17 +202,20 @@ def get_query_request_variables(environment_id: int, params: QueryParameters) ->
     """Get the GraphQL request variables for a given set of query parameters."""
     strict_params = validate_query_parameters(params)  # type: ignore
 
+    order_by_vars = []
+    for clause in strict_params.order_by or []:
+        if isinstance(clause, OrderByMetric):
+            order_by_vars.append({"metric": {"name": clause.name}, "descending": clause.descending})
+        else:
+            assert isinstance(clause, OrderByGroupBy)
+            order_by_vars.append(
+                {"groupBy": {"name": clause.name, "timeGranularity": clause.grain}, "descending": clause.descending}
+            )
+
     shared_vars = {
         "environmentId": environment_id,
         "where": [{"sql": sql} for sql in strict_params.where] if strict_params.where is not None else None,
-        "orderBy": [
-            {"metric": {"name": clause.name}, "descending": clause.descending}
-            if isinstance(clause, OrderByMetric)
-            else {"groupBy": {"name": clause.name, "timeGranularity": clause.grain}, "descending": clause.descending}
-            for clause in strict_params.order_by
-        ]
-        if strict_params.order_by is not None
-        else None,
+        "orderBy": order_by_vars if strict_params.order_by is not None else None,
         "limit": strict_params.limit,
         "readCache": strict_params.read_cache,
     }

dbtsl/api/shared/query_params.py

@@ -10,6 +10,16 @@ class OrderByMetric:
     descending: bool = False
 
 
+@dataclass(frozen=True)
+class GroupByParam:
+    """Spec for a group_by, i.e a dimension or an entity.
+
+    Not specifying a grain will defer the grain choice to the server.
+    """
+
+    name: str
+    grain: Optional[str]
+
 @dataclass(frozen=True)
 class OrderByGroupBy:
     """Spec for ordering by a group_by, i.e a dimension or an entity.
@@ -21,7 +31,6 @@ class OrderByGroupBy:
     grain: Optional[str]
     descending: bool = False
 
-
 OrderBySpec = Union[OrderByMetric, OrderByGroupBy]
 
 
@@ -33,7 +42,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]
@@ -44,8 +53,8 @@ class QueryParameters(TypedDict, total=False):
 class AdhocQueryParametersStrict:
     """The parameters of an adhoc query, strictly validated."""
 
-    metrics: Optional[List[str]]
-    group_by: Optional[List[str]]
+    metrics: Optional[List[Union[str, OrderByMetric]]]
+    group_by: Optional[List[Union[str, OrderByGroupBy]]]
     limit: Optional[int]
     order_by: Optional[List[OrderBySpec]]
     where: Optional[List[str]]
@@ -64,7 +73,9 @@ class SavedQueryQueryParametersStrict:
 
 
 def validate_order_by(
-    known_metrics: List[str], known_group_bys: List[str], clause: Union[OrderBySpec, str]
+    known_metrics: List[Union[str, OrderByMetric]],
+    known_group_bys: List[Union[str, OrderByGroupBy]],
+    clause: Union[OrderBySpec, str],
 ) -> OrderBySpec:
     """Validate an order by clause like `-metric_name`."""
     if isinstance(clause, OrderByMetric) or isinstance(clause, OrderByGroupBy):
@@ -74,10 +85,23 @@ def validate_order_by(
     if descending or clause.startswith("+"):
         clause = clause[1:]
 
-    if clause in known_metrics:
+    normalized_known_metrics: list[str] = []
+    for known_metric in known_metrics:
+        if isinstance(known_metric, OrderByMetric):
+            normalized_known_metrics.append(known_metric.name)
+        else:
+            normalized_known_metrics.append(known_metric)
+
+    if clause in normalized_known_metrics:
         return OrderByMetric(name=clause, descending=descending)
 
-    if clause in known_group_bys or clause == "metric_time":
+    normalized_known_group_bys: list[str] = []
+    for known_group_by in known_group_bys:
+        if isinstance(known_group_by, OrderByGroupBy):
+            normalized_known_group_bys.append(known_group_by.name)
+        else:
+            normalized_known_group_bys.append(known_group_by)
+    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/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 OrderByMetric, QueryParameters
 
 TEST_QUERIES: List[QueryParameters] = [
     # ad hoc query, all parameters
@@ -38,4 +38,39 @@
     {
         "saved_query": "order_metrics",
     },
+    # typed metric & group by
+    {
+        "metrics": [OrderByMetric(name="order_total", descending=False)],
+        "group_by": [OrderByGroupBy(name="customer__customer_type", grain="day", descending=False)],
+    },
+    # multiple typed metrics
+    {
+        "metrics": [
+            OrderByMetric(name="order_total", descending=True),
+            OrderByMetric(name="order_count", descending=False),
+        ],
+    },
+    # multiple typed group by
+    {
+        "metrics": [OrderByMetric(name="order_total", descending=False)],
+        "group_by": [
+            OrderByGroupBy(name="customer__customer_type", grain="day", descending=True),
+            OrderByGroupBy(name="order__status", grain="month", descending=False),
+        ],
+    },
+    # typed metrics with different grains
+    {
+        "metrics": [OrderByMetric(name="order_total", descending=True)],
+        "group_by": [
+            OrderByGroupBy(name="order__created_at", grain="day", descending=True),
+            OrderByGroupBy(name="order__created_at", grain="month", descending=False),
+        ],
+    },
+    # typed metrics with where clause and limit
+    {
+        "metrics": [OrderByMetric(name="order_total", descending=True)],
+        "group_by": [OrderByGroupBy(name="customer__customer_type", grain="day", descending=True)],
+        "where": ["order_total > 1000"],
+        "limit": 10,
+    },
 ]