[pac] Add get_asset_check_partitions storage method

Author: OwenKephart

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,28 @@ 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
+    last_execution_status: AssetCheckExecutionRecordStatus
+    # the storage id of the materialization that the last execution of the check targeted
+    last_execution_target_materialization_storage_id: Optional[int]
+    # the run id of the last planned event for the check
+    last_planned_run_id: str
+    # the storage id of the last event for the check
+    last_storage_id: int
+    # the storage id of the last materialization for the asset this check targets
+    last_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.last_materialization_storage_id is None
+            or self.last_materialization_storage_id
+            == self.last_execution_target_materialization_storage_id
+        )

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,15 @@ 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,
+    ) -> Sequence[AssetCheckPartitionInfo]:
+        """Get asset check partition records with execution status and planned run info."""
+        pass
+
     @abstractmethod
     def fetch_materializations(
         self,

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,106 @@ 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]
+    ) -> 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),
+        ]
+
+        # 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],
+        )
+
+        # 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"],
+                last_execution_status=AssetCheckExecutionRecordStatus(row["execution_status"]),
+                last_execution_target_materialization_storage_id=row[
+                    "materialization_event_storage_id"
+                ],
+                last_planned_run_id=row["run_id"],
+                last_storage_id=row["evaluation_event_storage_id"],
+                last_materialization_storage_id=row["latest_materialization_storage_id"],
+            )
+            for row in rows
+        ]
+
+    def get_asset_check_partition_info(
+        self,
+        keys: Sequence[AssetCheckKey],
+        after_storage_id: Optional[int] = None,
+    ) -> Sequence[AssetCheckPartitionInfo]:
+        check.list_param(keys, "keys", of_type=AssetCheckKey)
+        check.opt_int_param(after_storage_id, "after_storage_id")
+
+        infos = []
+        # 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))
+        return infos
+
     @property
     def supports_asset_checks(self):  # pyright: ignore[reportIncompatibleMethodOverride]
         return self.has_table(AssetCheckExecutionsTable.name)

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

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,15 @@ 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,
+    ) -> Sequence[AssetCheckPartitionInfo]:
+        return self._storage.event_log_storage.get_asset_check_partition_info(
+            keys=keys, after_storage_id=after_storage_id
+        )
+
 
 class LegacyScheduleStorage(ScheduleStorage, ConfigurableClass):
     def __init__(self, storage: DagsterStorage, inst_data: Optional[ConfigurableClassData] = None):

python_modules/dagster/dagster_tests/storage_tests/utils/event_log_storage.py

