fix: resolve asset check keys (#25686)

## Summary & Motivation
Update the logic that selects which asset checks are run with DA.

The logic is now to resolve to `asset_check_keys` that
- target one of the keys in `asset_keys`
- are in the same repo that we're targeting
- do not have an automation condition defined on it
- are explicitly requested

I've also included the optimization: if no checks on the `asset_key` set
have an automation condition, we pass None.

This is all behind a `AUTOMATION_RESOLVE_ASSET_CHECK_KEYS` feature gate.

Here's a summary of the behavior now when a DA run is kicked off (gate
OFF is the current behavior in production).

```
case                     gate  checks that EXIST                       checks that RUN       pills on the run
-----------------------  ----  --------------------------------------  --------------------  --------------------
unconditioned            OFF   row_count (u), not_null (u)             row_count, not_null   2 (row_count, not_null)
unconditioned            ON    row_count (u), not_null (u)             row_count, not_null   2 (row_count, not_null)
mixed                    OFF   row_count (c), not_null (u)             row_count             row_count
mixed                    ON    row_count (c), not_null (u)             row_count, not_null   2 (row_count, not_null)
conditioned_not_firing   OFF   yearly (c), not_null (u)                yearly, not_null      2 (yearly, not_null)
conditioned_not_firing   ON    yearly (c), not_null (u)                not_null              not_null
conditioned_only         OFF   row_count (c)                           row_count             row_count
conditioned_only         ON    row_count (c)                           row_count             row_count

(c) = conditioned (has its own automation_condition)   (u) = unconditioned (rides along)
```

## Test Plan
Unit & manual (local).

Internal-RevId: 421c76c3d19baa6cd46e9cac723c37ad7f8290b0

Author: clnoll

python_modules/dagster/dagster/_core/definitions/automation_tick_evaluation_context.py

@@ -14,6 +14,7 @@
 from dagster._core.definitions.asset_selection import AssetSelection
 from dagster._core.definitions.assets.graph.asset_graph_subset import AssetGraphSubset
 from dagster._core.definitions.assets.graph.base_asset_graph import BaseAssetGraph
