dagster 1.12.12__py3-none-any.whl → 1.12.13__py3-none-any.whl

dagster/_core/storage/asset_check_state.py

@@ -0,0 +1,263 @@
+from collections import defaultdict
+from collections.abc import Iterable, Mapping, Sequence
+from typing import TYPE_CHECKING, Optional, TypeAlias
+
+from dagster_shared.record import record
+from dagster_shared.serdes import whitelist_for_serdes
+
+from dagster._core.asset_graph_view.serializable_entity_subset import SerializableEntitySubset
+from dagster._core.definitions.asset_key import AssetCheckKey
+from dagster._core.definitions.partitions.definition.partitions_definition import (
+    PartitionsDefinition,
+)
+from dagster._core.loader import LoadableBy, LoadingContext
+from dagster._core.storage.asset_check_execution_record import (
+    AssetCheckExecutionRecordStatus,
+    AssetCheckExecutionResolvedStatus,
+    AssetCheckPartitionInfo,
+)
+from dagster._core.storage.dagster_run import FINISHED_STATUSES, DagsterRunStatus, RunsFilter
+
+if TYPE_CHECKING:
+    from dagster._core.instance import DagsterInstance
+
+StatusSubsets: TypeAlias = Mapping[
+    AssetCheckExecutionResolvedStatus, SerializableEntitySubset[AssetCheckKey]
+]
+InProgressRuns: TypeAlias = Mapping[str, SerializableEntitySubset[AssetCheckKey]]
+
+
+@whitelist_for_serdes
+@record
+class AssetCheckState(LoadableBy[tuple[AssetCheckKey, Optional[PartitionsDefinition]]]):
+    latest_storage_id: int
+    subsets: StatusSubsets
+    in_progress_runs: InProgressRuns
+
+    def compatible_with(self, partitions_def: Optional[PartitionsDefinition]) -> bool:
+        subset = next(iter(self.subsets.values()), None)
+        if subset is None:
+            return True
+        return subset.is_compatible_with_partitions_def(partitions_def)
+
+    @classmethod
+    def empty(cls) -> "AssetCheckState":
+        return cls(latest_storage_id=0, subsets={}, in_progress_runs={})
+
+    @classmethod
+    def _blocking_batch_load(
+        cls,
+        keys: Iterable[tuple[AssetCheckKey, Optional[PartitionsDefinition]]],
+        context: LoadingContext,
+    ) -> Iterable[Optional["AssetCheckState"]]:
+        keys = list(keys)
+        mapping = context.instance.event_log_storage.get_asset_check_state(keys)
+        return [mapping.get(check_key) for check_key, _ in keys]
+
+    def with_updates(
+        self,
+        key: AssetCheckKey,
+        partitions_def: Optional[PartitionsDefinition],
+        partition_records: Sequence[AssetCheckPartitionInfo],
+        run_statuses: Mapping[str, DagsterRunStatus],
+    ) -> "AssetCheckState":
+        latest_storage_id = max(
+            (
+                max(r.latest_check_event_storage_id, r.latest_materialization_storage_id or 0)
+                for r in partition_records
+            ),
+            default=self.latest_storage_id,
+        )
+
+        subsets = {
+            status: self.subsets.get(status, SerializableEntitySubset.empty(key, partitions_def))
+            for status in AssetCheckExecutionResolvedStatus
+        }
+        in_progress_runs = dict(self.in_progress_runs)
+
+        # update all subsets based on the new partition records
+        subsets, in_progress_runs = _process_partition_records(
+            key, partitions_def, subsets, in_progress_runs, partition_records
+        )
+        # then check the run statuses and resolve any previously in-progress runs that have completed
+        subsets, in_progress_runs = _process_run_statuses(
+            key, partitions_def, subsets, in_progress_runs, run_statuses
+        )
+        return AssetCheckState(
+            latest_storage_id=latest_storage_id,
+            subsets=subsets,
+            in_progress_runs=in_progress_runs,
+        )
+
+
+def bulk_update_asset_check_state(
+    instance: "DagsterInstance",
+    keys: Sequence[tuple[AssetCheckKey, Optional[PartitionsDefinition]]],
+    initial_states: Mapping[AssetCheckKey, "AssetCheckState"],
+) -> Mapping[AssetCheckKey, "AssetCheckState"]:
+    check_keys = [key for key, _ in keys]
+    partitions_defs_by_key = {key: partitions_def for key, partitions_def in keys}
+
+    # we prefer to do a single fetch for all keys, so we use the minimum storage id of the initial states
+    storage_id = min((state.latest_storage_id for state in initial_states.values()), default=0)
+    infos = instance.event_log_storage.get_asset_check_partition_info(
+        check_keys, after_storage_id=storage_id
+    )
+
+    # find the set of run ids we need to fetch to resolve the in-progress runs, and
+    # group the partition infos by check key
+    run_ids_to_fetch: set[str] = set().union(
+        *(state.in_progress_runs.keys() for state in initial_states.values())
+    )
+    infos_by_key: dict[AssetCheckKey, list[AssetCheckPartitionInfo]] = defaultdict(list)
+    for info in infos:
+        infos_by_key[info.check_key].append(info)
+        if info.latest_execution_status == AssetCheckExecutionRecordStatus.PLANNED:
+            run_ids_to_fetch.add(info.latest_planned_run_id)
+
+    # do a bulk fetch for runs across all states
+    finished_runs = (
+        instance.get_runs(
+            filters=RunsFilter(run_ids=list(run_ids_to_fetch), statuses=FINISHED_STATUSES)
+        )
+        if len(run_ids_to_fetch) > 0
+        else []
+    )
+    finished_runs_status_by_id = {run.run_id: run.status for run in finished_runs}
+    return {
+        key: initial_states[key].with_updates(
+            key, partitions_defs_by_key[key], infos_by_key[key], finished_runs_status_by_id
+        )
+        for key in check_keys
+    }
+
+
+def _valid_partition_key(
+    partition_key: Optional[str], partitions_def: Optional[PartitionsDefinition]
+) -> bool:
+    if partitions_def is None:
+        return partition_key is None
+    else:
+        return partition_key is not None and partitions_def.has_partition_key(partition_key)
+
+
+def _process_partition_records(
+    key: AssetCheckKey,
+    partitions_def: Optional[PartitionsDefinition],
+    subsets: StatusSubsets,
+    in_progress_runs: InProgressRuns,
+    partition_infos: Sequence[AssetCheckPartitionInfo],
+) -> tuple[StatusSubsets, InProgressRuns]:
+    """Returns a set of updated subsets based on new partition records and the latest materialization storage ids."""
+    new_subsets = dict(subsets)
+    new_in_progress_runs = dict(in_progress_runs)
+
+    for partition_record in partition_infos:
+        pk = partition_record.partition_key
+        if not _valid_partition_key(pk, partitions_def):
+            continue
+
+        partition_subset = SerializableEntitySubset.from_coercible_value(key, pk, partitions_def)
+
+        if partition_record.latest_execution_status == AssetCheckExecutionRecordStatus.PLANNED:
+            # Add to IN_PROGRESS and track run
+            new_subsets[AssetCheckExecutionResolvedStatus.IN_PROGRESS] = new_subsets[
+                AssetCheckExecutionResolvedStatus.IN_PROGRESS
+            ].compute_union(partition_subset)
+            run_id = partition_record.latest_planned_run_id
+            new_in_progress_runs[run_id] = new_in_progress_runs.get(
+                run_id, SerializableEntitySubset.empty(key, partitions_def)
+            ).compute_union(partition_subset)
+
+        elif partition_record.latest_execution_status in (
+            AssetCheckExecutionRecordStatus.SUCCEEDED,
+            AssetCheckExecutionRecordStatus.FAILED,
+        ):
+            if partition_record.is_current:
+                # Check is current, set appropriate status
+                status = (
+                    AssetCheckExecutionResolvedStatus.SUCCEEDED
+                    if partition_record.latest_execution_status
+                    == AssetCheckExecutionRecordStatus.SUCCEEDED
+                    else AssetCheckExecutionResolvedStatus.FAILED
+                )
+                new_subsets[status] = new_subsets[status].compute_union(partition_subset)
+            else:
+                # new materialization, clear the check status
+                for status in new_subsets:
+                    new_subsets[status] = new_subsets[status].compute_difference(partition_subset)
+
+    return new_subsets, new_in_progress_runs
+
+
+def _process_run_statuses(
+    key: AssetCheckKey,
+    partitions_def: Optional[PartitionsDefinition],
+    subsets: StatusSubsets,
+    in_progress_runs: InProgressRuns,
+    run_statuses: Mapping[str, DagsterRunStatus],
+) -> tuple[StatusSubsets, InProgressRuns]:
+    """Resolve in-progress runs that have completed.
+
+    This checks if any runs tracked in in_progress_runs have finished,
+    and moves their partitions to SKIPPED or EXECUTION_FAILED.
+    """
+    if not in_progress_runs:
+        return subsets, in_progress_runs
+
+    delta_skipped, delta_execution_failed, resolved_run_ids = _resolve_in_progress_subsets(
+        key, partitions_def, in_progress_runs, run_statuses
+    )
+
+    new_in_progress_runs = {
+        run_id: subset
+        for run_id, subset in in_progress_runs.items()
+        if run_id not in resolved_run_ids
+    }
+
+    new_subsets = dict(subsets)
+    new_subsets[AssetCheckExecutionResolvedStatus.IN_PROGRESS] = (
+        new_subsets[AssetCheckExecutionResolvedStatus.IN_PROGRESS]
+        .compute_difference(delta_skipped)
+        .compute_difference(delta_execution_failed)
+    )
+    new_subsets[AssetCheckExecutionResolvedStatus.SKIPPED] = new_subsets[
+        AssetCheckExecutionResolvedStatus.SKIPPED
+    ].compute_union(delta_skipped)
+    new_subsets[AssetCheckExecutionResolvedStatus.EXECUTION_FAILED] = new_subsets[
+        AssetCheckExecutionResolvedStatus.EXECUTION_FAILED
+    ].compute_union(delta_execution_failed)
+
+    return new_subsets, new_in_progress_runs
+
+
+def _resolve_in_progress_subsets(
+    key: AssetCheckKey,
+    partitions_def: Optional[PartitionsDefinition],
+    in_progress_runs: InProgressRuns,
+    run_statuses: Mapping[str, DagsterRunStatus],
+) -> tuple[
+    SerializableEntitySubset[AssetCheckKey],
+    SerializableEntitySubset[AssetCheckKey],
+    set[str],
+]:
+    """Resolve in-progress runs that have completed.
+
+    Returns:
+        Tuple of (delta_skipped, delta_execution_failed, resolved_run_ids)
+    """
+    empty_subset = SerializableEntitySubset.empty(key, partitions_def)
+    delta_skipped = empty_subset
+    delta_execution_failed = empty_subset
+    resolved_run_ids: set[str] = set()
+
+    for run_id, run_status in run_statuses.items():
+        if run_status in FINISHED_STATUSES and run_id in in_progress_runs:
+            resolved_run_ids.add(run_id)
+            run_subset = in_progress_runs[run_id]
+            if run_status == DagsterRunStatus.FAILURE:
+                delta_execution_failed = delta_execution_failed.compute_union(run_subset)
+            else:
+                delta_skipped = delta_skipped.compute_union(run_subset)
+
+    return delta_skipped, delta_execution_failed, resolved_run_ids

