diff --git a/python_modules/dagster/dagster/_core/definitions/multi_dimensional_partitions.py b/python_modules/dagster/dagster/_core/definitions/multi_dimensional_partitions.py index 68ccf339fa5ce..2cc6b6ab6cdad 100644 --- a/python_modules/dagster/dagster/_core/definitions/multi_dimensional_partitions.py +++ b/python_modules/dagster/dagster/_core/definitions/multi_dimensional_partitions.py @@ -425,13 +425,11 @@ def get_tags_for_partition_key(self, partition_key: str) -> Mapping[str, str]: @property def time_window_dimension(self) -> PartitionDimensionDefinition: - check.invariant(self.has_time_window_dimension, "Must have time window dimension") - check.inst( - self.primary_dimension.partitions_def, - TimeWindowPartitionsDefinition, - "Sanity check that assumption that primary is time-windowed if it exists", + time_window_dims = self._get_time_window_dims() + check.invariant( + len(time_window_dims) == 1, "Expected exactly one time window partitioned dimension" ) - return self.primary_dimension + return next(iter(time_window_dims)) def _get_time_window_dims(self) -> List[PartitionDimensionDefinition]: return [