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

dbt-labs · Oct 30, 2024 · d712d71 · d712d71
1 parent d18e478
commit d712d71
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 11 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,7 +31,7 @@ def get_entity_from_semantic_model(
         )
 
     @staticmethod
-    def resolved_primary_entity(semantic_model: SemanticModel) -> Optional[EntityReference]:
+    def resolved_primary_entity(semantic_model: SemanticModel) -> EntityReference:
         """Return the primary entity for dimensions in the model."""
         primary_entity_reference = semantic_model.primary_entity_reference
 
@@ -51,7 +51,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=}")
 
     @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
@@ -198,9 +198,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,
@@ -310,10 +307,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),),
                     )
                 )