dagster/_core/storage/dagster_run.py

@@ -41,6 +41,9 @@ from dagster._utils.tags import get_boolean_tag_value
 
 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,80 @@ class DagsterRun(
     def get_parent_run_id(self) -> Optional[str]:
         return self.tags.get(PARENT_RUN_ID_TAG)
 
+    @property
+    def is_partitioned(self) -> bool:
+        from dagster._core.storage.tags import (
+            ASSET_PARTITION_RANGE_END_TAG,
+            ASSET_PARTITION_RANGE_START_TAG,
+            PARTITION_NAME_TAG,
+        )
+
+        has_partition_tags = any(
+            self.tags.get(tag) is not None
+            for tag in [
+                PARTITION_NAME_TAG,
+                ASSET_PARTITION_RANGE_START_TAG,
+                ASSET_PARTITION_RANGE_END_TAG,
+            ]
+        )
+        return has_partition_tags or self.partitions_subset is not None
+
+    def get_resolved_partitions_subset_for_events(
+        self, partitions_def: Optional["PartitionsDefinition"]
+    ) -> Optional[PartitionsSubset]:
+        """Get the partitions subset targeted by a run based on its partition tags. Does not always use the
+        partition_subset that is directly stored on the run as this can contain a KeyRangesPartitionSubset
+        which can not be deserialized without additional information.
+        """
+        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,
+        )
+
+        # KeyRangesPartitionsSubset cannot be deserialized without access to the instance, so ignore it for now.
+        if self.partitions_subset is not None and not isinstance(
+            self.partitions_subset, KeyRangesPartitionsSubset
+        ):
+            return self.partitions_subset
+
+        # 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.

