/* PR_START p--short-term-perf 30 */ Change resolved_primary_entity.

dbt-labs · Oct 31, 2024 · 6b14dbd · 6b14dbd
1 parent bc06fb6
commit 6b14dbd
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 29 deletions.
diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_helper.py b/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_helper.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional, Sequence
+from typing import Sequence
 
 from dbt_semantic_interfaces.protocols import Dimension
 from dbt_semantic_interfaces.protocols.entity import Entity
@@ -31,8 +31,12 @@ def get_entity_from_semantic_model(
         )
 
     @staticmethod
-    def resolved_primary_entity(semantic_model: SemanticModel) -> Optional[EntityReference]:
-        """Return the primary entity for dimensions in the model."""
+    def resolved_primary_entity(semantic_model: SemanticModel) -> EntityReference:
+        """Return the primary entity for dimensions in the model.
+
+        Semantic models with measures or dimensions should have a primary entity as enforced by semantic manifest
+        validation.
+        """
         primary_entity_reference = semantic_model.primary_entity_reference
 
         entities_with_type_primary = tuple(
@@ -51,7 +55,7 @@ def resolved_primary_entity(semantic_model: SemanticModel) -> Optional[EntityRef
         if len(entities_with_type_primary) > 0:
             return entities_with_type_primary[0].reference
 
-        return None
+        raise ValueError(f"No primary entity found in {semantic_model.reference=}")
 
     @staticmethod
     def entity_links_for_local_elements(semantic_model: SemanticModel) -> Sequence[EntityReference]:

diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_lookup.py b/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_lookup.py
@@ -20,7 +20,6 @@
 
 from metricflow_semantics.errors.error_classes import InvalidSemanticModelError
 from metricflow_semantics.mf_logging.lazy_formattable import LazyFormat
-from metricflow_semantics.mf_logging.pretty_print import mf_pformat
 from metricflow_semantics.model.semantics.element_group import ElementGrouper
 from metricflow_semantics.model.semantics.semantic_model_helper import SemanticModelHelper
 from metricflow_semantics.model.spec_converters import MeasureConverter
@@ -198,9 +197,6 @@ def _add_semantic_model(self, semantic_model: SemanticModel) -> None:
                     f"Aggregation time dimension does not have a time granularity set: {agg_time_dimension}"
                 )
 
-            # TODO: do we need this here? This should be handled in validations
-            self.get_primary_entity_else_error(semantic_model)
-
             self._semantic_model_to_aggregation_time_dimensions[semantic_model.reference].add_value(
                 key=TimeDimensionReference(
                     element_name=agg_time_dimension.name,
@@ -246,22 +242,6 @@ def _add_semantic_model(self, semantic_model: SemanticModel) -> None:
 
         self._semantic_model_reference_to_semantic_model[semantic_model.reference] = semantic_model
 
-    def get_primary_entity_else_error(self, semantic_model: SemanticModel) -> EntityReference:
-        """Get primary entity from semantic model and error if it doesn't exist.
-
-        If there are dimensions in the semantic model, there must be a primary entity. If there are measures, we can
-        also assume there must be a primary entity because measures are required to have an `agg_time_dimension`
-        defined in the same semantic model.
-        """
-        # TODO: Move me.
-        primary_entity = SemanticModelHelper.resolved_primary_entity(semantic_model)
-        if primary_entity is None:
-            raise RuntimeError(
-                f"The semantic model should have a primary entity since there are dimensions, but it does not. "
-                f"Semantic model is:\n{mf_pformat(semantic_model)}"
-            )
-        return primary_entity
-
     def get_aggregation_time_dimensions_with_measures(
         self, semantic_model_reference: SemanticModelReference
     ) -> ElementGrouper[TimeDimensionReference, MeasureSpec]:
@@ -310,10 +290,6 @@ def _get_agg_time_dimension_specs_for_measure(
         # so we can assume the same semantic model for both measure and dimension.
         semantic_model = self.get_semantic_model_for_measure(measure_reference)
         entity_link = SemanticModelHelper.resolved_primary_entity(semantic_model)
-        assert entity_link is not None, (
-            f"Expected semantic model {semantic_model} to have a primary entity since it has a "
-            "measure requiring an agg_time_dimension, but found none.",
-        )
         return TimeDimensionSpec.generate_possible_specs_for_time_dimension(
             time_dimension_reference=agg_time_dimension,
             entity_links=(entity_link,),

diff --git a/metricflow/engine/metricflow_engine.py b/metricflow/engine/metricflow_engine.py
@@ -674,7 +674,7 @@ def list_dimensions(self) -> List[Dimension]:  # noqa: D102
                         pydantic_dimension=SemanticModelHelper.get_dimension_from_semantic_model(
                             semantic_model=semantic_model, dimension_reference=dimension_reference
                         ),
-                        entity_links=(semantic_model_lookup.get_primary_entity_else_error(semantic_model),),
+                        entity_links=(SemanticModelHelper.resolved_primary_entity(semantic_model),),
                     )
                 )