Add with_entity_prefix() helpers to LinkableInstances

dbt-labs · Nov 2, 2024 · fc0415f · fc0415f
1 parent 6691096
commit fc0415f
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 0 deletions.
diff --git a/metricflow-semantics/metricflow_semantics/specs/dimension_spec.py b/metricflow-semantics/metricflow_semantics/specs/dimension_spec.py
@@ -43,3 +43,6 @@ def element_path_key(self) -> ElementPathKey:
         return ElementPathKey(
             element_name=self.element_name, element_type=LinkableElementType.DIMENSION, entity_links=self.entity_links
         )
+
+    def with_entity_prefix(self, entity_prefix: EntityReference) -> DimensionSpec:  # noqa: D102
+        return DimensionSpec(element_name=self.element_name, entity_links=(entity_prefix,) + self.entity_links)
diff --git a/metricflow-semantics/metricflow_semantics/specs/entity_spec.py b/metricflow-semantics/metricflow_semantics/specs/entity_spec.py
@@ -53,6 +53,9 @@ def element_path_key(self) -> ElementPathKey:
             element_name=self.element_name, element_type=LinkableElementType.ENTITY, entity_links=self.entity_links
         )
 
+    def with_entity_prefix(self, entity_prefix: EntityReference) -> EntitySpec:  # noqa: D102
+        return EntitySpec(element_name=self.element_name, entity_links=(entity_prefix,) + self.entity_links)
+
 
 @dataclass(frozen=True)
 class LinklessEntitySpec(EntitySpec, SerializableDataclass):

diff --git a/metricflow-semantics/metricflow_semantics/specs/group_by_metric_spec.py b/metricflow-semantics/metricflow_semantics/specs/group_by_metric_spec.py
@@ -111,3 +111,10 @@ def element_path_key(self) -> ElementPathKey:
         return ElementPathKey(
             element_name=self.element_name, element_type=LinkableElementType.METRIC, entity_links=self.entity_links
         )
+
+    def with_entity_prefix(self, entity_prefix: EntityReference) -> GroupByMetricSpec:  # noqa: D102
+        return GroupByMetricSpec(
+            element_name=self.element_name,
+            entity_links=(entity_prefix,) + self.entity_links,
+            metric_subquery_entity_links=self.metric_subquery_entity_links,
+        )
diff --git a/metricflow-semantics/metricflow_semantics/specs/instance_spec.py b/metricflow-semantics/metricflow_semantics/specs/instance_spec.py
@@ -136,5 +136,10 @@ def element_path_key(self) -> ElementPathKey:
         """Return the ElementPathKey representation of the LinkableInstanceSpec subtype."""
         raise NotImplementedError()
 
+    @abstractmethod
+    def with_entity_prefix(self, entity_prefix: EntityReference) -> LinkableInstanceSpec:
+        """Add the selected entity prefix to the start of the entity links."""
+        raise NotImplementedError()
+
 
 SelfTypeT = TypeVar("SelfTypeT", bound="LinkableInstanceSpec")
diff --git a/metricflow-semantics/metricflow_semantics/specs/time_dimension_spec.py b/metricflow-semantics/metricflow_semantics/specs/time_dimension_spec.py
@@ -235,3 +235,12 @@ def generate_possible_specs_for_time_dimension(
     @property
     def is_metric_time(self) -> bool:  # noqa: D102
         return self.element_name == METRIC_TIME_ELEMENT_NAME
+
+    def with_entity_prefix(self, entity_prefix: EntityReference) -> TimeDimensionSpec:  # noqa: D102
+        return TimeDimensionSpec(
+            element_name=self.element_name,
+            entity_links=(entity_prefix,) + self.entity_links,
+            time_granularity=self.time_granularity,
+            date_part=self.date_part,
+            aggregation_state=self.aggregation_state,
+        )