dagster/_core/storage/event_log/base.py

@@ -19,6 +19,7 @@ from dagster._core.event_api import (
     EventLogRecord,
     EventRecordsFilter,
     EventRecordsResult,
+    PartitionKeyFilter,
     RunStatusChangeRecordsFilter,
 )
 from dagster._core.events import DagsterEventType
@@ -35,6 +36,7 @@ from dagster._core.loader import LoadableBy, LoadingContext
 from dagster._core.storage.asset_check_execution_record import (
     AssetCheckExecutionRecord,
     AssetCheckExecutionRecordStatus,
+    AssetCheckPartitionInfo,
 )
 from dagster._core.storage.dagster_run import DagsterRunStatsSnapshot
 from dagster._core.storage.partition_status_cache import get_and_update_asset_status_cache_value
@@ -48,6 +50,7 @@ from dagster._utils.warnings import deprecation_warning
 
 if TYPE_CHECKING:
     from dagster._core.events.log import EventLogEntry
+    from dagster._core.storage.asset_check_state import AssetCheckState
     from dagster._core.storage.partition_status_cache import AssetStatusCacheValue
 
 
@@ -633,17 +636,52 @@ class EventLogStorage(ABC, MayHaveInstanceWeakref[T_DagsterInstance]):
         limit: int,
         cursor: Optional[int] = None,
         status: Optional[Set[AssetCheckExecutionRecordStatus]] = None,
