[pac] Add partitions_subset to AssetCheckEvaluationPlanned, partition to AssetCheckEvaluation

Author: OwenKephart

python_modules/dagster/dagster/_core/asset_graph_view/serializable_entity_subset.py

@@ -50,6 +50,12 @@ class SerializableEntitySubset(Generic[T_EntityKey]):
     key: T_EntityKey
     value: EntitySubsetValue
 
+    @classmethod
+    def empty(
+        cls, key: T_EntityKey, partitions_def: Optional[PartitionsDefinition]
+    ) -> "SerializableEntitySubset[T_EntityKey]":
+        return cls(key=key, value=partitions_def.empty_subset() if partitions_def else False)
+
     @classmethod
     def from_coercible_value(
         cls,

python_modules/dagster/dagster/_core/definitions/asset_checks/asset_check_evaluation.py

@@ -9,6 +9,7 @@
 )
 from dagster._core.definitions.events import AssetKey, MetadataValue, RawMetadataValue
 from dagster._core.definitions.metadata import normalize_metadata
+from dagster._core.definitions.partitions.subset import PartitionsSubset
 from dagster._serdes import whitelist_for_serdes
 
 
@@ -19,6 +20,7 @@ class AssetCheckEvaluationPlanned:
 
     asset_key: AssetKey
     check_name: str
+    partitions_subset: Optional[PartitionsSubset] = None
 
     @property
     def asset_check_key(self) -> AssetCheckKey:
@@ -59,6 +61,8 @@ class AssetCheckEvaluation(IHaveNew):
             A text description of the result of the check evaluation.
         blocking (Optional[bool]):
             Whether the check is blocking.
