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

dagster/_core/definitions/decorators/asset_check_decorator.py

@@ -28,6 +28,9 @@ from dagster._core.definitions.decorators.decorator_assets_definition_builder im
 from dagster._core.definitions.decorators.op_decorator import _Op
 from dagster._core.definitions.events import AssetKey, CoercibleToAssetKey
 from dagster._core.definitions.output import Out
+from dagster._core.definitions.partitions.definition.partitions_definition import (
+    PartitionsDefinition,
+)
 from dagster._core.definitions.policy import RetryPolicy
 from dagster._core.definitions.source_asset import SourceAsset
 from dagster._core.definitions.utils import DEFAULT_OUTPUT
@@ -113,6 +116,7 @@ def asset_check(
     metadata: Optional[Mapping[str, Any]] = None,
     automation_condition: Optional[AutomationCondition[AssetCheckKey]] = None,
     pool: Optional[str] = None,
+    partitions_def: Optional[PartitionsDefinition] = None,
 ) -> Callable[[AssetCheckFunction], AssetChecksDefinition]:
     """Create a definition for how to execute an asset check.
 
@@ -151,6 +155,7 @@ def asset_check(
         automation_condition (Optional[AutomationCondition]): An AutomationCondition which determines
             when this check should be executed.
         pool (Optional[str]): A string that identifies the concurrency pool that governs this asset check's execution.
+        partitions_def (Optional[PartitionsDefinition]): The PartitionsDefinition for this asset check.
 
     Produces an :py:class:`AssetChecksDefinition` object.
 
@@ -218,6 +223,7 @@ def asset_check(
             blocking=blocking,
             metadata=metadata,
             automation_condition=automation_condition,
+            partitions_def=partitions_def,
         )
 
         resource_defs_for_execution = wrap_resources_for_execution(resource_defs)

dagster/_core/event_api.py

@@ -4,6 +4,7 @@ from datetime import datetime
 from enum import Enum
 from typing import Literal, NamedTuple, Optional, TypeAlias, Union
 
+from dagster_shared.record import record
 from dagster_shared.seven import json
 
 import dagster._check as check
@@ -342,6 +343,15 @@ class AssetRecordsFilter(
         return None
 
 
+@record
+class PartitionKeyFilter:
+    """Filter for the partition keys that should be included in the result. Allows for filtering on
+    unpartitioned assets, specific partition keys, or a combination of both.
+    """
+
+    key: Optional[str]
+
+
 @whitelist_for_serdes
 class RunStatusChangeRecordsFilter(
     NamedTuple(

dagster/_core/execution/context/asset_check_execution_context.py

@@ -6,6 +6,8 @@ from dagster._annotations import public
 from dagster._core.definitions.asset_checks.asset_check_spec import AssetCheckKey, AssetCheckSpec
 from dagster._core.definitions.job_definition import JobDefinition
 from dagster._core.definitions.op_definition import OpDefinition
+from dagster._core.definitions.partitions.partition_key_range import PartitionKeyRange
+from dagster._core.definitions.partitions.utils.time_window import TimeWindow
 from dagster._core.definitions.repository_definition.repository_definition import (
     RepositoryDefinition,
 )
@@ -135,6 +137,43 @@ class AssetCheckExecutionContext:
     def get_step_execution_context(self) -> StepExecutionContext:
         return self.op_execution_context.get_step_execution_context()
 
+    #### partition related
+    @public
+    @property
+    @_copy_docs_from_op_execution_context
+    def has_partition_key(self) -> bool:
+        return self.op_execution_context.has_partition_key
+
+    @public
+    @property
+    @_copy_docs_from_op_execution_context
+    def partition_key(self) -> str:
+        return self.op_execution_context.partition_key
+
+    @public
+    @property
+    @_copy_docs_from_op_execution_context
+    def partition_keys(self) -> Sequence[str]:
+        return self.op_execution_context.partition_keys
+
+    @public
+    @property
+    @_copy_docs_from_op_execution_context
+    def has_partition_key_range(self) -> bool:
+        return self.op_execution_context.has_partition_key_range
+
+    @public
+    @property
+    @_copy_docs_from_op_execution_context
+    def partition_key_range(self) -> PartitionKeyRange:
+        return self.op_execution_context.partition_key_range
+
+    @public
+    @property
+    @_copy_docs_from_op_execution_context
+    def partition_time_window(self) -> TimeWindow:
+        return self.op_execution_context.partition_time_window
+
     # misc
 
     @public

dagster/_core/execution/plan/execute_step.py

@@ -97,9 +97,6 @@ def _process_user_event(
         asset_key = _resolve_asset_result_asset_key(user_event, assets_def)
         output_name = assets_def.get_output_name_for_asset_key(asset_key)
 
-        for check_result in user_event.check_results or []:
-            yield from _process_user_event(step_context, check_result)
-
         with disable_dagster_warnings():
             if isinstance(user_event, MaterializeResult):
                 value = user_event.value
@@ -112,6 +109,10 @@ def _process_user_event(
                 data_version=user_event.data_version,
                 tags=user_event.tags,
             )
+
+        for check_result in user_event.check_results or []:
+            yield from _process_user_event(step_context, check_result)
+
     elif isinstance(user_event, AssetCheckResult):
         asset_check_evaluation = user_event.to_asset_check_evaluation(step_context)
         assets_def = _get_assets_def_for_step(step_context, user_event)

dagster/_core/instance/runs/run_domain.py

@@ -2,7 +2,7 @@ import logging
 import os
 import warnings
 from collections.abc import Mapping, Sequence, Set
-from typing import TYPE_CHECKING, Any, Optional, cast
+from typing import TYPE_CHECKING, Any, Optional, cast, overload
 
 import dagster._check as check
 from dagster._core.definitions.asset_checks.asset_check_evaluation import (
@@ -50,10 +50,13 @@ from dagster._utils.warnings import disable_dagster_warnings
 if TYPE_CHECKING:
     from dagster._core.definitions.asset_checks.asset_check_spec import AssetCheckKey
     from dagster._core.definitions.assets.graph.base_asset_graph import (
+        AssetCheckNode,
         BaseAssetGraph,
         BaseAssetNode,
+        BaseEntityNode,
     )
     from dagster._core.definitions.job_definition import JobDefinition
+    from dagster._core.definitions.partitions.definition import PartitionsDefinition
     from dagster._core.definitions.repository_definition.repository_definition import (
         RepositoryLoadData,
     )
@@ -66,10 +69,7 @@ if TYPE_CHECKING:
     from dagster._core.remote_representation.code_location import CodeLocation
     from dagster._core.remote_representation.external import RemoteJob
     from dagster._core.snap import ExecutionPlanSnapshot, JobSnap
-    from dagster._core.snap.execution_plan_snapshot import (
-        ExecutionStepOutputSnap,
-        ExecutionStepSnap,
-    )
+    from dagster._core.snap.execution_plan_snapshot import ExecutionStepSnap
     from dagster._core.workspace.context import BaseWorkspaceRequestContext
 
 
@@ -108,6 +108,7 @@ class RunDomain:
         """Create a run with the given parameters."""
         from dagster._core.definitions.asset_key import AssetCheckKey
         from dagster._core.definitions.assets.graph.remote_asset_graph import RemoteAssetGraph
+        from dagster._core.definitions.partitions.context import partition_loading_context
         from dagster._core.remote_origin import RemoteJobOrigin
         from dagster._core.snap import ExecutionPlanSnapshot, JobSnap
         from dagster._utils.tags import normalize_tags
@@ -256,7 +257,8 @@ class RunDomain:
         dagster_run = self._instance.run_storage.add_run(dagster_run)
 
         if execution_plan_snapshot and not assets_are_externally_managed(dagster_run):
-            self._log_asset_planned_events(dagster_run, execution_plan_snapshot, asset_graph)
+            with partition_loading_context(dynamic_partitions_store=self._instance):
+                self._log_asset_planned_events(dagster_run, execution_plan_snapshot, asset_graph)
 
         return dagster_run
 
@@ -315,8 +317,8 @@ class RunDomain:
                     adjusted_output = output
 
                     if asset_key:
-                        asset_node = self._get_repo_scoped_asset_node(
-                            asset_graph, asset_key, remote_job_origin
+                        asset_node = self._get_repo_scoped_entity_node(
+                            asset_key, asset_graph, remote_job_origin
                         )
                         if asset_node:
                             partitions_definition = asset_node.partitions_def
@@ -767,12 +769,28 @@ class RunDomain:
             {key for key in to_reexecute if isinstance(key, AssetCheckKey)},
         )
 
-    def _get_repo_scoped_asset_node(
+    @overload
+    def _get_repo_scoped_entity_node(
         self,
+        key: AssetKey,
+        asset_graph: "BaseAssetGraph",
+        remote_job_origin: Optional["RemoteJobOrigin"] = None,
+    ) -> Optional["BaseAssetNode"]: ...
+
+    @overload
+    def _get_repo_scoped_entity_node(
+        self,
+        key: "AssetCheckKey",
         asset_graph: "BaseAssetGraph",
-        asset_key: AssetKey,
         remote_job_origin: Optional["RemoteJobOrigin"] = None,
-    ) -> Optional["BaseAssetNode"]:
+    ) -> Optional["AssetCheckNode"]: ...
+
+    def _get_repo_scoped_entity_node(
+        self,
+        key: "EntityKey",
+        asset_graph: "BaseAssetGraph",
+        remote_job_origin: Optional["RemoteJobOrigin"] = None,
+    ) -> Optional["BaseEntityNode"]:
         from dagster._core.definitions.assets.graph.remote_asset_graph import (
             RemoteWorkspaceAssetGraph,
         )
@@ -783,16 +801,29 @@ class RunDomain:
         # in all cases, return the BaseAssetNode for the supplied asset key if it exists.
         if isinstance(asset_graph, RemoteWorkspaceAssetGraph):
             return cast(
-                "Optional[BaseAssetNode]",
+                "Optional[BaseEntityNode]",
                 asset_graph.get_repo_scoped_node(
-                    asset_key, check.not_none(remote_job_origin).repository_origin.get_selector()
+                    key, check.not_none(remote_job_origin).repository_origin.get_selector()
                 ),
             )
 
-        if not asset_graph.has(asset_key):
+        if not asset_graph.has(key):
             return None
 
-        return asset_graph.get(asset_key)
+        return asset_graph.get(key)
+
+    def _get_partitions_def(
+        self,
+        key: "EntityKey",
+        asset_graph: "BaseAssetGraph",
+        remote_job_origin: Optional["RemoteJobOrigin"],
+        run: "DagsterRun",
+    ) -> Optional["PartitionsDefinition"]:
+        # don't fetch the partitions def if the run is not partitioned
+        if not run.is_partitioned:
+            return None
+        entity_node = self._get_repo_scoped_entity_node(key, asset_graph, remote_job_origin)
+        return entity_node.partitions_def if entity_node else None
 
     def _log_asset_planned_events(
         self,
@@ -819,7 +850,7 @@ class RunDomain:
                     if asset_key:
                         events.extend(
                             self.get_materialization_planned_events_for_asset(
-                                dagster_run, asset_key, job_name, step, output, asset_graph
+                                dagster_run, asset_key, job_name, step, asset_graph
                             )
                         )
 
@@ -830,6 +861,13 @@ class RunDomain:
                         target_asset_key = asset_check_key.asset_key
                         check_name = asset_check_key.name
 
+                        partitions_def = self._get_partitions_def(
+                            asset_check_key, asset_graph, dagster_run.remote_job_origin, dagster_run
+                        )
+                        partitions_subset = dagster_run.get_resolved_partitions_subset_for_events(
+                            partitions_def
+                        )
+
                         event = DagsterEvent(
                             event_type_value=DagsterEventType.ASSET_CHECK_EVALUATION_PLANNED.value,
                             job_name=job_name,
@@ -838,8 +876,9 @@ class RunDomain:
                                 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,
                         )
@@ -878,94 +917,26 @@ class RunDomain:
         asset_key: AssetKey,
         job_name: str,
         step: "ExecutionStepSnap",
-        output: "ExecutionStepOutputSnap",
         asset_graph: "BaseAssetGraph[BaseAssetNode]",
     ) -> 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),
-        )
-
-        if partition_tag and (partition_range_start or partition_range_end):
-            raise DagsterInvariantViolationError(
-                f"Cannot have {ASSET_PARTITION_RANGE_START_TAG} or"
-                f" {ASSET_PARTITION_RANGE_END_TAG} set along with"
-                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 = check.not_none(
-                self._get_repo_scoped_asset_node(
-                    asset_graph, asset_key, dagster_run.remote_job_origin
-                )
-            )
-
-            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"
-                )
-
-            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"
+        partitions_def = self._get_partitions_def(
+            asset_key, asset_graph, dagster_run.remote_job_origin, dagster_run
         )
 
-        if not individual_partitions and not partitions_subset:
+        partitions_subset = dagster_run.get_resolved_partitions_subset_for_events(partitions_def)
+        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),
             )
             events.append(materialization_planned)
-        elif individual_partitions:
-            for individual_partition in individual_partitions:
-                materialization_planned = DagsterEvent.build_asset_materialization_planned_event(
-                    job_name,
-                    step.key,
-                    AssetMaterializationPlannedData(
-                        asset_key,
-                        partition=individual_partition,
-                        partitions_subset=partitions_subset,
-                    ),
-                )
-                events.append(materialization_planned)
-
-        else:
+        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,
@@ -974,6 +945,18 @@ class RunDomain:
                 ),
             )
             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
 

dagster/_core/remote_representation/external_data.py

@@ -899,6 +899,7 @@ class AssetCheckNodeSnap(IHaveNew):
     additional_asset_keys: Sequence[AssetKey]
     automation_condition: Optional[AutomationCondition]
     automation_condition_snapshot: Optional[AutomationConditionSnapshot]
+    partitions_def_snapshot: Optional[PartitionsSnap]
 
     def __new__(
         cls,
@@ -911,6 +912,7 @@ class AssetCheckNodeSnap(IHaveNew):
         additional_asset_keys: Optional[Sequence[AssetKey]] = None,
         automation_condition: Optional[AutomationCondition] = None,
         automation_condition_snapshot: Optional[AutomationConditionSnapshot] = None,
+        partitions_def_snapshot: Optional[PartitionsSnap] = None,
     ):
         return super().__new__(
             cls,
@@ -923,6 +925,7 @@ class AssetCheckNodeSnap(IHaveNew):
             additional_asset_keys=additional_asset_keys or [],
             automation_condition=automation_condition,
             automation_condition_snapshot=automation_condition_snapshot,
+            partitions_def_snapshot=partitions_def_snapshot,
         )
 
     @property
@@ -1211,6 +1214,9 @@ def asset_check_node_snaps_from_repo(repo: RepositoryDefinition) -> Sequence[Ass
                 additional_asset_keys=[dep.asset_key for dep in spec.additional_deps],
                 automation_condition=automation_condition,
                 automation_condition_snapshot=automation_condition_snapshot,
+                partitions_def_snapshot=PartitionsSnap.from_def(spec.partitions_def)
+                if spec.partitions_def
+                else None,
             )
         )
 

dagster/_core/storage/asset_check_execution_record.py

@@ -2,6 +2,7 @@ import enum
 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
@@ -60,6 +61,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,6 +74,7 @@ class AssetCheckExecutionRecord(
         status: AssetCheckExecutionRecordStatus,
         event: Optional[EventLogEntry],
         create_timestamp: float,
+        partition: Optional[str],
     ):
         check.inst_param(key, "key", AssetCheckKey)
         check.int_param(id, "id")
@@ -79,6 +82,7 @@ class AssetCheckExecutionRecord(
         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 +109,7 @@ class AssetCheckExecutionRecord(
             status=status,
             event=event,
             create_timestamp=create_timestamp,
+            partition=partition,
         )
 
     @property
@@ -129,6 +134,7 @@ class AssetCheckExecutionRecord(
                 else None
             ),
             create_timestamp=utc_datetime_from_naive(row["create_timestamp"]).timestamp(),
+            partition=row["partition"],
         )
 
     @classmethod
@@ -168,17 +174,28 @@ class AssetCheckExecutionRecord(
             check.failed(f"Unexpected status {self.status}")
 
     async def targets_latest_materialization(self, loading_context: LoadingContext) -> bool:
-        from dagster._core.storage.event_log.base import AssetRecord
+        from dagster._core.storage.event_log.base import AssetRecord, AssetRecordsFilter
 
         resolved_status = await self.resolve_status(loading_context)
         if resolved_status == AssetCheckExecutionResolvedStatus.IN_PROGRESS:
             # all in-progress checks execute against the latest version
             return True
 
-        asset_record = await AssetRecord.gen(loading_context, self.key.asset_key)
-        latest_materialization = (
-            asset_record.asset_entry.last_materialization_record if asset_record else None
-        )
+        if self.partition is None:
+            asset_record = await AssetRecord.gen(loading_context, self.key.asset_key)
+            latest_materialization = (
+                asset_record.asset_entry.last_materialization_record if asset_record else None
+            )
+        else:
+            records = loading_context.instance.fetch_materializations(
+                AssetRecordsFilter(
+                    asset_key=self.key.asset_key,
+                    asset_partitions=[self.partition],
+                ),
+                limit=1,
+            )
+            latest_materialization = records.records[0] if records.records else None
+
         if not latest_materialization:
             # no previous materialization, so it's executing against the lastest version
             return True
@@ -222,3 +239,30 @@ class AssetCheckExecutionRecord(
             )
         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
+        )