+        partition_filter: Optional[PartitionKeyFilter] = 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_filter: Optional[PartitionKeyFilter] = None,
     ) -> Mapping[AssetCheckKey, AssetCheckExecutionRecord]:
         """Get the latest executions for a list of asset checks."""
         pass
 
+    @abstractmethod
+    def get_asset_check_partition_info(
+        self,
+        keys: Sequence[AssetCheckKey],
+        after_storage_id: Optional[int] = None,
+        partition_keys: Optional[Sequence[str]] = None,
+    ) -> Sequence[AssetCheckPartitionInfo]:
+        """Get asset check partition records with execution status and planned run info."""
+        pass
+
+    def get_checkpointed_asset_check_state(
+        self, keys: Sequence[AssetCheckKey]
+    ) -> Mapping[AssetCheckKey, "AssetCheckState"]:
+        """Get the current stored asset check state for a list of asset checks and their
+        associated partitions definitions. This method is not guaranteed to return a
+        state object that is up to date with the latest events.
+        """
+        from dagster._core.storage.asset_check_state import AssetCheckState
+
+        return {key: AssetCheckState.empty() for key in keys}
+
+    def get_asset_check_state(
+        self, keys: Sequence[tuple[AssetCheckKey, Optional[PartitionsDefinition]]]
+    ) -> Mapping[AssetCheckKey, "AssetCheckState"]:
+        from dagster._core.storage.asset_check_state import bulk_update_asset_check_state
+
+        return bulk_update_asset_check_state(
+            self._instance,
+            keys,
+            initial_states=self.get_checkpointed_asset_check_state([key for key, _ in keys]),
+        )
+
     @abstractmethod
     def fetch_materializations(
         self,
@@ -740,3 +778,23 @@ class EventLogStorage(ABC, MayHaveInstanceWeakref[T_DagsterInstance]):
         # Base implementation of fetching pool config.  To be overriden for remote storage
         # implementations where the local instance might not match the remote instance.
         return self._instance.get_concurrency_config().pool_config
+
+    def _get_latest_unpartitioned_materialization_storage_ids(
+        self, keys: Sequence[AssetKey]
+    ) -> Mapping[AssetKey, int]:
+        # Returns a mapping of asset key to the latest recorded materialization storage id for the asset,
+        # ignoring partitioned assets. Used purely for the `get_asset_check_partition_info` method across
+        # different storage implementations.
+        asset_records = self.get_asset_records(keys)
+        latest_unpartitioned_materialization_storage_ids = {}
+        for asset_record in asset_records:
+            if (
+                asset_record.asset_entry.last_materialization_record is not None
+                and asset_record.asset_entry.last_materialization_record.event_log_entry.get_dagster_event().partition
+                is None
+            ):
+                latest_unpartitioned_materialization_storage_ids[
+                    asset_record.asset_entry.asset_key
+                ] = asset_record.asset_entry.last_materialization_storage_id
+
+        return latest_unpartitioned_materialization_storage_ids