feat: support for custom time granularities

dbtsl/api/adbc/protocol.py

@@ -9,6 +9,7 @@
     QueryParameters,
     validate_query_parameters,
 )
+from dbtsl.models.time import TimeGranularity
 
 
 class ADBCProtocol:
@@ -32,7 +33,7 @@ def _serialize_val(cls, val: Any) -> str:
         if isinstance(val, OrderByGroupBy):
             d = f'Dimension("{val.name}")'
             if val.grain:
-                grain_str = val.grain.name.lower()
+                grain_str = val.grain.name.lower() if isinstance(val.grain, TimeGranularity) else val.grain.lower()
                 d += f'.grain("{grain_str}")'
             if val.descending:
                 d += ".descending(True)"

dbtsl/api/shared/query_params.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import List, Optional, TypedDict, Union
 
-from dbtsl.models.time import TimeGranularity
+from dbtsl.models.time import Grain
 
 
 @dataclass(frozen=True)
@@ -20,7 +20,7 @@ class OrderByGroupBy:
     """
 
     name: str
-    grain: Optional[TimeGranularity]
+    grain: Optional[Grain]
     descending: bool = False
 
 

dbtsl/models/__init__.py

@@ -20,7 +20,7 @@
     SavedQueryQueryParams,
     SavedQueryWhereParam,
 )
-from .time import DatePart, TimeGranularity
+from .time import DatePart, Grain, TimeGranularity
 
 # Only importing this so it registers aliases
 _ = QueryResult
@@ -37,6 +37,7 @@
     "Export",
     "ExportConfig",
     "ExportDestinationType",
+    "Grain",
     "Measure",
     "Metric",
     "MetricType",

dbtsl/models/dimension.py

@@ -3,7 +3,7 @@
 from typing import List, Optional
 
 from dbtsl.models.base import BaseModel, GraphQLFragmentMixin
-from dbtsl.models.time import TimeGranularity
+from dbtsl.models.time import Grain
 
 
 class DimensionType(str, Enum):
@@ -24,4 +24,4 @@ class Dimension(BaseModel, GraphQLFragmentMixin):
     label: Optional[str]
     is_partition: bool
     expr: Optional[str]
-    queryable_granularities: List[TimeGranularity]
+    queryable_granularities: List[Grain]

dbtsl/models/metric.py

@@ -6,7 +6,7 @@
 from dbtsl.models.dimension import Dimension
 from dbtsl.models.entity import Entity
 from dbtsl.models.measure import Measure
-from dbtsl.models.time import TimeGranularity
+from dbtsl.models.time import Grain
 
 
 class MetricType(str, Enum):
@@ -29,6 +29,6 @@ class Metric(BaseModel, GraphQLFragmentMixin):
     dimensions: List[Dimension]
     measures: List[Measure]
     entities: List[Entity]
-    queryable_granularities: List[TimeGranularity]
+    queryable_granularities: List[Grain]
     label: str
     requires_metric_time: bool

dbtsl/models/saved_query.py

@@ -3,7 +3,7 @@
 from typing import List, Optional
 
 from dbtsl.models.base import BaseModel, GraphQLFragmentMixin
-from dbtsl.models.time import DatePart, TimeGranularity
+from dbtsl.models.time import DatePart, Grain
 
 
 class ExportDestinationType(str, Enum):
@@ -42,7 +42,7 @@ class SavedQueryGroupByParam(BaseModel, GraphQLFragmentMixin):
     """The groupBy param of a saved query."""
 
     name: str
-    grain: Optional[TimeGranularity]
+    grain: Optional[Grain]
     date_part: Optional[DatePart]
 
 

dbtsl/models/time.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from typing import Union
 
 
 class TimeGranularity(str, Enum):
@@ -17,6 +18,10 @@ class TimeGranularity(str, Enum):
     YEAR = "YEAR"
 
 
+Grain = Union[TimeGranularity, str]
+"""Either a standard TimeGranularity or a custom grain."""
+
+
 class DatePart(str, Enum):
     """Date part."""
 

tests/api/adbc/test_protocol.py

@@ -30,6 +30,10 @@ def test_serialize_val_OrderByGroupBy() -> None:
         ADBCProtocol._serialize_val(OrderByGroupBy(name="m", grain=TimeGranularity.WEEK, descending=True))
         == 'Dimension("m").grain("week").descending(True)'
     )
+    assert (
+        ADBCProtocol._serialize_val(OrderByGroupBy(name="m", grain="custom_grain"))
+        == 'Dimension("m").grain("custom_grain")'
+    )
 
 
 def test_serialize_query_params_metrics() -> None: