wip

dagster-io · Jan 8, 2025 · da5e923 · da5e923
1 parent 0ff03a6
commit da5e923
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 6 deletions.
diff --git a/js_modules/dagster-ui/packages/ui-core/src/automation/AutomationAssetSelectionFragment.tsx b/js_modules/dagster-ui/packages/ui-core/src/automation/AutomationAssetSelectionFragment.tsx
@@ -4,6 +4,12 @@ import {PYTHON_ERROR_FRAGMENT} from '../app/PythonErrorFragment';
 export const AUTOMATION_ASSET_SELECTION_FRAGMENT = gql`
   fragment AutomationAssetSelectionFragment on AssetSelection {
     assetSelectionString
+    assetChecks {
+      name
+      assetKey {
+        path
+      }
+    }
     assetsOrError {
       ... on AssetConnection {
         nodes {

diff --git a/js_modules/dagster-ui/packages/ui-core/src/automation/AutomationTargetList.tsx b/js_modules/dagster-ui/packages/ui-core/src/automation/AutomationTargetList.tsx
@@ -89,7 +89,7 @@ const AssetSelectionTag = ({
   assetSelection: AutomationAssetSelectionFragment;
   automationType: AutomationType;
 }) => {
-  const [showAssetSelection, setShowAssetSelection] = useState(false);
+  const [showDialog, setShowDialog] = useState(false);
 
   const error =
     assetSelection.assetsOrError.__typename === 'PythonError' ? assetSelection.assetsOrError : null;
@@ -129,9 +129,9 @@ const AssetSelectionTag = ({
   return (
     <>
       <Dialog
-        isOpen={showAssetSelection}
+        isOpen={showDialog}
         title="Targeted assets"
-        onClose={() => setShowAssetSelection(false)}
+        onClose={() => setShowDialog(false)}
         style={{width: '750px', maxWidth: '80vw', minWidth: '500px'}}
         canOutsideClickClose
         canEscapeKeyClose
@@ -162,7 +162,7 @@ const AssetSelectionTag = ({
           <Button
             intent="primary"
             onClick={() => {
-              setShowAssetSelection(false);
+              setShowDialog(false);
             }}
           >
             Close
@@ -181,7 +181,7 @@ const AssetSelectionTag = ({
                 body: <PythonErrorInfo error={error} />,
               });
             } else {
-              setShowAssetSelection(true);
+              setShowDialog(true);
             }
           }}
           color={error ? Colors.textRed() : Colors.linkDefault()}

diff --git a/python_modules/dagster-graphql/dagster_graphql/schema/asset_selections.py b/python_modules/dagster-graphql/dagster_graphql/schema/asset_selections.py
@@ -3,11 +3,12 @@
 import graphene
 from dagster._core.definitions.asset_key import AssetKey
 from dagster._core.definitions.asset_selection import AssetSelection, CodeLocationAssetSelection
+from dagster._core.definitions.asset_check_spec import AssetCheckKey
 from dagster._core.remote_representation.handle import RepositoryHandle
 
 from dagster_graphql.implementation.fetch_assets import get_asset
 from dagster_graphql.implementation.utils import capture_error
-from dagster_graphql.schema.entity_key import GrapheneAssetKey
+from dagster_graphql.schema.entity_key import GrapheneAssetKey, GrapheneAssetCheckHandle
 from dagster_graphql.schema.util import ResolveInfo, non_null_list
 
 if TYPE_CHECKING:
@@ -17,6 +18,7 @@
 class GrapheneAssetSelection(graphene.ObjectType):
     assetSelectionString = graphene.String()
     assetKeys = non_null_list(GrapheneAssetKey)
+    assetChecks = non_null_list(GrapheneAssetCheckHandle)
     assets = non_null_list("dagster_graphql.schema.pipelines.pipeline.GrapheneAsset")
     assetsOrError = graphene.NonNull("dagster_graphql.schema.roots.assets.GrapheneAssetsOrError")
 
@@ -28,6 +30,7 @@ def __init__(
         self._asset_selection = asset_selection
         self._repository_handle = repository_handle
         self._resolved_keys = None
+        self._resolved_checks = None
 
     def resolve_assetSelectionString(self, _graphene_info) -> str:
         return str(
@@ -43,6 +46,12 @@ def resolve_assetKeys(self, graphene_info: ResolveInfo):
             for asset_key in self._get_resolved_and_sorted_keys(graphene_info)
         ]
 
+    def resolve_assetChecks(self, graphene_info: ResolveInfo):
+        return [
+            GrapheneAssetCheckHandle(handle)
+            for handle in self._get_resolved_and_sorted_checks(graphene_info)
+        ]
+
     def _get_assets(self, graphene_info: ResolveInfo):
         return [
             get_asset(graphene_info, asset_key)
@@ -60,6 +69,14 @@ def _get_resolved_and_sorted_keys(self, graphene_info: ResolveInfo) -> Sequence[
 
         return self._resolved_keys
 
+    def _get_resolved_and_sorted_checks(self, graphene_info: ResolveInfo) -> Sequence[AssetCheckKey]:
+        """Use this to maintain stability in ordering."""
+        if self._resolved_checks is None:
+            repo = graphene_info.context.get_repository(self._repository_handle)
+            self._resolved_checks = sorted(self._asset_selection.resolve_checks(repo.asset_graph), key=str)
+
+        return self._resolved_checks
+
     @capture_error
     def resolve_assetsOrError(self, graphene_info) -> "GrapheneAssetConnection":
         from dagster_graphql.schema.roots.assets import GrapheneAssetConnection