[pac] Add get_asset_check_partitions storage method

OwenKephart · OwenKephart · commit 707bd1bce906 · 2026-01-22T13:21:34.000-08:00
diff --git a/python_modules/dagster/dagster/_core/storage/asset_check_execution_record.py b/python_modules/dagster/dagster/_core/storage/asset_check_execution_record.py
@@ -2,6 +2,7 @@
 from collections.abc import Iterable
 from typing import NamedTuple, Optional, cast
 
+from dagster_shared.record import record
 from dagster_shared.serdes import deserialize_value
 
 import dagster._check as check
@@ -238,3 +239,30 @@ async def targets_latest_materialization(self, loading_context: LoadingContext)
             )
         else:
             check.failed(f"Unexpected check status {resolved_status}")
+
+
+@record
+class AssetCheckPartitionInfo:
+    check_key: AssetCheckKey
+    partition_key: Optional[str]
+    # the status of the last execution of the check
+    latest_execution_status: AssetCheckExecutionRecordStatus
+    # the run id of the last planned event for the check
+    latest_planned_run_id: str
+    # the storage id of the last event (planned or evaluation) for the check
+    latest_check_event_storage_id: int
+    # the storage id of the last materialization for the asset / partition that this check targets
+    # this is the latest overall materialization, independent of if there has been a check event
+    # that targets it
+    latest_materialization_storage_id: Optional[int]
+    # the storage id of the materialization that the last execution of the check targeted
+    latest_target_materialization_storage_id: Optional[int]
+
+    @property
+    def is_current(self) -> bool:
+        """Returns True if the latest check execution targets the latest materialization event."""
+        return (
+            self.latest_materialization_storage_id is None
+            or self.latest_materialization_storage_id
+            == self.latest_target_materialization_storage_id
+        )
diff --git a/python_modules/dagster/dagster/_core/storage/event_log/base.py b/python_modules/dagster/dagster/_core/storage/event_log/base.py
@@ -36,6 +36,7 @@
 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