+from dagster._core.definitions.assets.graph.remote_asset_graph import RemoteWorkspaceAssetGraph
 from dagster._core.definitions.backfill_policy import BackfillPolicy, BackfillPolicyType
 from dagster._core.definitions.declarative_automation.automation_condition import (
     AutomationCondition,
@@ -30,6 +31,7 @@
 from dagster._core.definitions.partitions.context import use_partition_loading_context
 from dagster._core.definitions.partitions.definition import PartitionsDefinition
 from dagster._core.definitions.run_request import RunRequest
+from dagster._core.definitions.selector import RepositorySelector
 from dagster._core.instance import DynamicPartitionsStore
 from dagster._core.storage.tags import (
     ASSET_PARTITION_RANGE_END_TAG,
@@ -81,6 +83,8 @@ def __init__(
         self._auto_observe_asset_keys = auto_observe_asset_keys or set()
         self._partition_loading_context = self._evaluator.asset_graph_view.partition_loading_context
         self._logger = logger
+        # Feature-gated: when off, run-request asset check resolution uses the legacy behavior.
+        self._resolve_check_keys = instance.automation_resolve_asset_check_keys_enabled()
 
     @property
     def cursor(self) -> AssetDaemonCursor:
@@ -137,6 +141,7 @@ def _build_run_requests(self, entity_subsets: Sequence[EntitySubset]) -> Sequenc
             asset_graph=self.asset_graph,
             run_tags=self._materialize_run_tags,
             emit_backfills=self._evaluator.emit_backfills,
+            resolve_check_keys_enabled=self._resolve_check_keys,
         )
 
     def _get_updated_cursor(
@@ -315,6 +320,8 @@ def build_run_requests(
     asset_graph: BaseAssetGraph,
     run_tags: Mapping[str, str] | None,
     emit_backfills: bool,
+    *,
+    resolve_check_keys_enabled: bool,
 ) -> Sequence[RunRequest]:
     """For a single asset in a given tick, the asset will only be part of a run or a backfill, not both.
     If the asset is targetd by a backfill, there will only be one backfill that targets the asset.
@@ -330,17 +337,65 @@ def build_run_requests(
         _get_mapping_from_entity_subsets(entity_subsets, asset_graph),
         asset_graph,
         run_tags,
+        resolve_check_keys_enabled=resolve_check_keys_enabled,
     )
     if backfill_run_request:
         run_requests = [backfill_run_request, *run_requests]
 
     return run_requests
 
 
+def _any_check_uses_automation_condition(
+    asset_graph: BaseAssetGraph, asset_keys: AbstractSet[AssetKey]
+) -> bool:
+    return any(
+        asset_graph.get(check_key).automation_condition is not None
+        for check_key in asset_graph.get_check_keys_for_assets(asset_keys)
+    )
+
+
+def _ride_along_check_keys_for_assets(
+    asset_graph: BaseAssetGraph, asset_keys: AbstractSet[AssetKey]
+) -> AbstractSet[AssetCheckKey]:
+    """The checks that ride along by default with a set of assets sharing a single repository.
+
+    These are the checks that target the selected assets, restricted to those that:
+
+    - do NOT own an automation condition: a check with its own automation condition is scheduled
+      independently by declarative automation, so it only runs when its own condition fires (i.e.
+      when it is explicitly requested), not whenever its asset is materialized; and
+    - are defined in the same repository as the assets: a check can target an asset that lives in
+      a different code location, but such a check is not part of this repository's implicit asset
+      job, so it cannot be grouped into the run (otherwise downstream the run fails to build an
+      execution plan for it).
+    """
+    candidate_check_keys = {
+        check_key
+        for check_key in asset_graph.get_check_keys_for_assets(asset_keys)
+        if asset_graph.get(check_key).automation_condition is None
+    }
+    if not asset_keys or not isinstance(asset_graph, RemoteWorkspaceAssetGraph):
+        # A non-workspace graph only ever spans a single repository, so every candidate check
+        # is necessarily defined alongside the assets.
+        return candidate_check_keys
+
+    def repo_selector(entity_key: EntityKey) -> RepositorySelector:
+        return asset_graph.get_repository_handle(entity_key).to_selector()
+
+    # All assets in `asset_keys` share a repository (the mapping is split per-repository before
+    # this is called), so any of them identifies the target repository.
+    target_repo = repo_selector(next(iter(asset_keys)))
+    return {
+        check_key for check_key in candidate_check_keys if repo_selector(check_key) == target_repo
+    }
+
+
 def _build_run_requests_from_partitions_def_mapping(
     mapping: _PartitionsDefKeyMapping,
     asset_graph: BaseAssetGraph,
     run_tags: Mapping[str, str] | None,
+    *,
+    resolve_check_keys_enabled: bool,
 ) -> Sequence[RunRequest]:
     run_requests = []
 
@@ -352,19 +407,30 @@ def _build_run_requests_from_partitions_def_mapping(
             tags.update({**partitions_def.get_tags_for_partition_key(partition_key)})
 
         for entity_keys_for_repo in asset_graph.split_entity_keys_by_repository(entity_keys):
-            asset_check_keys = [k for k in entity_keys_for_repo if isinstance(k, AssetCheckKey)]
+            asset_keys = [k for k in entity_keys_for_repo if isinstance(k, AssetKey)]
+            requested_check_keys = [k for k in entity_keys_for_repo if isinstance(k, AssetCheckKey)]
+            if resolve_check_keys_enabled and _any_check_uses_automation_condition(
+                asset_graph, set(asset_keys)
+            ):
+                ride_along_check_keys = _ride_along_check_keys_for_assets(
+                    asset_graph, set(asset_keys)
+                )
+                resolved_asset_check_keys = list({*requested_check_keys, *ride_along_check_keys})
+            else:
+                # Reached when the feature is gated off, or when no check on these assets uses DA:
+                # pass the explicitly-requested checks, or `None` to let the code server expand to
+                # all of the assets' checks at execution time.
+                resolved_asset_check_keys = requested_check_keys or None
             run_requests.append(
                 # Do not call run_request.with_resolved_tags_and_config as the partition key is
                 # valid and there is no config.
                 # Calling with_resolved_tags_and_config is costly in asset reconciliation as it
                 # checks for valid partition keys.
                 RunRequest(
-                    asset_selection=[k for k in entity_keys_for_repo if isinstance(k, AssetKey)],
+                    asset_selection=asset_keys,
                     partition_key=partition_key,
                     tags=tags,
-                    # if selecting no asset_check_keys, just pass in `None` to allow required
-                    # checks to be included
-                    asset_check_keys=asset_check_keys or None,
+                    asset_check_keys=resolved_asset_check_keys,
                 )
             )
 
@@ -380,6 +446,8 @@ def build_run_requests_with_backfill_policies(
     asset_partitions: Iterable[AssetKeyPartitionKey],
     asset_graph: BaseAssetGraph,
     dynamic_partitions_store: DynamicPartitionsStore,
+    *,
+    resolve_check_keys_enabled: bool,
 ) -> Sequence[RunRequest]:
     """Build run requests for a selection of asset partitions based on the associated BackfillPolicies."""
     run_requests = []
@@ -430,7 +498,12 @@ def build_run_requests_with_backfill_policies(
                 (partitions_def, pk): entity_keys for pk in (partition_keys or [None])
             }
             run_requests.extend(
-                _build_run_requests_from_partitions_def_mapping(mapping, asset_graph, run_tags={})
+                _build_run_requests_from_partitions_def_mapping(
+                    mapping,
+                    asset_graph,
+                    run_tags={},
+                    resolve_check_keys_enabled=resolve_check_keys_enabled,
+                )
             )
         else:
             run_requests.extend(

python_modules/dagster/dagster/_core/execution/asset_backfill.py

@@ -1661,6 +1661,7 @@ def _execute_asset_backfill_iteration_inner(
         asset_partitions=asset_partitions_to_request,
         asset_graph=asset_graph,
         dynamic_partitions_store=instance_queryer,
+        resolve_check_keys_enabled=instance_queryer.instance.automation_resolve_asset_check_keys_enabled(),
     )
     if run_config is not None:
         run_requests = [rr._replace(run_config=run_config) for rr in run_requests]

python_modules/dagster/dagster/_core/instance/methods/asset_methods.py

@@ -456,6 +456,10 @@ def internal_asset_freshness_enabled(self) -> bool:
         """Check if internal asset freshness is enabled - moved from AssetMixin.internal_asset_freshness_enabled()."""
         return os.getenv("DAGSTER_ASSET_FRESHNESS_ENABLED", "").lower() != "false"
 
+    def automation_resolve_asset_check_keys_enabled(self) -> bool:
+        """Whether declarative automation explicitly resolves the asset checks recorded on its run requests."""
+        return os.getenv("DAGSTER_AUTOMATION_RESOLVE_ASSET_CHECK_KEYS", "").lower() == "true"
+
     def streamline_read_asset_health_supported(self, streamline_name: StreamlineName) -> bool:
         """Check if streamline read asset health is supported - moved from AssetMixin.streamline_read_asset_health_supported()."""
         return False

python_modules/dagster/dagster_tests/declarative_automation_tests/daemon_tests/definitions/check_conditioned_sibling_not_firing.py

@@ -0,0 +1,29 @@
+import dagster as dg
+
+
+@dg.asset
+def raw_files() -> None: ...
+
+
+@dg.asset(automation_condition=dg.AutomationCondition.eager(), deps=[raw_files])
+def processed_files() -> None: ...
+
+
+# Has an automation condition, but a time-gated one that will NOT fire during the test window
+# (the cron is far from the frozen test time). DA therefore never individually requests this
+# check on the ticks exercised by the test. It is also excluded from the "default" ride-along
+# set, because it owns an automation condition.
+@dg.asset_check(
+    asset=processed_files, automation_condition=dg.AutomationCondition.on_cron("0 0 1 1 *")
+)
+def yearly_check() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+# No automation condition -> rides along whenever `processed_files` is materialized.
+@dg.asset_check(asset=processed_files)
+def non_null() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+defs = dg.Definitions(assets=[raw_files, processed_files], asset_checks=[yearly_check, non_null])

python_modules/dagster/dagster_tests/declarative_automation_tests/daemon_tests/definitions/check_on_other_asset_excluded.py

@@ -0,0 +1,43 @@
+import dagster as dg
+
+any_dep_newly_updated = dg.AutomationCondition.any_deps_match(
+    dg.AutomationCondition.newly_updated() | dg.AutomationCondition.will_be_requested()
+)
+
+
+@dg.asset
+def raw_files() -> None: ...
+
+
+@dg.asset(automation_condition=dg.AutomationCondition.eager(), deps=[raw_files])
+def processed_files() -> None: ...
+
+
+# A conditioned check on `processed_files` -- fires alongside it, so the run takes the DA
+# check-resolution path (rather than the no-DA-checks fallback).
+@dg.asset_check(asset=processed_files, automation_condition=any_dep_newly_updated)
+def row_count() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+# An unconditioned check on `processed_files` -- rides along.
+@dg.asset_check(asset=processed_files)
+def non_null() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+# A DIFFERENT asset with its own unconditioned check. This check targets `other_asset`, not
+# `processed_files`, so it must never be attached to a run that targets `processed_files`.
+@dg.asset
+def other_asset() -> None: ...
+
+
+@dg.asset_check(asset=other_asset)
+def other_check() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+defs = dg.Definitions(
+    assets=[raw_files, processed_files, other_asset],
+    asset_checks=[row_count, non_null, other_check],
+)

python_modules/dagster/dagster_tests/declarative_automation_tests/daemon_tests/definitions/check_subset_with_unconditioned_sibling.py

@@ -0,0 +1,26 @@
+import dagster as dg
+
+any_dep_newly_updated = dg.AutomationCondition.any_deps_match(
+    dg.AutomationCondition.newly_updated() | dg.AutomationCondition.will_be_requested()
+)
+
+
+@dg.asset
+def raw_files() -> None: ...
+
+
+@dg.asset(automation_condition=dg.AutomationCondition.eager(), deps=[raw_files])
+def processed_files() -> None: ...
+
+
+@dg.asset_check(asset=processed_files, automation_condition=any_dep_newly_updated)
+def row_count() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+@dg.asset_check(asset=processed_files)
+def non_null() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+defs = dg.Definitions(assets=[raw_files, processed_files], asset_checks=[row_count, non_null])

python_modules/dagster/dagster_tests/declarative_automation_tests/daemon_tests/definitions/check_without_condition.py

@@ -0,0 +1,25 @@
+import dagster as dg
+
+
+@dg.asset
+def raw_files() -> None: ...
+
+
+@dg.asset(automation_condition=dg.AutomationCondition.eager(), deps=[raw_files])
+def processed_files() -> None: ...
+
+
+# These checks have NO automation condition of their own. They are not individually
+# requested by the automation tick -- they ride along with any run that materializes
+# `processed_files`. The run request should therefore record all of them explicitly.
+@dg.asset_check(asset=processed_files)
+def row_count() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+@dg.asset_check(asset=processed_files)
+def non_null() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+defs = dg.Definitions(assets=[raw_files, processed_files], asset_checks=[row_count, non_null])

python_modules/dagster/dagster_tests/declarative_automation_tests/daemon_tests/definitions/cross_location_da_check_asset.py

@@ -0,0 +1,22 @@
+import dagster as dg
+
+any_dep_newly_updated = dg.AutomationCondition.any_deps_match(
+    dg.AutomationCondition.newly_updated() | dg.AutomationCondition.will_be_requested()
+)
+
+
+@dg.asset
+def raw_files() -> None: ...
+
+
+@dg.asset(automation_condition=dg.AutomationCondition.eager(), deps=[raw_files])
+def processed_files() -> None: ...
+
+
+# A SAME-location check that owns an automation condition and fires alongside `processed_files`.
+@dg.asset_check(asset=processed_files, automation_condition=any_dep_newly_updated)
+def row_count() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+defs = dg.Definitions(assets=[raw_files, processed_files], asset_checks=[row_count])

python_modules/dagster/dagster_tests/declarative_automation_tests/daemon_tests/definitions/cross_location_da_check_other.py

@@ -0,0 +1,12 @@
+import dagster as dg
+
+
+# `processed_files` is defined in a different code location (cross_location_da_check_asset.py).
+# This unconditioned check targets it but lives here, so it must be excluded from the ride-along
+# set resolved for `processed_files` in the other location (repo-scoping).
+@dg.asset_check(asset=dg.AssetKey("processed_files"), name="external_not_null")
+def external_not_null() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+defs = dg.Definitions(asset_checks=[external_not_null])

python_modules/dagster/dagster_tests/declarative_automation_tests/daemon_tests/definitions/cross_location_unconditioned_check.py

@@ -0,0 +1,10 @@
+import dagster as dg
+
+
+# `processed_files` is defined in a different code location
+@dg.asset_check(asset=dg.AssetKey("processed_files"))
+def external_non_null() -> dg.AssetCheckResult:
+    return dg.AssetCheckResult(passed=True)
+
+
+defs = dg.Definitions(asset_checks=[external_non_null])