Methods to look up agg_time_dimensions

metricflow/model/semantics/linkable_spec_resolver.py

@@ -26,7 +26,6 @@
 from metricflow.mf_logging.pretty_print import mf_pformat
 from metricflow.model.semantics.linkable_element_properties import LinkableElementProperties
 from metricflow.model.semantics.semantic_model_join_evaluator import SemanticModelJoinEvaluator
-from metricflow.model.semantics.semantic_model_lookup import SemanticModelLookup
 from metricflow.protocols.semantics import SemanticModelAccessor
 from metricflow.specs.specs import (
     DEFAULT_TIME_GRANULARITY,
@@ -524,9 +523,7 @@ def __init__(
         linkable_element_sets_to_merge: List[LinkableElementSet] = []
 
         for semantic_model in semantic_manifest.semantic_models:
-            linkable_element_sets_to_merge.append(
-                ValidLinkableSpecResolver._get_elements_in_semantic_model(semantic_model)
-            )
+            linkable_element_sets_to_merge.append(self._get_elements_in_semantic_model(semantic_model))
 
         metric_time_elements_for_no_metrics = self._get_metric_time_elements(measure_reference=None)
         self._no_metric_linkable_element_set = LinkableElementSet.merge_by_path_key(
@@ -550,8 +547,7 @@ def _get_semantic_model_for_measure(self, measure_reference: MeasureReference) -
             )
         return semantic_models_where_measure_was_found[0]
 
-    @staticmethod
-    def _get_elements_in_semantic_model(semantic_model: SemanticModel) -> LinkableElementSet:
+    def _get_elements_in_semantic_model(self, semantic_model: SemanticModel) -> LinkableElementSet:
         """Gets the elements in the semantic model, without requiring any joins.
 
         Elements related to metric_time are handled separately in _get_metric_time_elements().
@@ -568,7 +564,7 @@ def _get_elements_in_semantic_model(semantic_model: SemanticModel) -> LinkableEl
                     properties=frozenset({LinkableElementProperties.LOCAL, LinkableElementProperties.ENTITY}),
                 )
             )
-            for entity_link in SemanticModelLookup.entity_links_for_local_elements(semantic_model):
+            for entity_link in self._semantic_model_lookup.entity_links_for_local_elements(semantic_model):
                 linkable_entities.append(
                     LinkableEntity(
                         semantic_model_origin=semantic_model.reference,
@@ -579,7 +575,7 @@ def _get_elements_in_semantic_model(semantic_model: SemanticModel) -> LinkableEl
                     )
                 )
 
-        for entity_link in SemanticModelLookup.entity_links_for_local_elements(semantic_model):
+        for entity_link in self._semantic_model_lookup.entity_links_for_local_elements(semantic_model):
             dimension_properties = frozenset({LinkableElementProperties.LOCAL})
             for dimension in semantic_model.dimensions:
                 dimension_type = dimension.type

metricflow/model/semantics/metric_lookup.py

@@ -6,15 +6,19 @@
 from dbt_semantic_interfaces.enum_extension import assert_values_exhausted
 from dbt_semantic_interfaces.protocols.metric import Metric, MetricInputMeasure, MetricType
 from dbt_semantic_interfaces.protocols.semantic_manifest import SemanticManifest
-from dbt_semantic_interfaces.references import MeasureReference, MetricReference
+from dbt_semantic_interfaces.references import MeasureReference, MetricReference, TimeDimensionReference
 
 from metricflow.errors.errors import DuplicateMetricError, MetricNotFoundError, NonExistentMeasureError
 from metricflow.model.semantics.linkable_element_properties import LinkableElementProperties
-from metricflow.model.semantics.linkable_spec_resolver import LinkableElementSet, ValidLinkableSpecResolver
+from metricflow.model.semantics.linkable_spec_resolver import (
+    ElementPathKey,
+    LinkableElementSet,
+    ValidLinkableSpecResolver,
+)
 from metricflow.model.semantics.semantic_model_join_evaluator import MAX_JOIN_HOPS
 from metricflow.model.semantics.semantic_model_lookup import SemanticModelLookup
 from metricflow.protocols.semantics import MetricAccessor
-from metricflow.specs.specs import LinkableInstanceSpec
+from metricflow.specs.specs import LinkableInstanceSpec, TimeDimensionSpec
 
 logger = logging.getLogger(__name__)
 
@@ -157,3 +161,32 @@ def contains_cumulative_or_time_offset_metric(self, metric_references: Sequence[
                     if input_metric.offset_window or input_metric.offset_to_grain:
                         return True
         return False
+
+    def _get_agg_time_dimension_path_keys_for_metric(
+        self, metric_reference: MetricReference
+    ) -> Sequence[ElementPathKey]:
+        """Retrieves the aggregate time dimensions associated with the metric's measures."""
+        metric = self.get_metric(metric_reference)
+        path_keys = set()
+        for input_measure in metric.input_measures:
+            path_key = self._semantic_model_lookup.get_agg_time_dimension_path_key_for_measure(
+                measure_reference=input_measure.measure_reference
+            )
+            path_keys.add(path_key)
+        return list(path_keys)
+
+    def get_valid_agg_time_dimensions_for_metric(
+        self, metric_reference: MetricReference
+    ) -> Sequence[TimeDimensionSpec]:
+        """Get the agg time dimension specs that can be used in place of metric time for this metric, if applicable."""
+        agg_time_dimension_element_path_keys = self._get_agg_time_dimension_path_keys_for_metric(metric_reference)
+        if len(agg_time_dimension_element_path_keys) != 1:
+            # If the metric's input measures have different agg_time_dimensions, user must use metric_time.
+            return []
+
+        path_key = agg_time_dimension_element_path_keys[0]
+        valid_agg_time_dimension_specs = TimeDimensionSpec.generate_possible_specs_for_time_dimension(
+            time_dimension_reference=TimeDimensionReference(element_name=path_key.element_name),
+            entity_links=path_key.entity_links,
+        )
+        return valid_agg_time_dimension_specs

metricflow/model/semantics/semantic_model_lookup.py

@@ -25,6 +25,7 @@
 from metricflow.errors.errors import InvalidSemanticModelError
 from metricflow.mf_logging.pretty_print import mf_pformat
 from metricflow.model.semantics.element_group import ElementGrouper
+from metricflow.model.semantics.linkable_spec_resolver import ElementPathKey
 from metricflow.model.spec_converters import MeasureConverter
 from metricflow.protocols.semantics import SemanticModelAccessor
 from metricflow.specs.specs import (
@@ -214,7 +215,7 @@ def _add_semantic_model(self, semantic_model: SemanticModel) -> None:
                     f"Aggregation time dimension does not have a time granularity set: {agg_time_dimension}"
                 )
 
-            primary_entity = SemanticModelLookup._resolved_primary_entity(semantic_model)
+            primary_entity = SemanticModelLookup.resolved_primary_entity(semantic_model)
 
             if primary_entity is None:
                 raise RuntimeError(
@@ -291,7 +292,7 @@ def get_entity_from_semantic_model(
         )
 
     @staticmethod
-    def _resolved_primary_entity(semantic_model: SemanticModel) -> Optional[EntityReference]:
+    def resolved_primary_entity(semantic_model: SemanticModel) -> Optional[EntityReference]:
         """Return the primary entity for dimensions in the model."""
         primary_entity_reference = semantic_model.primary_entity_reference
 
@@ -300,14 +301,12 @@ def _resolved_primary_entity(semantic_model: SemanticModel) -> Optional[EntityRe
         )
 
         # This should be caught by the validation, but adding a sanity check.
-        assert len(entities_with_type_primary) <= 1, f"Found >1 primary entity in {semantic_model}"
+        assert len(entities_with_type_primary) <= 1, f"Found > 1 primary entity in {semantic_model}"
         if primary_entity_reference is not None:
             assert len(entities_with_type_primary) == 0, (
                 f"The primary_entity field was set to {primary_entity_reference}, but there are non-zero entities with "
                 f"type {EntityType.PRIMARY} in {semantic_model}"
             )
-
-        if primary_entity_reference is not None:
             return primary_entity_reference
 
         if len(entities_with_type_primary) > 0:
@@ -339,3 +338,37 @@ def get_element_spec_for_name(self, element_name: str) -> LinkableInstanceSpec:
             return self._entity_ref_to_spec[EntityReference(element_name=element_name)]
         else:
             raise ValueError(f"Unable to find linkable element {element_name} in manifest")
+
+    def get_agg_time_dimension_path_key_for_measure(self, measure_reference: MeasureReference) -> ElementPathKey:
+        """Get the agg time dimension associated with the measure."""
+        agg_time_dimension = self.get_agg_time_dimension_for_measure(measure_reference)
+
+        # A measure's agg_time_dimension is required to be in the same semantic model as the measure,
+        # so we can assume the same semantic model for both measure and dimension.
+        semantic_models = self.get_semantic_models_for_measure(measure_reference)
+        assert (
+            len(semantic_models) == 1
+        ), f"Expected exactly one semantic model for measure {measure_reference}, but found semantic models {semantic_models}."
+        semantic_model = semantic_models[0]
+
+        entity_link = self.resolved_primary_entity(semantic_model)
+        assert (
+            entity_link is not None
+        ), f"Expected semantic model {semantic_model} to have a primary entity since is contains dimensions, but found none."
+
+        return ElementPathKey(
+            element_name=agg_time_dimension.element_name,
+            entity_links=(entity_link,),
+            time_granularity=None,
+            date_part=None,
+        )
+
+    def get_agg_time_dimension_specs_for_measure(
+        self, measure_reference: MeasureReference
+    ) -> Sequence[TimeDimensionSpec]:
+        """Get the agg time dimension specs that can be used in place of metric time for this measure."""
+        path_key = self.get_agg_time_dimension_path_key_for_measure(measure_reference)
+        return TimeDimensionSpec.generate_possible_specs_for_time_dimension(
+            time_dimension_reference=TimeDimensionReference(element_name=path_key.element_name),
+            entity_links=path_key.entity_links,
+        )

metricflow/specs/specs.py

@@ -449,6 +449,37 @@ def comparison_key(self, exclude_fields: Sequence[TimeDimensionSpecField] = ())
             exclude_fields=exclude_fields,
         )
 
+    @staticmethod
+    def generate_possible_specs_for_time_dimension(
+        time_dimension_reference: TimeDimensionReference, entity_links: Tuple[EntityReference, ...]
+    ) -> List[TimeDimensionSpec]:
+        """Generate a list of time dimension specs with all combinations of granularity & date part."""
+        time_dimension_specs: List[TimeDimensionSpec] = []
+        for time_granularity in TimeGranularity:
+            time_dimension_specs.append(
+                TimeDimensionSpec(
+                    element_name=time_dimension_reference.element_name,
+                    entity_links=entity_links,
+                    time_granularity=time_granularity,
+                    date_part=None,
+                )
+            )
+        for date_part in DatePart:
+            for time_granularity in date_part.compatible_granularities:
+                time_dimension_specs.append(
+                    TimeDimensionSpec(
+                        element_name=time_dimension_reference.element_name,
+                        entity_links=entity_links,
+                        time_granularity=time_granularity,
+                        date_part=date_part,
+                    )
+                )
+        return time_dimension_specs
+
+    @property
+    def is_metric_time(self) -> bool:  # noqa: D
+        return self.element_name == METRIC_TIME_ELEMENT_NAME
+
 
 @dataclass(frozen=True)
 class NonAdditiveDimensionSpec(SerializableDataclass):
@@ -635,7 +666,7 @@ def metric_time_specs(self) -> Sequence[TimeDimensionSpec]:
         return tuple(
             time_dimension_spec
             for time_dimension_spec in self.time_dimension_specs
-            if time_dimension_spec.element_name == METRIC_TIME_ELEMENT_NAME
+            if time_dimension_spec.is_metric_time
         )
 
     @property