From 50fed79fd9b5d6fb66ebe2fc308073f9306d1da9 Mon Sep 17 00:00:00 2001
From: Courtney Holcomb <courtneyeholcomb@gmail.com>
Date: Fri, 14 Jun 2024 14:37:28 -0700
Subject: [PATCH] Set default_grain in transformations, if not set in YAML

---
 .../implementations/metric.py                 |  2 +-
 dbt_semantic_interfaces/protocols/metric.py   |  2 +-
 .../transformations/default_grain.py          | 48 +++++++++++++++++++
 3 files changed, 50 insertions(+), 2 deletions(-)
 create mode 100644 dbt_semantic_interfaces/transformations/default_grain.py

diff --git a/dbt_semantic_interfaces/implementations/metric.py b/dbt_semantic_interfaces/implementations/metric.py
index fa3ab892..bb0ef2a7 100644
--- a/dbt_semantic_interfaces/implementations/metric.py
+++ b/dbt_semantic_interfaces/implementations/metric.py
@@ -202,7 +202,7 @@ class PydanticMetric(HashableBaseModel, ModelWithMetadataParsing):
     metadata: Optional[PydanticMetadata]
     label: Optional[str] = None
     config: Optional[PydanticMetricConfig]
-    default_grain: TimeGranularity = TimeGranularity.DAY
+    default_grain: Optional[TimeGranularity] = None
 
     @property
     def input_measures(self) -> Sequence[PydanticMetricInputMeasure]:
diff --git a/dbt_semantic_interfaces/protocols/metric.py b/dbt_semantic_interfaces/protocols/metric.py
index 37734c3c..cbb67671 100644
--- a/dbt_semantic_interfaces/protocols/metric.py
+++ b/dbt_semantic_interfaces/protocols/metric.py
@@ -328,7 +328,7 @@ def label(self) -> Optional[str]:
 
     @property
     @abstractmethod
-    def default_grain(self) -> TimeGranularity:
+    def default_grain(self) -> Optional[TimeGranularity]:
         """Default grain used for the metric.
 
         This will be used in a couple of circumstances:
diff --git a/dbt_semantic_interfaces/transformations/default_grain.py b/dbt_semantic_interfaces/transformations/default_grain.py
new file mode 100644
index 00000000..ec15ef8e
--- /dev/null
+++ b/dbt_semantic_interfaces/transformations/default_grain.py
@@ -0,0 +1,48 @@
+from typing import Set
+
+from typing_extensions import override
+
+from dbt_semantic_interfaces.implementations.semantic_manifest import (
+    PydanticSemanticManifest,
+)
+from dbt_semantic_interfaces.protocols import ProtocolHint
+from dbt_semantic_interfaces.references import (
+    DimensionReference,
+    TimeDimensionReference,
+)
+from dbt_semantic_interfaces.transformations.transform_rule import (
+    SemanticManifestTransformRule,
+)
+from dbt_semantic_interfaces.type_enums.time_granularity import TimeGranularity
+
+
+class SetDefaultGrainRule(ProtocolHint[SemanticManifestTransformRule[PydanticSemanticManifest]]):
+    """If default_grain is not set for a metric, set it to DAY if available, else the smallest available grain."""
+
+    @override
+    def _implements_protocol(self) -> SemanticManifestTransformRule[PydanticSemanticManifest]:  # noqa: D
+        return self
+
+    @staticmethod
+    def transform_model(semantic_manifest: PydanticSemanticManifest) -> PydanticSemanticManifest:
+        """For each metric, set default_grain to DAY or the smallest granularity supported by all agg_time_dims."""
+        for metric in semantic_manifest.metrics:
+            if metric.default_grain:
+                continue
+
+            default_grain = TimeGranularity.DAY
+            seen_agg_time_dimensions: Set[TimeDimensionReference] = set()
+            for semantic_model in semantic_manifest.semantic_models:
+                for measure_ref in set(metric.measure_references).intersection(semantic_model.measure_references):
+                    agg_time_dimension_ref = semantic_model.checked_agg_time_dimension_for_measure(measure_ref)
+                    if agg_time_dimension_ref in seen_agg_time_dimensions:
+                        continue
+                    seen_agg_time_dimensions.add(agg_time_dimension_ref)
+                    dimension = semantic_model.get_dimension(DimensionReference(agg_time_dimension_ref.element_name))
+                    if (
+                        dimension.type_params
+                        and dimension.type_params.time_granularity.to_int() > default_grain.to_int()
+                    ):
+                        default_grain = dimension.type_params.time_granularity
+
+        return semantic_manifest