@@ -7092,6 +7092,47 @@ def test_asset_check_partitioned_planned_and_evaluation(
         assert latest_c[check_key].partition == "c"
         assert latest_c[check_key].status == AssetCheckExecutionRecordStatus.PLANNED
 
+        # Test get_asset_check_partition_records - returns all partitions with latest status
+        partition_records = storage.get_asset_check_partition_info([check_key])
+        assert len(partition_records) == 3
+
+        records_by_partition = {r.partition_key: r for r in partition_records}
+        assert set(records_by_partition.keys()) == {"a", "b", "c"}
+
+        # Verify each partition has correct status
+        assert (
+            records_by_partition["a"].last_execution_status
+            == AssetCheckExecutionRecordStatus.SUCCEEDED
+        )
+        assert (
+            records_by_partition["b"].last_execution_status
+            == AssetCheckExecutionRecordStatus.FAILED
+        )
+        assert (
+            records_by_partition["c"].last_execution_status
+            == AssetCheckExecutionRecordStatus.PLANNED
+        )
+
+        # Verify all records have the same run_id (all planned in same run)
+        assert records_by_partition["a"].last_planned_run_id == run_id
+        assert records_by_partition["b"].last_planned_run_id == run_id
+        assert records_by_partition["c"].last_planned_run_id == run_id
+
+        # Verify last_storage_id is set for all records (this is the row id in the table)
+        assert records_by_partition["a"].last_storage_id is not None
+        assert records_by_partition["b"].last_storage_id is not None
+        assert records_by_partition["c"].last_storage_id is not None
+
+        # Verify last_materialization_storage_id is None for all records (no materializations)
+        assert records_by_partition["a"].last_materialization_storage_id is None
+        assert records_by_partition["b"].last_materialization_storage_id is None
+        assert records_by_partition["c"].last_materialization_storage_id is None
+
+        filtered_records = storage.get_asset_check_partition_info(
+            [check_key], after_storage_id=999999
+        )
+        assert len(filtered_records) == 0
+
     def test_asset_check_partitioned_multiple_runs_same_partition(
         self,
         storage: EventLogStorage,
@@ -7107,22 +7148,54 @@ def test_asset_check_partitioned_multiple_runs_same_partition(
         partitions_def = dg.StaticPartitionsDefinition(["a"])
         partitions_subset = partitions_def.subset_with_partition_keys(["a"])
 
-        # Run 1: Store planned + evaluation for partition "a" with passed=True
+        # Run 1: Store planned event for partition "a"
         storage.store_event(
             _create_check_planned_event(run_id_1, check_key, partitions_subset=partitions_subset)
         )
+
+        # status for partition "a" should be PLANNED
+        partition_records = storage.get_asset_check_partition_info([check_key])
+        assert len(partition_records) == 1
+        record = partition_records[0]
+        assert record.partition_key == "a"
+        assert record.last_execution_status == AssetCheckExecutionRecordStatus.PLANNED
+        assert record.last_planned_run_id == run_id_1
+
+        # Run 1: Now store evaluation event for partition "a" with passed=True
         storage.store_event(
             _create_check_evaluation_event(run_id_1, check_key, passed=True, partition="a")
         )
 
+        # status for partition "a" should be SUCCEEDED
+        partition_records = storage.get_asset_check_partition_info([check_key])
+        assert len(partition_records) == 1
+        record = partition_records[0]
+        assert record.partition_key == "a"
+        assert record.last_execution_status == AssetCheckExecutionRecordStatus.SUCCEEDED
+        assert record.last_planned_run_id == run_id_1
+
         # Run 2: Store planned + evaluation for partition "a" with passed=False
         storage.store_event(
             _create_check_planned_event(run_id_2, check_key, partitions_subset=partitions_subset)
         )
+
+        # back to PLANNED
+        partition_records = storage.get_asset_check_partition_info([check_key])
+        record = partition_records[0]
+        assert record.last_execution_status == AssetCheckExecutionRecordStatus.PLANNED
+        assert record.last_planned_run_id == run_id_2
+
+        # Run 2: Now store evaluation event for partition "a" with passed=False
         storage.store_event(
             _create_check_evaluation_event(run_id_2, check_key, passed=False, partition="a")
         )
 
+        # onto FAILED
+        partition_records = storage.get_asset_check_partition_info([check_key])
+        record = partition_records[0]
+        assert record.last_execution_status == AssetCheckExecutionRecordStatus.FAILED
+        assert record.last_planned_run_id == run_id_2
+
         # Verify get_asset_check_execution_history returns 2 records for partition "a"
         checks = storage.get_asset_check_execution_history(check_key, limit=10, partition="a")
         assert len(checks) == 2
@@ -7149,6 +7222,16 @@ def test_asset_check_partitioned_multiple_runs_same_partition(
         assert check_key in latest_overall
         assert latest_overall[check_key].run_id == run_id_2
 
+        # Test get_asset_check_partition_records returns only the latest record per partition
+        partition_records = storage.get_asset_check_partition_info([check_key])
+        assert len(partition_records) == 1  # Only partition "a" exists
+
+        record = partition_records[0]
+        assert record.partition_key == "a"
+        # Should be the latest execution (run_id_2, FAILED)
+        assert record.last_execution_status == AssetCheckExecutionRecordStatus.FAILED
+        assert record.last_planned_run_id == run_id_2
+
     def test_asset_check_partitioned_with_target_materialization(
         self,
         storage: EventLogStorage,
@@ -7210,6 +7293,12 @@ def test_asset_check_partitioned_with_target_materialization(
         assert check_data.target_materialization_data
         assert check_data.target_materialization_data.storage_id == m1_storage_id
 
+        # Verify get_asset_check_partition_records returns M1 as the latest materialization
+        partition_records = storage.get_asset_check_partition_info([check_key])
+        assert len(partition_records) == 1
+        record = partition_records[0]
+        assert record.last_materialization_storage_id == m1_storage_id
+
         # Store materialization M2 for partition "a"
         storage.store_event(_create_materialization_event(run_id_2, asset_key, partition="a"))
 
@@ -7237,6 +7326,19 @@ def test_asset_check_partitioned_with_target_materialization(
         assert check_data.target_materialization_data.storage_id == m1_storage_id
         assert check_data.target_materialization_data.storage_id != m2_storage_id
 
+        # Test get_asset_check_partition_records includes target_materialization_storage_id
+        partition_records = storage.get_asset_check_partition_info([check_key])
+        assert len(partition_records) == 1
+
+        record = partition_records[0]
+        assert record.partition_key == "a"
+        assert record.last_execution_status == AssetCheckExecutionRecordStatus.SUCCEEDED
+        assert record.last_planned_run_id == run_id_1
+        # Verify last_execution_target_materialization_storage_id matches M1 (what check targeted)
+        assert record.last_execution_target_materialization_storage_id == m1_storage_id
+        # Verify last_materialization_storage_id matches M2 (the current latest materialization)
+        assert record.last_materialization_storage_id == m2_storage_id
+
 
 def _create_check_planned_event(
     run_id: str,