@@ -648,6 +649,16 @@ def get_latest_asset_check_execution_by_key(
         """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
+
     @abstractmethod
     def fetch_materializations(
         self,
@@ -744,3 +755,23 @@ def get_pool_config(self) -> PoolConfig:
         # 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
diff --git a/python_modules/dagster/dagster/_core/storage/event_log/sql_event_log.py b/python_modules/dagster/dagster/_core/storage/event_log/sql_event_log.py
@@ -59,6 +59,7 @@
     COMPLETED_ASSET_CHECK_EXECUTION_RECORD_STATUSES,
     AssetCheckExecutionRecord,
     AssetCheckExecutionRecordStatus,
+    AssetCheckPartitionInfo,
 )
 from dagster._core.storage.dagster_run import DagsterRunStatsSnapshot
 from dagster._core.storage.event_log.base import (
@@ -3008,6 +3009,7 @@ def _store_asset_check_evaluation_planned(
                             execution_status=AssetCheckExecutionRecordStatus.PLANNED.value,
                             evaluation_event=serialize_value(event),
                             evaluation_event_timestamp=self._event_insert_timestamp(event),
+                            evaluation_event_storage_id=event_id,
                         )
                         for partition_key in partition_keys
                     ]
@@ -3232,6 +3234,135 @@ def get_latest_asset_check_execution_by_key(
             results[check_key] = AssetCheckExecutionRecord.from_db_row(row, key=check_key)
         return results
 
+    def _get_asset_check_partition_info_for_key(
+        self,
+        check_key: AssetCheckKey,
+        after_storage_id: Optional[int],
+        partition_keys: Optional[Sequence[str]],
+        latest_unpartitioned_materialization_storage_ids: Mapping[AssetKey, int],
+    ) -> Sequence[AssetCheckPartitionInfo]:
+        # Build the base filter conditions
+        filter_conditions = [
+            AssetCheckExecutionsTable.c.asset_key == check_key.asset_key.to_string(),
+            AssetCheckExecutionsTable.c.check_name == check_key.name,
+            # Historical records may have NULL in the evaluation_event_storage_id column for
+            # PLANNED events
+            AssetCheckExecutionsTable.c.evaluation_event_storage_id.isnot(None),
+        ]
+        if partition_keys is not None:
+            filter_conditions.append(AssetCheckExecutionsTable.c.partition.in_(partition_keys))
+
+        # Subquery to find the max id for each partition
+        latest_check_ids_subquery = db_subquery(
+            db_select(
+                [
+                    db.func.max(AssetCheckExecutionsTable.c.id).label("id"),
+                    AssetCheckExecutionsTable.c.partition.label("partition"),
+                ]
+            )
+            .where(db.and_(*filter_conditions))
+            .group_by(AssetCheckExecutionsTable.c.partition),
+            "latest_check_ids_subquery",
+        )
+
+        # Subquery to find the latest materialization storage id for each partition of the
+        # target asset. Note: we don't filter by after_storage_id here because we always want
+        # to return the latest materialization storage id, even if it's older than after_storage_id.
+        latest_materialization_ids_subquery = self._latest_event_ids_by_partition_subquery(
+            check_key.asset_key,
+            [DagsterEventType.ASSET_MATERIALIZATION],
+            asset_partitions=partition_keys,
+        )
+
+        # Main query to get all columns for the latest records, joined with latest
+        # materialization storage ids
+        query = db_select(
+            [
+                AssetCheckExecutionsTable.c.id,
+                AssetCheckExecutionsTable.c.partition,
+                AssetCheckExecutionsTable.c.execution_status,
+                AssetCheckExecutionsTable.c.evaluation_event_storage_id,
+                AssetCheckExecutionsTable.c.materialization_event_storage_id,
+                AssetCheckExecutionsTable.c.run_id,
+                latest_materialization_ids_subquery.c.id.label("latest_materialization_storage_id"),
+            ]
+        ).select_from(
+            AssetCheckExecutionsTable.join(
+                latest_check_ids_subquery,
+                AssetCheckExecutionsTable.c.id == latest_check_ids_subquery.c.id,
+            ).join(
+                latest_materialization_ids_subquery,
+                AssetCheckExecutionsTable.c.partition
+                == latest_materialization_ids_subquery.c.partition,
+                isouter=True,
+            )
+        )
+
+        # these filters are applied to the main query rather than the individual subqueries to ensure
+        # we don't miss records that only have a new materialization or a new check execution but not both
+        if after_storage_id is not None:
+            query = query.where(
+                db.or_(
+                    AssetCheckExecutionsTable.c.evaluation_event_storage_id > after_storage_id,
+                    latest_materialization_ids_subquery.c.id > after_storage_id,
+                )
+            )
+
+        with self.index_connection() as conn:
+            rows = db_fetch_mappings(conn, query)
+
+        return [
+            AssetCheckPartitionInfo(
+                check_key=check_key,
+                partition_key=row["partition"],
+                latest_execution_status=AssetCheckExecutionRecordStatus(row["execution_status"]),
+                latest_target_materialization_storage_id=row["materialization_event_storage_id"],
+                latest_planned_run_id=row["run_id"],
+                latest_check_event_storage_id=row["evaluation_event_storage_id"],
+                latest_materialization_storage_id=max(
+                    filter(
+                        None,
+                        [
+                            row["latest_materialization_storage_id"],
+                            latest_unpartitioned_materialization_storage_ids.get(
+                                check_key.asset_key
+                            ),
+                        ],
+                    ),
+                    default=None,
+                ),
+            )
+            for row in rows
+        ]
+
+    def get_asset_check_partition_info(
+        self,
+        keys: Sequence[AssetCheckKey],
+        after_storage_id: Optional[int] = None,
+        partition_keys: Optional[Sequence[str]] = None,
+    ) -> Sequence[AssetCheckPartitionInfo]:
+        check.list_param(keys, "keys", of_type=AssetCheckKey)
+        check.opt_int_param(after_storage_id, "after_storage_id")
+
+        infos = []
+        latest_unpartitioned_materialization_storage_ids = (
+            self._get_latest_unpartitioned_materialization_storage_ids(
+                list(set(key.asset_key for key in keys))
+            )
+        )
+        # the inner query is not feasible to join in a single query because the latest materialization ids subquery,
+        # so for now we fetch the info for each key separately
+        for key in keys:
+            infos.extend(
+                self._get_asset_check_partition_info_for_key(
+                    key,
+                    after_storage_id,
+                    partition_keys,
+                    latest_unpartitioned_materialization_storage_ids,
+                )
+            )
+        return infos
+
     @property
     def supports_asset_checks(self):  # pyright: ignore[reportIncompatibleMethodOverride]
         return self.has_table(AssetCheckExecutionsTable.name)
diff --git a/python_modules/dagster/dagster/_core/storage/event_log/sqlite/sqlite_event_log.py b/python_modules/dagster/dagster/_core/storage/event_log/sqlite/sqlite_event_log.py
@@ -276,7 +276,12 @@ def store_event(self, event: EventLogEntry) -> None:
             self.store_asset_event_tags([event], [event_id])
 
         if event.is_dagster_event and event.dagster_event_type in ASSET_CHECK_EVENTS:
-            self.store_asset_check_event(event, None)
+            # mirror the event in the cross-run index database
+            with self.index_connection() as conn:
+                result = conn.execute(insert_event_statement)
+                event_id = result.inserted_primary_key[0]
+
+            self.store_asset_check_event(event, event_id)
 
         if event.is_dagster_event and event.dagster_event_type in EVENT_TYPE_TO_PIPELINE_RUN_STATUS:
             # should mirror run status change events in the index shard
diff --git a/python_modules/dagster/dagster/_core/storage/legacy_storage.py b/python_modules/dagster/dagster/_core/storage/legacy_storage.py
@@ -16,6 +16,7 @@
 from dagster._core.storage.asset_check_execution_record import (
     AssetCheckExecutionRecord,
     AssetCheckExecutionRecordStatus,
+    AssetCheckPartitionInfo,
 )
 from dagster._core.storage.base_storage import DagsterStorage
 from dagster._core.storage.event_log.base import (
@@ -765,6 +766,16 @@ def get_latest_asset_check_execution_by_key(
             check_keys, partition=partition
         )
 
+    def get_asset_check_partition_info(
+        self,
+        keys: Sequence["AssetCheckKey"],
+        after_storage_id: Optional[int] = None,
+        partition_keys: Optional[Sequence[str]] = None,
+    ) -> Sequence[AssetCheckPartitionInfo]:
+        return self._storage.event_log_storage.get_asset_check_partition_info(
+            keys=keys, after_storage_id=after_storage_id, partition_keys=partition_keys
+        )
+
 
 class LegacyScheduleStorage(ScheduleStorage, ConfigurableClass):
     def __init__(self, storage: DagsterStorage, inst_data: Optional[ConfigurableClassData] = None):
diff --git a/python_modules/dagster/dagster_tests/storage_tests/utils/event_log_storage.py b/python_modules/dagster/dagster_tests/storage_tests/utils/event_log_storage.py