+        partition (Optional[str]):
+            The partition that the check was evaluated on, if applicable.
     """
 
     asset_key: AssetKey
@@ -69,6 +73,7 @@ class AssetCheckEvaluation(IHaveNew):
     severity: AssetCheckSeverity
     description: Optional[str]
     blocking: Optional[bool]
+    partition: Optional[str]
 
     def __new__(
         cls,
@@ -80,6 +85,7 @@ def __new__(
         severity: AssetCheckSeverity = AssetCheckSeverity.ERROR,
         description: Optional[str] = None,
         blocking: Optional[bool] = None,
+        partition: Optional[str] = None,
     ):
         return super().__new__(
             cls,
@@ -91,6 +97,7 @@ def __new__(
             severity=severity,
             description=description,
             blocking=blocking,
+            partition=partition,
         )
 
     @property

python_modules/dagster/dagster/_core/definitions/declarative_automation/legacy/valid_asset_subset.py

@@ -98,14 +98,14 @@ def all(
                     key=asset_key, value=AllPartitionsSubset(partitions_def, ctx)
                 )
 
-    @staticmethod
+    @classmethod
     def empty(
-        asset_key: AssetKey, partitions_def: Optional[PartitionsDefinition]
+        cls, key: AssetKey, partitions_def: Optional[PartitionsDefinition]
     ) -> "ValidAssetSubset":
         if partitions_def is None:
-            return ValidAssetSubset(key=asset_key, value=False)
+            return cls(key=key, value=False)
         else:
-            return ValidAssetSubset(key=asset_key, value=partitions_def.empty_subset())
+            return cls(key=key, value=partitions_def.empty_subset())
 
     @staticmethod
     def from_asset_partitions_set(

python_modules/dagster/dagster/_core/instance/runs/run_domain.py

@@ -830,6 +830,15 @@ def _log_asset_planned_events(
                         target_asset_key = asset_check_key.asset_key
                         check_name = asset_check_key.name
 
+                        partitions_def = (
+                            asset_graph.get(asset_check_key).partitions_def
+                            if asset_graph.get(asset_check_key)
+                            else None
+                        )
+                        partitions_subset = dagster_run.get_resolved_partitions_subset(
+                            partitions_def
+                        )
+
                         event = DagsterEvent(
                             event_type_value=DagsterEventType.ASSET_CHECK_EVALUATION_PLANNED.value,
                             job_name=job_name,
@@ -838,8 +847,9 @@ def _log_asset_planned_events(
                                 f" asset {target_asset_key.to_string()}"
                             ),
                             event_specific_data=AssetCheckEvaluationPlanned(
-                                target_asset_key,
+                                asset_key=target_asset_key,
                                 check_name=check_name,
+                                partitions_subset=partitions_subset,
                             ),
                             step_key=step.key,
                         )
@@ -883,16 +893,13 @@ def get_materialization_planned_events_for_asset(
     ) -> Sequence["DagsterEvent"]:
         """Moved from DagsterInstance._log_materialization_planned_event_for_asset."""
         from dagster._core.definitions.partitions.context import partition_loading_context
-        from dagster._core.definitions.partitions.definition import DynamicPartitionsDefinition
         from dagster._core.events import AssetMaterializationPlannedData, DagsterEvent
 
         events = []
 
         partition_tag = dagster_run.tags.get(PARTITION_NAME_TAG)
-        partition_range_start, partition_range_end = (
-            dagster_run.tags.get(ASSET_PARTITION_RANGE_START_TAG),
-            dagster_run.tags.get(ASSET_PARTITION_RANGE_END_TAG),
-        )
+        partition_range_start = dagster_run.tags.get(ASSET_PARTITION_RANGE_START_TAG)
+        partition_range_end = dagster_run.tags.get(ASSET_PARTITION_RANGE_END_TAG)
 
         if partition_tag and (partition_range_start or partition_range_end):
             raise DagsterInvariantViolationError(
@@ -901,79 +908,46 @@ def get_materialization_planned_events_for_asset(
                 f" {PARTITION_NAME_TAG}"
             )
 
-        partitions_subset = None
-        individual_partitions = None
-        if partition_range_start or partition_range_end:
-            if not partition_range_start or not partition_range_end:
-                raise DagsterInvariantViolationError(
-                    f"Cannot have {ASSET_PARTITION_RANGE_START_TAG} or"
-                    f" {ASSET_PARTITION_RANGE_END_TAG} set without the other"
-                )
+        asset_node = self._get_repo_scoped_asset_node(
+            asset_graph, asset_key, dagster_run.remote_job_origin
+        )
+        partitions_def = asset_node.partitions_def if asset_node else None
 
-            asset_node = check.not_none(
-                self._get_repo_scoped_asset_node(
-                    asset_graph, asset_key, dagster_run.remote_job_origin
-                )
-            )
+        partitions_subset = dagster_run.get_resolved_partitions_subset(partitions_def)
 
-            partitions_def = asset_node.partitions_def
-            if (
-                isinstance(partitions_def, DynamicPartitionsDefinition)
-                and partitions_def.name is None
-            ):
-                raise DagsterInvariantViolationError(
-                    "Creating a run targeting a partition range is not supported for assets partitioned with function-based dynamic partitions"
+        with partition_loading_context(dynamic_partitions_store=self._instance):
+            if partitions_subset is None:
+                materialization_planned = DagsterEvent.build_asset_materialization_planned_event(
+                    job_name,
+                    step.key,
+                    AssetMaterializationPlannedData(
+                        asset_key, partition=None, partitions_subset=None
+                    ),
                 )
-
-            if partitions_def is not None:
-                with partition_loading_context(dynamic_partitions_store=self._instance):
-                    if self._instance.event_log_storage.supports_partition_subset_in_asset_materialization_planned_events:
-                        partitions_subset = partitions_def.subset_with_partition_keys(
-                            partitions_def.get_partition_keys_in_range(
-                                PartitionKeyRange(partition_range_start, partition_range_end),
-                            )
-                        ).to_serializable_subset()
-                        individual_partitions = []
-                    else:
-                        individual_partitions = partitions_def.get_partition_keys_in_range(
-                            PartitionKeyRange(partition_range_start, partition_range_end),
-                        )
-        elif check.not_none(output.properties).is_asset_partitioned and partition_tag:
-            individual_partitions = [partition_tag]
-
-        assert not (individual_partitions and partitions_subset), (
-            "Should set either individual_partitions or partitions_subset, but not both"
-        )
-
-        if not individual_partitions and not partitions_subset:
-            materialization_planned = DagsterEvent.build_asset_materialization_planned_event(
-                job_name,
-                step.key,
-                AssetMaterializationPlannedData(asset_key, partition=None, partitions_subset=None),
-            )
-            events.append(materialization_planned)
-        elif individual_partitions:
-            for individual_partition in individual_partitions:
+                events.append(materialization_planned)
+            elif self._instance.event_log_storage.supports_partition_subset_in_asset_materialization_planned_events:
                 materialization_planned = DagsterEvent.build_asset_materialization_planned_event(
                     job_name,
                     step.key,
                     AssetMaterializationPlannedData(
-                        asset_key,
-                        partition=individual_partition,
-                        partitions_subset=partitions_subset,
+                        asset_key, partition=None, partitions_subset=partitions_subset
                     ),
                 )
                 events.append(materialization_planned)
-
-        else:
-            materialization_planned = DagsterEvent.build_asset_materialization_planned_event(
-                job_name,
-                step.key,
-                AssetMaterializationPlannedData(
-                    asset_key, partition=None, partitions_subset=partitions_subset
-                ),
-            )
-            events.append(materialization_planned)
+            else:
+                for partition_key in partitions_subset.get_partition_keys():
+                    materialization_planned = (
+                        DagsterEvent.build_asset_materialization_planned_event(
+                            job_name,
+                            step.key,
+                            AssetMaterializationPlannedData(
+                                asset_key,
+                                partition=partition_key,
+                                partitions_subset=None,
+                            ),
+                        )
+                    )
+                    events.append(materialization_planned)
 
         return events
 

python_modules/dagster/dagster/_core/storage/asset_check_execution_record.py

@@ -60,6 +60,7 @@ class AssetCheckExecutionRecord(
             # Old records won't have an event if the status is PLANNED.
             ("event", Optional[EventLogEntry]),
             ("create_timestamp", float),
+            ("partition", Optional[str]),
         ],
     ),
     LoadableBy[AssetCheckKey],
@@ -72,13 +73,15 @@ def __new__(
         status: AssetCheckExecutionRecordStatus,
         event: Optional[EventLogEntry],
         create_timestamp: float,
+        partition: Optional[str],
     ):
         check.inst_param(key, "key", AssetCheckKey)
         check.int_param(id, "id")
         check.str_param(run_id, "run_id")
         check.inst_param(status, "status", AssetCheckExecutionRecordStatus)
         check.opt_inst_param(event, "event", EventLogEntry)
         check.float_param(create_timestamp, "create_timestamp")
+        check.opt_str_param(partition, "partition")
 
         event_type = event.dagster_event_type if event else None
         if status == AssetCheckExecutionRecordStatus.PLANNED:
@@ -105,6 +108,7 @@ def __new__(
             status=status,
             event=event,
             create_timestamp=create_timestamp,
+            partition=partition,
         )
 
     @property
@@ -129,6 +133,7 @@ def from_db_row(cls, row, key: AssetCheckKey) -> "AssetCheckExecutionRecord":
                 else None
             ),
             create_timestamp=utc_datetime_from_naive(row["create_timestamp"]).timestamp(),
+            partition=row["partition"],
         )
 
     @classmethod

python_modules/dagster/dagster/_core/storage/dagster_run.py

@@ -41,6 +41,9 @@
 
 if TYPE_CHECKING:
     from dagster._core.definitions.assets.graph.base_asset_graph import EntityKey
+    from dagster._core.definitions.partitions.definition.partitions_definition import (
+        PartitionsDefinition,
+    )
     from dagster._core.definitions.schedule_definition import ScheduleDefinition
     from dagster._core.definitions.sensor_definition import SensorDefinition
     from dagster._core.remote_representation.external import RemoteSchedule, RemoteSensor
@@ -372,6 +375,57 @@ def get_root_run_id(self) -> Optional[str]:
     def get_parent_run_id(self) -> Optional[str]:
         return self.tags.get(PARENT_RUN_ID_TAG)
 
+    def get_resolved_partitions_subset(
+        self, partitions_def: Optional["PartitionsDefinition"]
+    ) -> Optional[PartitionsSubset]:
+        """Get the partitions subset targeted by a run based on its partition tags."""
+        from dagster._core.definitions.partitions.definition import DynamicPartitionsDefinition
+        from dagster._core.definitions.partitions.partition_key_range import PartitionKeyRange
+        from dagster._core.errors import DagsterInvariantViolationError
+        from dagster._core.storage.tags import (
+            ASSET_PARTITION_RANGE_END_TAG,
+            ASSET_PARTITION_RANGE_START_TAG,
+            PARTITION_NAME_TAG,
+        )
+
+        # some runs store this information directly
+        if self.partitions_subset is not None:
+            return self.partitions_subset
+
+        # otherwise, fetch information from the tags
+        partition_tag = self.tags.get(PARTITION_NAME_TAG)
+        partition_range_start = self.tags.get(ASSET_PARTITION_RANGE_START_TAG)
+        partition_range_end = self.tags.get(ASSET_PARTITION_RANGE_END_TAG)
+
+        if partition_range_start or partition_range_end:
+            if not partition_range_start or not partition_range_end:
+                raise DagsterInvariantViolationError(
+                    f"Cannot have {ASSET_PARTITION_RANGE_START_TAG} or"
+                    f" {ASSET_PARTITION_RANGE_END_TAG} set without the other"
+                )
+
+            if (
+                isinstance(partitions_def, DynamicPartitionsDefinition)
+                and partitions_def.name is None
+            ):
+                raise DagsterInvariantViolationError(
+                    "Creating a run targeting a partition range is not supported for assets "
+                    "partitioned with function-based dynamic partitions"
+                )
+
+            if partitions_def is not None:
+                return partitions_def.subset_with_partition_keys(
+                    partitions_def.get_partition_keys_in_range(
+                        PartitionKeyRange(partition_range_start, partition_range_end),
+                    )
+                ).to_serializable_subset()
+        elif partition_tag and partitions_def is not None:
+            return partitions_def.subset_with_partition_keys(
+                [partition_tag]
+            ).to_serializable_subset()
+
+        return None
+
     @cached_property
     def dagster_execution_info(self) -> Mapping[str, str]:
         """Key-value pairs encoding metadata about the current Dagster run, typically attached to external execution resources.

python_modules/dagster/dagster/_core/storage/event_log/base.py

@@ -1,6 +1,7 @@
 import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Mapping, Sequence, Set
+from types import EllipsisType
 from typing import TYPE_CHECKING, Annotated, NamedTuple, Optional, Union
 
 from dagster_shared.record import ImportFrom, record
@@ -633,13 +634,16 @@ def get_asset_check_execution_history(
         limit: int,
         cursor: Optional[int] = None,
         status: Optional[Set[AssetCheckExecutionRecordStatus]] = None,
+        partition: Optional[str] = None,
     ) -> Sequence[AssetCheckExecutionRecord]:
         """Get executions for one asset check, sorted by recency."""
         pass
 
     @abstractmethod
     def get_latest_asset_check_execution_by_key(
-        self, check_keys: Sequence[AssetCheckKey]
+        self,
+        check_keys: Sequence[AssetCheckKey],
+        partition: Union[str, None, EllipsisType] = ...,
     ) -> Mapping[AssetCheckKey, AssetCheckExecutionRecord]:
         """Get the latest executions for a list of asset checks."""
         pass